Zod Schema Converter for Google Generative AI Compatibility

gemini.object-converter.ts

/**
 * Gemini generated object to original schema compatible object converter
 * 
 * Utility to convert Gemini generated objects back to their original schema format.
 * Pairs with gemini.schema-converter.ts
 */
import { z } from "zod";

function findDiscriminatorKey(schema: z.ZodTypeAny): string | null {
  if (schema instanceof z.ZodUnion) {
    const options = (schema as any)._def.options as z.ZodTypeAny[];

    if (options.every((opt: z.ZodTypeAny) => opt instanceof z.ZodObject)) {
      const firstOption = options[0];
      const possibleKeys = Object.keys(firstOption.shape);

      return (
        possibleKeys.find((key) =>
          options.every(
            (opt) => key in opt.shape && opt.shape[key] instanceof z.ZodLiteral
          )
        ) || null
      );
    }
  } else if (schema instanceof z.ZodObject) {
    for (const [_, value] of Object.entries(schema.shape)) {
      const discriminator = findDiscriminatorKey(value as z.ZodTypeAny);
      if (discriminator) return discriminator;
    }
  }

  return null;
}

export function convertBackToOriginalFormat(
  convertedObj: any,
  originalSchema: z.ZodTypeAny
): any {
  if (!convertedObj) {
    return convertedObj;
  }

  if (typeof convertedObj !== "object") {
    return convertedObj;
  }

  if (Array.isArray(convertedObj)) {
    return convertedObj;
  }

  const result: any = {};
  const discriminatorKey = findDiscriminatorKey(originalSchema);

  for (const [key, value] of Object.entries(convertedObj)) {
    if (value && typeof value === "object") {
      if (!Array.isArray(value)) {
        const innerKeys = Object.keys(value);
        if (innerKeys.length === 1 && discriminatorKey) {
          const discriminatorValue = innerKeys[0];
          const innerValue = (value as Record<string, any>)[discriminatorValue];

          result[key] = {
            [discriminatorKey]: discriminatorValue,
            ...innerValue,
          };
        } else {
          result[key] = value;
        }
      } else {
        result[key] = value;
      }
    } else {
      result[key] = value;
    }
  }

  return result;
}

gemini.schema-converter.ts

/**
 * Zod Schema Converter for Google Gemini AI
 * 
 * This utility converts Zod schema into Google's Gemini compatible format.
 * Handles conversion of unions, objects, arrays, and optional fields.
 * @see https://sdk.vercel.ai/providers/ai-sdk-providers/google-generative-ai#troubleshooting-schema-limitations
 * @see https://github.com/vercel/ai/issues/2974
 * 
 * Usual error thrown by Gemini:
 * ```
 * GenerateContentRequest.generation_config.response_schema.properties[occupation].type: must be specified
 * ```
 */

import { z } from "zod";

function convertFieldsWithDescriptions(
  fields: Record<string, z.ZodTypeAny>
): Record<string, z.ZodTypeAny> {
  return Object.fromEntries(
    Object.entries(fields).map(([key, value]) => {
      const fieldDescription = (value as any)._def.description;
      const converted = convertToGeminiCompatibleSchema(value);
      return [
        key,
        fieldDescription ? converted.describe(fieldDescription) : converted,
      ];
    })
  );
}

export function convertToGeminiCompatibleSchema(
  schema: z.ZodTypeAny
): z.ZodTypeAny {
  if (schema instanceof z.ZodUnion) {
    const options = (schema as any)._def.options;
    const description = (schema as any)._def.description;

    if (options.every((opt: z.ZodTypeAny) => opt instanceof z.ZodObject)) {
      const discriminatorKey = findDiscriminatorKey(options);
      if (!discriminatorKey) {
        throw new Error(
          "Cannot convert union without common discriminator key"
        );
      }

      const result: Record<string, z.ZodTypeAny> = {};

      options.forEach((option: z.ZodObject<any>) => {
        const discriminatorValue = option.shape[discriminatorKey];
        if (discriminatorValue instanceof z.ZodLiteral) {
          const key = discriminatorValue.value;
          const { [discriminatorKey]: _, ...otherFields } = option.shape;

          const convertedFields = convertFieldsWithDescriptions(otherFields);
          const optionDescription =
            (option as any)._def.description ||
            (option.shape[discriminatorKey] as any)._def.description;

          result[key] = z
            .object(convertedFields)
            .optional()
            .describe(optionDescription || "");
        }
      });

      return z
        .object(result)
        .describe(description || "")
        .refine((data) => Object.keys(data).length === 1, {
          message: "Exactly one option must be specified",
        });
    }
  } else if (schema instanceof z.ZodObject) {
    const convertedShape = convertFieldsWithDescriptions(schema.shape);
    const description = (schema as any)._def.description;
    return z.object(convertedShape).describe(description || "");
  } else if (schema instanceof z.ZodOptional) {
    const innerSchema = convertToGeminiCompatibleSchema(schema.unwrap());
    const description = (schema as any)._def.description;
    return z.optional(innerSchema).describe(description || "");
  } else if (schema instanceof z.ZodArray) {
    const innerSchema = convertToGeminiCompatibleSchema(schema.element);
    const description = (schema as any)._def.description;
    return z.array(innerSchema).describe(description || "");
  }

  return schema;
}

function findDiscriminatorKey(options: z.ZodObject<any>[]): string | null {
  if (options.length === 0) return null;

  const firstOption = options[0];
  const possibleKeys = Object.keys(firstOption.shape);

  return (
    possibleKeys.find((key) =>
      options.every(
        (opt) => key in opt.shape && opt.shape[key] instanceof z.ZodLiteral
      )
    ) || null
  );
}