Skip to content

Commit 7ea196e

Browse files
edelaunaedelauna
committed
fix(gemini): preserve keyword-named tool parameters during schema sanitization
1 parent eddfcb0 commit 7ea196e

2 files changed

Lines changed: 53 additions & 26 deletions

File tree

‎src/api/providers/__tests__/gemini-handler.spec.ts‎

Lines changed: 41 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@ describe("GeminiHandler backend support", () => {
3030
} as ApiHandlerOptions
3131
const handler = new GeminiHandler(options)
3232
const stub = vi.fn().mockReturnValue((async function* () {})())
33-
// @ts-ignore access private client
3433
handler["client"].models.generateContentStream = stub
3534
await handler.createMessage("instr", [] as any).next()
3635
const config = stub.mock.calls[0][0].config
@@ -47,7 +46,6 @@ describe("GeminiHandler backend support", () => {
4746
} as ApiHandlerOptions
4847
const handler = new GeminiHandler(options)
4948
const stub = vi.fn().mockResolvedValue({ text: "ok" })
50-
// @ts-ignore access private client
5149
handler["client"].models.generateContent = stub
5250
const res = await handler.completePrompt("hi")
5351
expect(res).toBe("ok")
@@ -78,7 +76,6 @@ describe("GeminiHandler backend support", () => {
7876
}
7977

8078
const stub = vi.fn().mockReturnValue(mockStream())
81-
// @ts-ignore access private client
8279
handler["client"].models.generateContentStream = stub
8380

8481
const messages = []
@@ -117,7 +114,6 @@ describe("GeminiHandler backend support", () => {
117114
}
118115

119116
const stub = vi.fn().mockReturnValue(mockStream())
120-
// @ts-ignore access private client
121117
handler["client"].models.generateContentStream = stub
122118

123119
const messages = []
@@ -152,7 +148,6 @@ describe("GeminiHandler backend support", () => {
152148

153149
const mockError = new Error("API rate limit exceeded")
154150
const stub = vi.fn().mockRejectedValue(mockError)
155-
// @ts-ignore access private client
156151
handler["client"].models.generateContentStream = stub
157152

158153
await expect(async () => {
@@ -196,7 +191,6 @@ describe("GeminiHandler backend support", () => {
196191
} as ApiHandlerOptions
197192
const handler = new GeminiHandler(options)
198193
const stub = vi.fn().mockReturnValue((async function* () {})())
199-
// @ts-ignore access private client
200194
handler["client"].models.generateContentStream = stub
201195

202196
await handler
@@ -217,7 +211,6 @@ describe("GeminiHandler backend support", () => {
217211
} as ApiHandlerOptions
218212
const handler = new GeminiHandler(options)
219213
const stub = vi.fn().mockReturnValue((async function* () {})())
220-
// @ts-ignore access private client
221214
handler["client"].models.generateContentStream = stub
222215

223216
await handler
@@ -240,7 +233,6 @@ describe("GeminiHandler backend support", () => {
240233
} as ApiHandlerOptions
241234
const handler = new GeminiHandler(options)
242235
const stub = vi.fn().mockReturnValue((async function* () {})())
243-
// @ts-ignore access private client
244236
handler["client"].models.generateContentStream = stub
245237

246238
const manyTools = Array.from({ length: 30 }, (_, index) => ({
@@ -271,7 +263,6 @@ describe("GeminiHandler backend support", () => {
271263
} as ApiHandlerOptions
272264
const handler = new GeminiHandler(options)
273265
const stub = vi.fn().mockReturnValue((async function* () {})())
274-
// @ts-ignore access private client
275266
handler["client"].models.generateContentStream = stub
276267

277268
const manyTools = Array.from({ length: 30 }, (_, index) => ({
@@ -308,7 +299,6 @@ describe("GeminiHandler backend support", () => {
308299
} as ApiHandlerOptions
309300
const handler = new GeminiHandler(options)
310301
const stub = vi.fn().mockReturnValue((async function* () {})())
311-
// @ts-ignore access private client
312302
handler["client"].models.generateContentStream = stub
313303

314304
await handler
@@ -331,7 +321,6 @@ describe("GeminiHandler backend support", () => {
331321
} as ApiHandlerOptions
332322
const handler = new GeminiHandler(options)
333323
const stub = vi.fn().mockReturnValue((async function* () {})())
334-
// @ts-ignore access private client
335324
handler["client"].models.generateContentStream = stub
336325

337326
await handler
@@ -355,7 +344,6 @@ describe("GeminiHandler backend support", () => {
355344
} as ApiHandlerOptions
356345
const handler = new GeminiHandler(options)
357346
const stub = vi.fn().mockReturnValue((async function* () {})())
358-
// @ts-ignore access private client
359347
handler["client"].models.generateContentStream = stub
360348

361349
await handler
@@ -378,7 +366,6 @@ describe("GeminiHandler backend support", () => {
378366
} as ApiHandlerOptions
379367
const handler = new GeminiHandler(options)
380368
const stub = vi.fn().mockReturnValue((async function* () {})())
381-
// @ts-ignore access private client
382369
handler["client"].models.generateContentStream = stub
383370

384371
await handler
@@ -439,7 +426,6 @@ describe("GeminiHandler backend support", () => {
439426
} as ApiHandlerOptions
440427
const handler = new GeminiHandler(options)
441428
const stub = vi.fn().mockReturnValue((async function* () {})())
442-
// @ts-ignore access private client
443429
handler["client"].models.generateContentStream = stub
444430

445431
await handler
@@ -498,7 +484,6 @@ describe("GeminiHandler backend support", () => {
498484
const options = { apiProvider: "gemini" } as ApiHandlerOptions
499485
const handler = new GeminiHandler(options)
500486
const stub = vi.fn().mockReturnValue((async function* () {})())
501-
// @ts-ignore access private client
502487
handler["client"].models.generateContentStream = stub
503488

504489
await handler
@@ -543,7 +528,6 @@ describe("GeminiHandler backend support", () => {
543528
const options = { apiProvider: "gemini" } as ApiHandlerOptions
544529
const handler = new GeminiHandler(options)
545530
const stub = vi.fn().mockReturnValue((async function* () {})())
546-
// @ts-ignore access private client
547531
handler["client"].models.generateContentStream = stub
548532

549533
await handler
@@ -592,7 +576,6 @@ describe("GeminiHandler backend support", () => {
592576
const options = { apiProvider: "gemini" } as ApiHandlerOptions
593577
const handler = new GeminiHandler(options)
594578
const stub = vi.fn().mockReturnValue((async function* () {})())
595-
// @ts-ignore access private client
596579
handler["client"].models.generateContentStream = stub
597580

598581
await handler
@@ -634,7 +617,6 @@ describe("GeminiHandler backend support", () => {
634617
const options = { apiProvider: "gemini" } as ApiHandlerOptions
635618
const handler = new GeminiHandler(options)
636619
const stub = vi.fn().mockReturnValue((async function* () {})())
637-
// @ts-ignore access private client
638620
handler["client"].models.generateContentStream = stub
639621

640622
await handler
@@ -681,5 +663,46 @@ describe("GeminiHandler backend support", () => {
681663
required: ["value"],
682664
})
683665
})
666+
667+
it("should preserve parameter names that collide with stripped schema keywords", async () => {
668+
const options = { apiProvider: "gemini" } as ApiHandlerOptions
669+
const handler = new GeminiHandler(options)
670+
const stub = vi.fn().mockReturnValue((async function* () {})())
671+
handler["client"].models.generateContentStream = stub
672+
673+
await handler
674+
.createMessage("test", [] as any, {
675+
taskId: "test-task",
676+
tools: [
677+
{
678+
type: "function",
679+
function: {
680+
name: "keyword_param_tool",
681+
description: "Tool whose parameter names match JSON Schema keywords",
682+
parameters: {
683+
type: "object",
684+
properties: {
685+
default: { type: "string" },
686+
additionalProperties: { type: "boolean" },
687+
$schema: { type: "string" },
688+
normal: { type: "integer" },
689+
},
690+
required: ["default", "additionalProperties"],
691+
},
692+
},
693+
},
694+
],
695+
})
696+
.next()
697+
698+
const schema = stub.mock.calls[0][0].config.tools[0].functionDeclarations[0].parametersJsonSchema
699+
expect(schema.properties).toEqual({
700+
default: { type: "string" },
701+
additionalProperties: { type: "boolean" },
702+
$schema: { type: "string" },
703+
normal: { type: "integer" },
704+
})
705+
expect(schema.required).toEqual(expect.arrayContaining(["default", "additionalProperties"]))
706+
})
684707
})
685708
})

‎src/api/providers/gemini.ts‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -132,12 +132,16 @@ function sanitizeSchemaForGemini(
132132
}
133133

134134
if (key === "properties" && value && typeof value === "object" && !Array.isArray(value)) {
135-
const sanitizedProperties = sanitizeSchemaForGemini(value, resolvedDefs, activeRefs)
136-
if (sanitizedProperties && typeof sanitizedProperties === "object" && !Array.isArray(sanitizedProperties)) {
137-
result.properties = {
138-
...(result.properties as Record<string, unknown> | undefined),
139-
...(sanitizedProperties as Record<string, unknown>),
140-
}
135+
// Iterate the property map directly so that property names that happen
136+
// to match schema keywords (e.g. "default", "additionalProperties") are
137+
// preserved as-is; only each property's schema value is sanitized.
138+
const sanitizedProperties: Record<string, unknown> = {}
139+
for (const [propName, propSchema] of Object.entries(value as Record<string, unknown>)) {
140+
sanitizedProperties[propName] = sanitizeSchemaForGemini(propSchema, resolvedDefs, activeRefs)
141+
}
142+
result.properties = {
143+
...(result.properties as Record<string, unknown> | undefined),
144+
...sanitizedProperties,
141145
}
142146
continue
143147
}
@@ -301,8 +305,8 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
301305
}
302306

303307
// Do not pass metadata.allowedFunctionNames to Gemini. Live API testing showed
304-
// that declarations can exceed 25 entries, but allowedFunctionNames starts
305-
// returning generic 400 INVALID_ARGUMENT responses at 26 names. It can also
308+
// that allowedFunctionNames triggers a generic 400 INVALID_ARGUMENT at 26 or more
309+
// names. It can also
306310
// reject prior function calls if their names are absent from the current
307311
// allowed list. We still pass all declarations for history compatibility;
308312
// mode/tool restrictions are enforced by the tool execution layer.

0 commit comments

Comments
 (0)