updates google ai safety settings schema

src/shared/api-schemas/google-ai.ts

@@ -10,6 +10,24 @@ const GoogleAIV1ContentSchema = z.object({
   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();
+
 // https://developers.generativeai.google/api/rest/generativelanguage/models/generateContent
 export const GoogleAIV1GenerateContentSchema = z
   .object({
@@ -17,10 +35,11 @@ export const GoogleAIV1GenerateContentSchema = z
     stream: z.boolean().optional().default(false), // also used for router
     contents: z.array(GoogleAIV1ContentSchema),
     tools: z.array(z.object({})).max(0).optional(),
-    safetySettings: z.array(z.object({})).optional(),
+    safetySettings: SafetySettingsSchema,
     systemInstruction: GoogleAIV1ContentSchema.optional(),
-    generationConfig: z.object({
-      temperature: z.number().optional(),
+    generationConfig: z
+      .object({
+        temperature: z.number().min(0).max(2).optional(),
         maxOutputTokens: z.coerce
           .number()
           .int()
@@ -28,10 +47,11 @@ export const GoogleAIV1GenerateContentSchema = z
           .default(16)
           .transform((v) => Math.min(v, 4096)), // TODO: Add config
         candidateCount: z.literal(1).optional(),
-      topP: z.number().optional(),
-      topK: z.number().optional(),
+        topP: z.number().min(0).max(1).optional(),
+        topK: z.number().min(1).max(40).optional(),
         stopSequences: z.array(z.string().max(500)).max(5).optional(),
-    }).default({}),
+      })
+      .default({}),
   })
   .strip();
 export type GoogleAIChatMessage = z.infer<
@@ -120,6 +140,7 @@ export const transformOpenAIToGoogleAI: APIFormatTransformer<
       { category: "HARM_CATEGORY_HATE_SPEECH", threshold: "BLOCK_NONE" },
       { category: "HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold: "BLOCK_NONE" },
       { category: "HARM_CATEGORY_DANGEROUS_CONTENT", threshold: "BLOCK_NONE" },
+      { category: "HARM_CATEGORY_CIVIC_INTEGRITY", threshold: "BLOCK_NONE" },
     ],
   };
 };