From ba4532b38db0a630a1b80f7c2a93d8163f1e1ac0 Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Wed, 9 Oct 2024 17:11:53 -0500
Subject: [PATCH] more fixes for annoying gemini API design that allows arrays
 or single objects for contents parts

---
 src/shared/api-schemas/google-ai.ts | 50 ++++++++++++++++++-----------
 1 file changed, 31 insertions(+), 19 deletions(-)

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(),