From cab726a327dd197d0c35502cd95982a3bbed1122 Mon Sep 17 00:00:00 2001 From: Enrico Ros Date: Tue, 9 Jul 2024 19:31:52 -0700 Subject: [PATCH] AIX: fix Gemini with JSON outputs --- .../aix/server/dispatch/anthropic/anthropic.wiretypes.ts | 2 +- src/modules/aix/server/dispatch/gemini/gemini.wiretypes.ts | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/src/modules/aix/server/dispatch/anthropic/anthropic.wiretypes.ts b/src/modules/aix/server/dispatch/anthropic/anthropic.wiretypes.ts index 9cecb3cac..e698ea9e3 100644 --- a/src/modules/aix/server/dispatch/anthropic/anthropic.wiretypes.ts +++ b/src/modules/aix/server/dispatch/anthropic/anthropic.wiretypes.ts @@ -123,7 +123,7 @@ export const anthropicWireMessageCreateSchema = z.object({ ]).optional(), /** - * + * (optional) Tools that the model can use to generate the response. */ tools: z.array(anthropicWire_ToolUL_Schema).optional(), diff --git a/src/modules/aix/server/dispatch/gemini/gemini.wiretypes.ts b/src/modules/aix/server/dispatch/gemini/gemini.wiretypes.ts index 0c13bdc4e..6a7f50f7d 100644 --- a/src/modules/aix/server/dispatch/gemini/gemini.wiretypes.ts +++ b/src/modules/aix/server/dispatch/gemini/gemini.wiretypes.ts @@ -176,8 +176,15 @@ const geminiSafetySettingSchema = z.object({ threshold: geminiBlockSafetyLevelEnum, }); +const geminiResponseMimeTypeEnum = z.enum([ + 'text/plain', + 'application/json', +]); + const geminiGenerationConfigSchema = z.object({ stopSequences: z.array(z.string()).optional(), + responseMimeType: geminiResponseMimeTypeEnum.optional(), // defaults to 'text/plain' + responseSchema: z.record(z.any()).optional(), // This is a simplified representation of the Schema object candidateCount: z.number().int().optional(), maxOutputTokens: z.number().int().optional(), temperature: z.number().optional(),