adjusts gemini pro model assignment

This commit is contained in:
nai-degen 2024-06-26 13:37:23 -05:00
parent e3d1ab51d1
commit 994b30dcce
2 changed files with 21 additions and 7 deletions

View File

@ -30,7 +30,12 @@ const getModelsResponse = () => {
if (!config.googleAIKey) return { object: "list", data: [] }; if (!config.googleAIKey) return { object: "list", data: [] };
const googleAIVariants = ["gemini-pro", "gemini-1.0-pro", "gemini-1.5-pro"]; const googleAIVariants = [
"gemini-pro",
"gemini-1.0-pro",
"gemini-1.5-pro",
"gemini-1.5-pro-latest",
];
const models = googleAIVariants.map((id) => ({ const models = googleAIVariants.map((id) => ({
id, id,
@ -124,12 +129,21 @@ googleAIRouter.get("/v1/models", handleModelRequest);
googleAIRouter.post( googleAIRouter.post(
"/v1/chat/completions", "/v1/chat/completions",
ipLimiter, ipLimiter,
createPreprocessorMiddleware({ createPreprocessorMiddleware(
inApi: "openai", { inApi: "openai", outApi: "google-ai", service: "google-ai" },
outApi: "google-ai", { afterTransform: [maybeReassignModel] }
service: "google-ai", ),
}),
googleAIProxy googleAIProxy
); );
/** Replaces requests for non-Google AI models with gemini-pro-1.5-latest. */
function maybeReassignModel(req: Request) {
const requested = req.body.model;
if (requested.includes("gemini")) {
return;
}
req.log.info({ requested }, "Reassigning model to gemini-pro-1.5-latest");
req.body.model = "gemini-pro-1.5-latest";
}
export const googleAI = googleAIRouter; export const googleAI = googleAIRouter;

View File

@ -103,7 +103,7 @@ export const transformOpenAIToGoogleAI: APIFormatTransformer<
stops = [...new Set(stops)].slice(0, 5); stops = [...new Set(stops)].slice(0, 5);
return { return {
model: "gemini-pro", model: req.body.model,
stream: rest.stream, stream: rest.stream,
contents, contents,
tools: [], tools: [],