fix(js/plugins/googleai): Handle Gemini API thought signatures. (#3014)

Author: mbleigh

js/plugins/googleai/src/gemini.ts

@@ -660,6 +660,17 @@ function fromCodeExecutionResult(part: GeminiPart): Part {
   };
 }
 
+function fromThought(part: {
+  thought: boolean;
+  text?: string;
+  thoughtSignature?: string;
+}): Part {
+  return {
+    reasoning: part.text || '',
+    metadata: { thoughtSignature: (part as any).thoughtSignature },
+  };
+}
+
 function toCustomPart(part: Part): GeminiPart {
   if (!part.custom) {
     throw new Error('Invalid GeminiPart: missing custom');
@@ -673,6 +684,14 @@ function toCustomPart(part: Part): GeminiPart {
   throw new Error('Unsupported Custom Part type');
 }
 
+function toThought(part: Part) {
+  const outPart: any = { thought: true };
+  if (part.metadata?.thoughtSignature)
+    outPart.thoughtSignature = part.metadata.thoughtSignature;
+  if (part.reasoning?.length) outPart.text = part.reasoning;
+  return outPart;
+}
+
 function toGeminiPart(part: Part): GeminiPart {
   if (part.text !== undefined) return { text: part.text || ' ' };
   if (part.media) {
@@ -682,6 +701,7 @@ function toGeminiPart(part: Part): GeminiPart {
   if (part.toolRequest) return toFunctionCall(part);
   if (part.toolResponse) return toFunctionResponse(part);
   if (part.custom) return toCustomPart(part);
+  if (typeof part.reasoning === 'string') return toThought(part);
   throw new Error('Unsupported Part type' + JSON.stringify(part));
 }
 
@@ -690,19 +710,16 @@ function fromGeminiPart(
   jsonMode: boolean,
   ref: string
 ): Part {
-  if (part.text !== undefined) {
-    if ((part as any).thought === true) {
-      return { reasoning: part.text };
-    }
-    return { text: part.text };
-  }
+  if ('thought' in part) return fromThought(part as any);
+  if (typeof part.text === 'string') return { text: part.text };
   if (part.inlineData) return fromInlineData(part);
   if (part.functionCall) return fromFunctionCall(part, ref);
   if (part.functionResponse) return fromFunctionResponse(part);
   if (part.executableCode) return fromExecutableCode(part);
   if (part.codeExecutionResult) return fromCodeExecutionResult(part);
-  throw new Error('Unsupported GeminiPart type');
+  throw new Error('Unsupported GeminiPart type: ' + JSON.stringify(part));
 }
+
 export function toGeminiMessage(
   message: MessageData,
   model?: ModelReference<z.ZodTypeAny>

js/plugins/googleai/tests/gemini_test.ts

@@ -136,6 +136,17 @@ describe('toGeminiMessages', () => {
         ],
       },
     },
+    {
+      should: 'should re-populate thoughtSignature from reasoning metadata',
+      inputMessage: {
+        role: 'model',
+        content: [{ reasoning: '', metadata: { thoughtSignature: 'abc123' } }],
+      },
+      expectedOutput: {
+        role: 'model',
+        parts: [{ thought: true, thoughtSignature: 'abc123' }],
+      },
+    },
   ];
   for (const test of testCases) {
     it(test.should, () => {
@@ -365,6 +376,102 @@ describe('fromGeminiCandidate', () => {
         },
       },
     },
+    {
+      should:
+        'should transform gemini candidate to genkit candidate (thought parts) correctly',
+      geminiCandidate: {
+        content: {
+          role: 'model',
+          parts: [
+            {
+              thought: true,
+              thoughtSignature: 'abc123',
+            },
+            {
+              thought: true,
+              text: 'thought with text',
+              thoughtSignature: 'def456',
+            },
+          ],
+        },
+        finishReason: 'STOP',
+        safetyRatings: [
+          {
+            category: 'HARM_CATEGORY_HATE_SPEECH',
+            probability: 'NEGLIGIBLE',
+            probabilityScore: 0.11858909,
+            severity: 'HARM_SEVERITY_NEGLIGIBLE',
+            severityScore: 0.11456649,
+          },
+          {
+            category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+            probability: 'NEGLIGIBLE',
+            probabilityScore: 0.13857833,
+            severity: 'HARM_SEVERITY_NEGLIGIBLE',
+            severityScore: 0.11417085,
+          },
+          {
+            category: 'HARM_CATEGORY_HARASSMENT',
+            probability: 'NEGLIGIBLE',
+            probabilityScore: 0.28012377,
+            severity: 'HARM_SEVERITY_NEGLIGIBLE',
+            severityScore: 0.112405084,
+          },
+          {
+            category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+            probability: 'NEGLIGIBLE',
+          },
+        ],
+      },
+      expectedOutput: {
+        index: 0,
+        message: {
+          role: 'model',
+          content: [
+            {
+              reasoning: '',
+              metadata: { thoughtSignature: 'abc123' },
+            },
+            {
+              reasoning: 'thought with text',
+              metadata: { thoughtSignature: 'def456' },
+            },
+          ],
+        },
+        finishReason: 'stop',
+        finishMessage: undefined,
+        custom: {
+          citationMetadata: undefined,
+          safetyRatings: [
+            {
+              category: 'HARM_CATEGORY_HATE_SPEECH',
+              probability: 'NEGLIGIBLE',
+              probabilityScore: 0.11858909,
+              severity: 'HARM_SEVERITY_NEGLIGIBLE',
+              severityScore: 0.11456649,
+            },
+            {
+              category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+              probability: 'NEGLIGIBLE',
+              probabilityScore: 0.13857833,
+              severity: 'HARM_SEVERITY_NEGLIGIBLE',
+              severityScore: 0.11417085,
+            },
+            {
+              category: 'HARM_CATEGORY_HARASSMENT',
+              probability: 'NEGLIGIBLE',
+              probabilityScore: 0.28012377,
+              severity: 'HARM_SEVERITY_NEGLIGIBLE',
+              severityScore: 0.112405084,
+            },
+            {
+              category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+              probability: 'NEGLIGIBLE',
+            },
+          ],
+        },
+      },
+    },
   ];
   for (const test of testCases) {
     it(test.should, () => {