feat(provider): add Gemini tool schema sanitization (#1132)

This commit is contained in:
Rami Chowdhury 2025-07-18 16:02:54 -04:00 committed by GitHub
parent 5c9d1910af
commit f1da70b1de
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -494,7 +494,10 @@ export namespace Provider {
...t,
parameters: optionalToNullable(t.parameters),
})),
google: TOOLS,
google: TOOLS.map((t) => ({
...t,
parameters: sanitizeGeminiParameters(t.parameters),
})),
}
export async function tools(providerID: string) {
@ -508,6 +511,60 @@ export namespace Provider {
return TOOL_MAPPING[providerID] ?? TOOLS
}
function sanitizeGeminiParameters(schema: z.ZodTypeAny, visited = new Set()): z.ZodTypeAny {
if (!schema || visited.has(schema)) {
return schema
}
visited.add(schema)
if (schema instanceof z.ZodDefault) {
const innerSchema = schema.removeDefault()
// Handle Gemini's incompatibility with `default` on `anyOf` (unions).
if (innerSchema instanceof z.ZodUnion) {
// The schema was `z.union(...).default(...)`, which is not allowed.
// We strip the default and return the sanitized union.
return sanitizeGeminiParameters(innerSchema, visited)
}
// Otherwise, the default is on a regular type, which is allowed.
// We recurse on the inner type and then re-apply the default.
return sanitizeGeminiParameters(innerSchema, visited).default(schema._def.defaultValue())
}
if (schema instanceof z.ZodOptional) {
return z.optional(sanitizeGeminiParameters(schema.unwrap(), visited))
}
if (schema instanceof z.ZodObject) {
const newShape: Record<string, z.ZodTypeAny> = {}
for (const [key, value] of Object.entries(schema.shape)) {
newShape[key] = sanitizeGeminiParameters(value as z.ZodTypeAny, visited)
}
return z.object(newShape)
}
if (schema instanceof z.ZodArray) {
return z.array(sanitizeGeminiParameters(schema.element, visited))
}
if (schema instanceof z.ZodUnion) {
// This schema corresponds to `anyOf` in JSON Schema.
// We recursively sanitize each option in the union.
const sanitizedOptions = schema.options.map((option: z.ZodTypeAny) => sanitizeGeminiParameters(option, visited))
return z.union(sanitizedOptions as [z.ZodTypeAny, z.ZodTypeAny, ...z.ZodTypeAny[]])
}
if (schema instanceof z.ZodString) {
const newSchema = z.string({ description: schema.description })
const safeChecks = ["min", "max", "length", "regex", "startsWith", "endsWith", "includes", "trim"]
// rome-ignore lint/suspicious/noExplicitAny: <explanation>
;(newSchema._def as any).checks = (schema._def as z.ZodStringDef).checks.filter((check) =>
safeChecks.includes(check.kind),
)
return newSchema
}
return schema
}
function optionalToNullable(schema: z.ZodTypeAny): z.ZodTypeAny {
if (schema instanceof z.ZodObject) {
const shape = schema.shape