fix: fixed gemini model handling of nullish optional fields in respon…

…se schema (#1411)
firebase · Nov 26, 2024 · 9a8a1d4 · 9a8a1d4
1 parent f4261f9
commit 9a8a1d4
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 2 deletions.
diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts
@@ -449,7 +449,7 @@ export function fromGeminiCandidate(
   };
 }
 
-function cleanSchema(schema: JSONSchema): JSONSchema {
+export function cleanSchema(schema: JSONSchema): JSONSchema {
   const out = structuredClone(schema);
   for (const key in out) {
     if (key === '$schema' || key === 'additionalProperties') {
@@ -459,6 +459,12 @@ function cleanSchema(schema: JSONSchema): JSONSchema {
     if (typeof out[key] === 'object') {
       out[key] = cleanSchema(out[key]);
     }
+    // Zod nullish() and picoschema optional fields will produce type `["string", "null"]`
+    // which is not supported by the model API. Convert them to just `"string"`.
+    if (key === 'type' && Array.isArray(out[key])) {
+      // find the first that's not `null`.
+      out[key] = out[key].find((t) => t !== 'null');
+    }
   }
   return out;
 }

diff --git a/js/plugins/googleai/tests/gemini_test.ts b/js/plugins/googleai/tests/gemini_test.ts
@@ -19,6 +19,7 @@ import { GenerateContentCandidate } from '@google/generative-ai';
 import assert from 'node:assert';
 import { describe, it } from 'node:test';
 import {
+  cleanSchema,
   fromGeminiCandidate,
   toGeminiMessage,
   toGeminiSystemInstruction,
@@ -345,3 +346,35 @@ describe('fromGeminiCandidate', () => {
     });
   }
 });
+
+describe('cleanSchema', () => {
+  it('strips nulls from type', () => {
+    const cleaned = cleanSchema({
+      type: 'object',
+      properties: {
+        title: {
+          type: 'string',
+        },
+        subtitle: {
+          type: ['string', 'null'],
+        },
+      },
+      required: ['title'],
+      additionalProperties: true,
+      $schema: 'http://json-schema.org/draft-07/schema#',
+    });
+
+    assert.deepStrictEqual(cleaned, {
+      type: 'object',
+      properties: {
+        title: {
+          type: 'string',
+        },
+        subtitle: {
+          type: 'string',
+        },
+      },
+      required: ['title'],
+    });
+  });
+});
diff --git a/js/plugins/vertexai/src/gemini.ts b/js/plugins/vertexai/src/gemini.ts
@@ -415,7 +415,7 @@ const convertSchemaProperty = (property) => {
   }
 };
 
-function cleanSchema(schema: JSONSchema): JSONSchema {
+export function cleanSchema(schema: JSONSchema): JSONSchema {
   const out = structuredClone(schema);
   for (const key in out) {
     if (key === '$schema' || key === 'additionalProperties') {
@@ -425,6 +425,12 @@ function cleanSchema(schema: JSONSchema): JSONSchema {
     if (typeof out[key] === 'object') {
       out[key] = cleanSchema(out[key]);
     }
+    // Zod nullish() and picoschema optional fields will produce type `["string", "null"]`
+    // which is not supported by the model API. Convert them to just `"string"`.
+    if (key === 'type' && Array.isArray(out[key])) {
+      // find the first that's not `null`.
+      out[key] = out[key].find((t) => t !== 'null');
+    }
   }
   return out;
 }

diff --git a/js/plugins/vertexai/tests/gemini_test.ts b/js/plugins/vertexai/tests/gemini_test.ts
@@ -19,6 +19,7 @@ import { MessageData } from 'genkit';
 import assert from 'node:assert';
 import { describe, it } from 'node:test';
 import {
+  cleanSchema,
   fromGeminiCandidate,
   toGeminiMessage,
   toGeminiSystemInstruction,
@@ -348,3 +349,35 @@ describe('fromGeminiCandidate', () => {
     });
   }
 });
+
+describe('cleanSchema', () => {
+  it('strips nulls from type', () => {
+    const cleaned = cleanSchema({
+      type: 'object',
+      properties: {
+        title: {
+          type: 'string',
+        },
+        subtitle: {
+          type: ['string', 'null'],
+        },
+      },
+      required: ['title'],
+      additionalProperties: true,
+      $schema: 'http://json-schema.org/draft-07/schema#',
+    });
+
+    assert.deepStrictEqual(cleaned, {
+      type: 'object',
+      properties: {
+        title: {
+          type: 'string',
+        },
+        subtitle: {
+          type: 'string',
+        },
+      },
+      required: ['title'],
+    });
+  });
+});