Merge pull request #8865 from uinstinct/gemini-thought-signature

feat: add support for gemini thought signature

Author: RomneyDa

core/llm/llms/Gemini.ts

@@ -23,6 +23,14 @@ import {
   convertContinueToolToGeminiFunction,
 } from "./gemini-types";
 
+interface GeminiToolCallDelta extends ToolCallDelta {
+  extra_content?: {
+    google?: {
+      thought_signature?: string;
+    };
+  };
+}
+
 class Gemini extends BaseLLM {
   static providerName = "gemini";
 
@@ -266,18 +274,37 @@ class Gemini extends BaseLLM {
                   ? [{ text: msg.content }]
                   : msg.content.map(this.continuePartToGeminiPart),
             };
-            if (msg.toolCalls) {
-              msg.toolCalls.forEach((toolCall) => {
-                if (toolCall.function?.name) {
-                  assistantMsg.parts.push({
-                    functionCall: {
-                      name: toolCall.function.name,
-                      args: safeParseToolCallArgs(toolCall),
-                    },
-                  });
-                }
-              });
+
+            if (msg.toolCalls && msg.toolCalls.length) {
+              (msg.toolCalls as GeminiToolCallDelta[]).forEach(
+                (toolCall, index) => {
+                  if (toolCall.function?.name) {
+                    const signatureForCall =
+                      toolCall?.extra_content?.google?.thought_signature;
+
+                    let thoughtSignature: string | undefined;
+                    if (index === 0) {
+                      if (typeof signatureForCall === "string") {
+                        thoughtSignature = signatureForCall;
+                      } else {
+                        // Fallback per https://ai.google.dev/gemini-api/docs/thought-signatures
+                        // for histories that were not generated by Gemini or are missing signatures.
+                        thoughtSignature = "skip_thought_signature_validator";
+                      }
+                    }
+
+                    assistantMsg.parts.push({
+                      functionCall: {
+                        name: toolCall.function.name,
+                        args: safeParseToolCallArgs(toolCall),
+                      },
+                      ...(thoughtSignature && { thoughtSignature }),
+                    });
+                  }
+                },
+              );
             }
+
             return assistantMsg;
           }
           return {
@@ -370,6 +397,7 @@ class Gemini extends BaseLLM {
             if ("text" in part) {
               textParts.push({ type: "text", text: part.text });
             } else if ("functionCall" in part) {
+              const thoughtSignature = part.thoughtSignature;
               toolCalls.push({
                 type: "function",
                 id: part.functionCall.id ?? uuidv4(),
@@ -380,6 +408,13 @@ class Gemini extends BaseLLM {
                       ? part.functionCall.args
                       : JSON.stringify(part.functionCall.args),
                 },
+                ...(thoughtSignature && {
+                  extra_content: {
+                    google: {
+                      thought_signature: thoughtSignature,
+                    },
+                  },
+                }),
               });
             } else {
               // Note: function responses shouldn't be streamed, images not supported

core/llm/llms/gemini-types.ts

@@ -192,6 +192,7 @@ export type GeminiFunctionCallContentPart = {
     name: string;
     args: JSONSchema7Object;
   };
+  thoughtSignature?: string;
 };
 
 export type GeminiFunctionResponseContentPart = {

packages/openai-adapters/src/apis/Gemini.ts

@@ -44,6 +44,24 @@ type UsageInfo = Pick<
   "total_tokens" | "completion_tokens" | "prompt_tokens"
 >;
 
+interface GeminiToolCall
+  extends OpenAI.Chat.Completions.ChatCompletionMessageFunctionToolCall {
+  extra_content?: {
+    google?: {
+      thought_signature?: string;
+    };
+  };
+}
+
+interface GeminiToolDelta
+  extends OpenAI.Chat.Completions.ChatCompletionChunk.Choice.Delta {
+  extra_content?: {
+    google?: {
+      thought_signature?: string;
+    };
+  };
+}
+
 export class GeminiApi implements BaseLlmApi {
   apiBase: string = "https://generativelanguage.googleapis.com/v1beta/";
 
@@ -143,25 +161,43 @@ export class GeminiApi implements BaseLlmApi {
 
           return {
             role: "model" as const,
-            parts: msg.tool_calls.map((toolCall) => {
-              // Type guard for function tool calls
-              if (toolCall.type === "function" && "function" in toolCall) {
-                return {
-                  functionCall: {
-                    id: includeToolCallIds ? toolCall.id : undefined,
-                    name: toolCall.function.name,
-                    args: safeParseArgs(
-                      toolCall.function.arguments,
-                      `Call: ${toolCall.function.name} ${toolCall.id}`,
-                    ),
-                  },
-                };
-              } else {
+            parts: (msg.tool_calls as GeminiToolCall[]).map(
+              (toolCall, index) => {
+                if (toolCall.type === "function" && "function" in toolCall) {
+                  let thoughtSignature: string | undefined;
+                  if (index === 0) {
+                    const rawSignature =
+                      toolCall?.extra_content?.google?.thought_signature;
+
+                    if (
+                      typeof rawSignature === "string" &&
+                      rawSignature.length > 0
+                    ) {
+                      thoughtSignature = rawSignature;
+                    } else {
+                      // Fallback per https://ai.google.dev/gemini-api/docs/thought-signatures
+                      // for histories that were not generated by Gemini or are missing signatures.
+                      thoughtSignature = "skip_thought_signature_validator";
+                    }
+                  }
+
+                  return {
+                    functionCall: {
+                      id: includeToolCallIds ? toolCall.id : undefined,
+                      name: toolCall.function.name,
+                      args: safeParseArgs(
+                        toolCall.function.arguments,
+                        `Call: ${toolCall.function.name} ${toolCall.id}`,
+                      ),
+                    },
+                    ...(thoughtSignature && { thoughtSignature }),
+                  };
+                }
                 throw new Error(
                   `Unsupported tool call type in Gemini: ${toolCall.type}`,
                 );
-              }
-            }),
+              },
+            ),
           };
         }
 
@@ -328,11 +364,27 @@ export class GeminiApi implements BaseLlmApi {
         if (contentParts) {
           for (const part of contentParts) {
             if ("text" in part) {
+              const thoughtSignature = part?.thoughtSignature;
+              if (thoughtSignature) {
+                yield chatChunkFromDelta({
+                  model,
+                  delta: {
+                    role: "assistant",
+                    extra_content: {
+                      google: {
+                        thought_signature: thoughtSignature,
+                      },
+                    },
+                  } as GeminiToolDelta,
+                });
+              }
+
               yield chatChunk({
                 content: part.text,
                 model,
               });
             } else if ("functionCall" in part) {
+              const thoughtSignature = part?.thoughtSignature;
               yield chatChunkFromDelta({
                 model,
                 delta: {
@@ -345,6 +397,13 @@ export class GeminiApi implements BaseLlmApi {
                         name: part.functionCall.name,
                         arguments: JSON.stringify(part.functionCall.args),
                       },
+                      ...(thoughtSignature && {
+                        extra_content: {
+                          google: {
+                            thought_signature: thoughtSignature,
+                          },
+                        },
+                      }),
                     },
                   ],
                 },