From 3beea5dcfc6a6fbe7e46e54340a4696f1d732fb7 Mon Sep 17 00:00:00 2001
From: reanon <85157-reanon@users.noreply.gitgud.io>
Date: Wed, 23 Apr 2025 19:22:46 +0000
Subject: [PATCH] wtf

---
 src/shared/api-schemas/google-ai.ts | 57 ++++++++++++++++-------------
 1 file changed, 32 insertions(+), 25 deletions(-)

diff --git a/src/shared/api-schemas/google-ai.ts b/src/shared/api-schemas/google-ai.ts
index 7686865..4e3f775 100644
--- a/src/shared/api-schemas/google-ai.ts
+++ b/src/shared/api-schemas/google-ai.ts
@@ -62,6 +62,13 @@ export const GoogleAIV1GenerateContentSchema = z
     // except for system_instruction where it randomly uses snake case.
     // google api evidently accepts either case.
     system_instruction: GoogleAIV1ContentSchema.optional(),
+    // These fields are used by SillyTavern but aren't part of the actual API schema
+    // They'll be removed by the transform but used to determine thinkingBudget
+    enable_thinking: z.boolean().optional(),
+    thinking_budget: z.union([
+      z.literal("auto"),
+      z.number().int().min(0).max(24576)
+    ]).optional(),
     generationConfig: z
       .object({
         temperature: z.number().min(0).max(2).optional(),
@@ -76,29 +83,13 @@ export const GoogleAIV1GenerateContentSchema = z
         topK: z.number().min(1).max(40).optional(),
         stopSequences: z.array(z.string().max(500)).max(5).optional(),
         thinkingConfig: z.object({
-          enable_thinking: z.boolean().optional(),
-          includeThoughts: z.boolean().optional(),
           thinkingBudget: z.union([
             z.literal("auto"),
             z.number().int().min(0).max(24576)
           ]).optional()
         }).optional()
       })
-      .default({})
-      .transform((config) => {
-        // Auto-enable thinking when thinkingBudget is set
-        if (config.thinkingConfig?.thinkingBudget !== undefined && 
-            config.thinkingConfig.enable_thinking === undefined) {
-          return {
-            ...config,
-            thinkingConfig: {
-              ...config.thinkingConfig,
-              enable_thinking: true
-            }
-          };
-        }
-        return config;
-      }),
+      .default({}),
   })
   .strip();
 export type GoogleAIChatMessage = z.infer<
@@ -170,18 +161,34 @@ export const transformOpenAIToGoogleAI: APIFormatTransformer<
   stops.push(...Array.from(foundNames).map((name) => `\n${name}:`));
   stops = [...new Set(stops)].slice(0, 5);
 
+  // Setup generation config
+  const generationConfig: any = {
+    maxOutputTokens: rest.max_tokens,
+    stopSequences: stops,
+    topP: rest.top_p,
+    topK: 40, // openai schema doesn't have this, google ai defaults to 40
+    temperature: rest.temperature,
+  };
+
+  // Handle thinking budget from SillyTavern's UI
+  if (req.body.enable_thinking !== undefined || req.body.thinking_budget !== undefined) {
+    const thinkingEnabled = req.body.enable_thinking !== false; // Default to true if not specified
+    const thinkingBudget = req.body.thinking_budget ?? 1000; // Default budget if not specified
+    
+    // Only set thinkingBudget if thinking is enabled
+    if (thinkingEnabled && thinkingBudget !== 0) {
+      generationConfig.thinkingConfig = { 
+        thinkingBudget: thinkingBudget 
+      };
+    }
+  }
+
   return {
     model: req.body.model,
     stream: rest.stream,
     contents,
     tools: [],
-    generationConfig: {
-      maxOutputTokens: rest.max_tokens,
-      stopSequences: stops,
-      topP: rest.top_p,
-      topK: 40, // openai schema doesn't have this, google ai defaults to 40
-      temperature: rest.temperature,
-    },
+    generationConfig,
     safetySettings: [
       { category: "HARM_CATEGORY_HARASSMENT", threshold: "BLOCK_NONE" },
       { category: "HARM_CATEGORY_HATE_SPEECH", threshold: "BLOCK_NONE" },
@@ -197,4 +204,4 @@ export function containsImageContent(contents: GoogleAIChatMessage[]): boolean {
     const parts = Array.isArray(content.parts) ? content.parts : [content.parts];
     return parts.some(part => 'inlineData' in part);
   });
-}
\ No newline at end of file
+}