fix(js/plugins/googleai): add lenient parsing of JSON in data part

firebase · Oct 3, 2024 · 8e4c6fc · 8e4c6fc
1 parent b82af41
commit 8e4c6fc
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts
@@ -15,38 +15,39 @@
  */
 
 import {
+  Content as GeminiMessage,
   FileDataPart,
   FunctionCallPart,
   FunctionDeclaration,
   FunctionDeclarationSchemaType,
   FunctionResponsePart,
   GenerateContentCandidate as GeminiCandidate,
-  Content as GeminiMessage,
-  Part as GeminiPart,
   GenerateContentResponse,
   GenerationConfig,
   GoogleGenerativeAI,
   InlineDataPart,
+  Part as GeminiPart,
   RequestOptions,
   StartChatParams,
   Tool,
 } from '@google/generative-ai';
 import { GENKIT_CLIENT_HEADER, z } from 'genkit';
+import { extractJson } from 'genkit/extract';
 import {
   CandidateData,
+  defineModel,
   GenerationCommonConfigSchema,
+  getBasicUsageStats,
   MediaPart,
   MessageData,
   ModelAction,
   ModelMiddleware,
+  modelRef,
   ModelReference,
   Part,
   ToolDefinitionSchema,
   ToolRequestPart,
   ToolResponsePart,
-  defineModel,
-  getBasicUsageStats,
-  modelRef,
 } from 'genkit/model';
 import {
   downloadRequestMedia,
@@ -387,7 +388,7 @@ function toGeminiPart(part: Part): GeminiPart {
 
 function fromGeminiPart(part: GeminiPart, jsonMode: boolean): Part {
   if (jsonMode && part.text !== undefined) {
-    return { data: JSON.parse(part.text) };
+    return { data: extractJson(part.text) };
   }
   if (part.text !== undefined) return { text: part.text };
   if (part.inlineData) return fromInlineData(part);