fix structured outputs: json schema for gemini models

Author: pnkvalavala

src/providers/google-vertex-ai/transformGenerationConfig.ts

@@ -1,5 +1,8 @@
 import { Params } from '../../types/requestBody';
-import { derefer, recursivelyDeleteUnsupportedParameters } from './utils';
+import {
+  recursivelyDeleteUnsupportedParameters,
+  transformGeminiToolParameters,
+} from './utils';
 import { GoogleEmbedParams } from './embed';
 import { EmbedInstancesData } from './types';
 /**
@@ -39,20 +42,11 @@ export function transformGenerationConfig(params: Params) {
   }
   if (params?.response_format?.type === 'json_schema') {
     generationConfig['responseMimeType'] = 'application/json';
-    recursivelyDeleteUnsupportedParameters(
-      params?.response_format?.json_schema?.schema
-    );
     let schema =
       params?.response_format?.json_schema?.schema ??
       params?.response_format?.json_schema;
-    if (Object.keys(schema).includes('$defs')) {
-      schema = derefer(schema);
-      delete schema['$defs'];
-    }
-    if (Object.hasOwn(schema, '$schema')) {
-      delete schema['$schema'];
-    }
-    generationConfig['responseSchema'] = schema;
+    recursivelyDeleteUnsupportedParameters(schema);
+    generationConfig['responseSchema'] = transformGeminiToolParameters(schema);
   }
 
   if (params?.thinking) {

src/providers/google-vertex-ai/utils.ts

@@ -267,18 +267,25 @@ export const transformGeminiToolParameters = (
       if (key === 'enum' && Array.isArray(value)) {
         transformed.enum = value;
         transformed.format = 'enum';
-      } else if ((key === 'anyOf' || key === 'oneOf') && Array.isArray(value)) {
+        continue;
+      }
+
+      if ((key === 'anyOf' || key === 'oneOf') && Array.isArray(value)) {
         const nonNullItems = value.filter((item) => !isNullTypeNode(item));
-        if (nonNullItems.length < value.length) {
-          // remove `null` type in schema and set nullable: true
-          transformed[key] = transformNode(nonNullItems);
-          transformed.nullable = true;
-        } else {
-          transformed[key] = transformNode(value);
+        const hadNull = nonNullItems.length < value.length;
+        if (nonNullItems.length === 1 && hadNull) {
+          // Flatten to single schema: get rid of anyOf/oneOf and set nullable: true
+          const single = transformNode(nonNullItems[0]);
+          if (single && typeof single === 'object') single.nullable = true;
+          return single;
         }
-      } else {
-        transformed[key] = transformNode(value);
+
+        transformed[key] = transformNode(hadNull ? nonNullItems : value);
+        if (hadNull) transformed.nullable = true;
+        continue;
       }
+
+      transformed[key] = transformNode(value);
     }
     return transformed;
   };

src/providers/google/chatComplete.ts

@@ -11,7 +11,6 @@ import {
 } from '../../types/requestBody';
 import { buildGoogleSearchRetrievalTool } from '../google-vertex-ai/chatComplete';
 import {
-  derefer,
   getMimeType,
   recursivelyDeleteUnsupportedParameters,
   transformGeminiToolParameters,
@@ -63,17 +62,11 @@ const transformGenerationConfig = (params: Params) => {
   }
   if (params?.response_format?.type === 'json_schema') {
     generationConfig['responseMimeType'] = 'application/json';
-    recursivelyDeleteUnsupportedParameters(
-      params?.response_format?.json_schema?.schema
-    );
     let schema =
       params?.response_format?.json_schema?.schema ??
       params?.response_format?.json_schema;
-    if (Object.keys(schema).includes('$defs')) {
-      schema = derefer(schema);
-      delete schema['$defs'];
-    }
-    generationConfig['responseSchema'] = schema;
+    recursivelyDeleteUnsupportedParameters(schema);
+    generationConfig['responseSchema'] = transformGeminiToolParameters(schema);
   }
   if (params?.thinking) {
     const thinkingConfig: Record<string, any> = {};