diff --git a/src/shared/api-schemas/google-ai.ts b/src/shared/api-schemas/google-ai.ts index 5919ab5..fb1e82e 100644 --- a/src/shared/api-schemas/google-ai.ts +++ b/src/shared/api-schemas/google-ai.ts @@ -6,27 +6,39 @@ import { import { APIFormatTransformer } from "./index"; const GoogleAIV1ContentSchema = z.object({ - parts: z.array(z.object({ text: z.string() })), // TODO: add other media types + parts: z + .union([ + z.array(z.object({ text: z.string() })), + z.object({ text: z.string() }), + ]) + // Google allows parts to be an array or a single object, which is really + // annoying for downstream code. We will coerce it to an array here. + .transform((val) => (Array.isArray(val) ? val : [val])), + // TODO: add other media types role: z.enum(["user", "model"]).optional(), }); -const SafetySettingsSchema = z.array(z.object({ - category: z.enum([ - "HARM_CATEGORY_HARASSMENT", - "HARM_CATEGORY_HATE_SPEECH", - "HARM_CATEGORY_SEXUALLY_EXPLICIT", - "HARM_CATEGORY_DANGEROUS_CONTENT", - "HARM_CATEGORY_CIVIC_INTEGRITY", - ]), - threshold: z.enum([ - "OFF", - "BLOCK_NONE", - "BLOCK_ONLY_HIGH", - "BLOCK_MEDIUM_AND_ABOVE", - "BLOCK_LOW_AND_ABOVE", - "HARM_BLOCK_THRESHOLD_UNSPECIFIED", - ]), -})).optional(); +const SafetySettingsSchema = z + .array( + z.object({ + category: z.enum([ + "HARM_CATEGORY_HARASSMENT", + "HARM_CATEGORY_HATE_SPEECH", + "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "HARM_CATEGORY_DANGEROUS_CONTENT", + "HARM_CATEGORY_CIVIC_INTEGRITY", + ]), + threshold: z.enum([ + "OFF", + "BLOCK_NONE", + "BLOCK_ONLY_HIGH", + "BLOCK_MEDIUM_AND_ABOVE", + "BLOCK_LOW_AND_ABOVE", + "HARM_BLOCK_THRESHOLD_UNSPECIFIED", + ]), + }) + ) + .optional(); // https://developers.generativeai.google/api/rest/generativelanguage/models/generateContent export const GoogleAIV1GenerateContentSchema = z @@ -40,7 +52,7 @@ export const GoogleAIV1GenerateContentSchema = z // quick fix for SillyTavern, which uses camel case field names for everything // except for system_instruction where it randomly uses snake case. // google api evidently accepts either case. - "system_instruction": GoogleAIV1ContentSchema.optional(), + system_instruction: GoogleAIV1ContentSchema.optional(), generationConfig: z .object({ temperature: z.number().min(0).max(2).optional(),