Skip to content

Commit 7479672

Browse files
roomoteroomote
committed
fix: prevent urlContext from being added to Vertex AI requests
- Added isVertex property to GeminiHandler to track provider type - Modified tool configuration to only add urlContext for regular Gemini, not Vertex AI - Added comprehensive tests to verify the fix works correctly - Fixes issue where Gemini requests were incorrectly routed to Vertex AI after both providers were configured Fixes #7968
1 parent d09689b commit 7479672

File tree

2 files changed

+308
-2
lines changed

2 files changed

+308
-2
lines changed

‎src/api/providers/__tests__/vertex-gemini-urlcontext.spec.ts‎

Lines changed: 302 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,302 @@
1+
// npx vitest run src/api/providers/__tests__/vertex-gemini-urlcontext.spec.ts
2+
3+
import { Anthropic } from "@anthropic-ai/sdk"
4+
import { GeminiHandler } from "../gemini"
5+
import { VertexHandler } from "../vertex"
6+
7+
describe("Vertex vs Gemini urlContext handling", () => {
8+
describe("GeminiHandler", () => {
9+
it("should include urlContext tool when enableUrlContext is true", async () => {
10+
const mockGenerateContentStream = vitest.fn()
11+
12+
const handler = new GeminiHandler({
13+
geminiApiKey: "test-key",
14+
apiModelId: "gemini-1.5-flash",
15+
enableUrlContext: true,
16+
})
17+
18+
// Replace the client with our mock
19+
handler["client"] = {
20+
models: {
21+
generateContentStream: mockGenerateContentStream,
22+
},
23+
} as any
24+
25+
// Setup mock to return an async generator
26+
mockGenerateContentStream.mockResolvedValue({
27+
[Symbol.asyncIterator]: async function* () {
28+
yield { text: "Test response" }
29+
yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 } }
30+
},
31+
})
32+
33+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
34+
35+
const stream = handler.createMessage("System prompt", messages)
36+
const chunks = []
37+
for await (const chunk of stream) {
38+
chunks.push(chunk)
39+
}
40+
41+
// Verify that generateContentStream was called with urlContext in tools
42+
expect(mockGenerateContentStream).toHaveBeenCalledWith(
43+
expect.objectContaining({
44+
config: expect.objectContaining({
45+
tools: expect.arrayContaining([{ urlContext: {} }]),
46+
}),
47+
}),
48+
)
49+
})
50+
51+
it("should not include urlContext tool when enableUrlContext is false", async () => {
52+
const mockGenerateContentStream = vitest.fn()
53+
54+
const handler = new GeminiHandler({
55+
geminiApiKey: "test-key",
56+
apiModelId: "gemini-1.5-flash",
57+
enableUrlContext: false,
58+
})
59+
60+
// Replace the client with our mock
61+
handler["client"] = {
62+
models: {
63+
generateContentStream: mockGenerateContentStream,
64+
},
65+
} as any
66+
67+
// Setup mock to return an async generator
68+
mockGenerateContentStream.mockResolvedValue({
69+
[Symbol.asyncIterator]: async function* () {
70+
yield { text: "Test response" }
71+
yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 } }
72+
},
73+
})
74+
75+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
76+
77+
const stream = handler.createMessage("System prompt", messages)
78+
const chunks = []
79+
for await (const chunk of stream) {
80+
chunks.push(chunk)
81+
}
82+
83+
// Verify that generateContentStream was called without urlContext in tools
84+
expect(mockGenerateContentStream).toHaveBeenCalledWith(
85+
expect.objectContaining({
86+
config: expect.not.objectContaining({
87+
tools: expect.anything(),
88+
}),
89+
}),
90+
)
91+
})
92+
93+
it("should include urlContext in completePrompt when enableUrlContext is true", async () => {
94+
const mockGenerateContent = vitest.fn()
95+
96+
const handler = new GeminiHandler({
97+
geminiApiKey: "test-key",
98+
apiModelId: "gemini-1.5-flash",
99+
enableUrlContext: true,
100+
})
101+
102+
// Replace the client with our mock
103+
handler["client"] = {
104+
models: {
105+
generateContent: mockGenerateContent,
106+
},
107+
} as any
108+
109+
// Mock the response
110+
mockGenerateContent.mockResolvedValue({
111+
text: "Test response",
112+
})
113+
114+
await handler.completePrompt("Test prompt")
115+
116+
// Verify that generateContent was called with urlContext in tools
117+
expect(mockGenerateContent).toHaveBeenCalledWith(
118+
expect.objectContaining({
119+
config: expect.objectContaining({
120+
tools: expect.arrayContaining([{ urlContext: {} }]),
121+
}),
122+
}),
123+
)
124+
})
125+
})
126+
127+
describe("VertexHandler", () => {
128+
it("should NOT include urlContext tool even when enableUrlContext is true", async () => {
129+
const mockGenerateContentStream = vitest.fn()
130+
131+
const handler = new VertexHandler({
132+
vertexProjectId: "test-project",
133+
vertexRegion: "us-central1",
134+
apiModelId: "gemini-1.5-pro-001",
135+
enableUrlContext: true, // This should be ignored for Vertex
136+
})
137+
138+
// Replace the client with our mock
139+
handler["client"] = {
140+
models: {
141+
generateContentStream: mockGenerateContentStream,
142+
},
143+
} as any
144+
145+
// Setup mock to return an async generator
146+
mockGenerateContentStream.mockResolvedValue({
147+
[Symbol.asyncIterator]: async function* () {
148+
yield { text: "Test response" }
149+
yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 } }
150+
},
151+
})
152+
153+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
154+
155+
const stream = handler.createMessage("System prompt", messages)
156+
const chunks = []
157+
for await (const chunk of stream) {
158+
chunks.push(chunk)
159+
}
160+
161+
// Verify that generateContentStream was called WITHOUT urlContext in tools
162+
// even though enableUrlContext was true
163+
const callArgs = mockGenerateContentStream.mock.calls[0][0]
164+
if (callArgs.config.tools) {
165+
// If tools array exists, it should not contain urlContext
166+
expect(callArgs.config.tools).not.toContainEqual({ urlContext: {} })
167+
}
168+
})
169+
170+
it("should NOT include urlContext in completePrompt even when enableUrlContext is true", async () => {
171+
const mockGenerateContent = vitest.fn()
172+
173+
const handler = new VertexHandler({
174+
vertexProjectId: "test-project",
175+
vertexRegion: "us-central1",
176+
apiModelId: "gemini-1.5-pro-001",
177+
enableUrlContext: true, // This should be ignored for Vertex
178+
})
179+
180+
// Replace the client with our mock
181+
handler["client"] = {
182+
models: {
183+
generateContent: mockGenerateContent,
184+
},
185+
} as any
186+
187+
// Mock the response
188+
mockGenerateContent.mockResolvedValue({
189+
text: "Test response",
190+
})
191+
192+
await handler.completePrompt("Test prompt")
193+
194+
// Verify that generateContent was called WITHOUT urlContext in tools
195+
const callArgs = mockGenerateContent.mock.calls[0][0]
196+
if (callArgs.config.tools) {
197+
// If tools array exists, it should not contain urlContext
198+
expect(callArgs.config.tools).not.toContainEqual({ urlContext: {} })
199+
}
200+
})
201+
202+
it("should still include googleSearch tool when enableGrounding is true", async () => {
203+
const mockGenerateContentStream = vitest.fn()
204+
205+
const handler = new VertexHandler({
206+
vertexProjectId: "test-project",
207+
vertexRegion: "us-central1",
208+
apiModelId: "gemini-1.5-pro-001",
209+
enableUrlContext: true, // Should be ignored
210+
enableGrounding: true, // Should be respected
211+
})
212+
213+
// Replace the client with our mock
214+
handler["client"] = {
215+
models: {
216+
generateContentStream: mockGenerateContentStream,
217+
},
218+
} as any
219+
220+
// Setup mock to return an async generator
221+
mockGenerateContentStream.mockResolvedValue({
222+
[Symbol.asyncIterator]: async function* () {
223+
yield { text: "Test response" }
224+
yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 } }
225+
},
226+
})
227+
228+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
229+
230+
const stream = handler.createMessage("System prompt", messages)
231+
const chunks = []
232+
for await (const chunk of stream) {
233+
chunks.push(chunk)
234+
}
235+
236+
// Verify that googleSearch is included but urlContext is not
237+
const callArgs = mockGenerateContentStream.mock.calls[0][0]
238+
expect(callArgs.config.tools).toContainEqual({ googleSearch: {} })
239+
expect(callArgs.config.tools).not.toContainEqual({ urlContext: {} })
240+
})
241+
})
242+
243+
describe("Integration test - switching between providers", () => {
244+
it("should correctly handle urlContext based on provider type", async () => {
245+
const mockGenerateContentStream = vitest.fn()
246+
247+
// Setup mock to return an async generator
248+
mockGenerateContentStream.mockResolvedValue({
249+
[Symbol.asyncIterator]: async function* () {
250+
yield { text: "Test response" }
251+
yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 } }
252+
},
253+
})
254+
255+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
256+
257+
// Test with Gemini handler
258+
const geminiHandler = new GeminiHandler({
259+
geminiApiKey: "test-key",
260+
apiModelId: "gemini-1.5-flash",
261+
enableUrlContext: true,
262+
})
263+
geminiHandler["client"] = {
264+
models: { generateContentStream: mockGenerateContentStream },
265+
} as any
266+
267+
const geminiStream = geminiHandler.createMessage("System prompt", messages)
268+
for await (const chunk of geminiStream) {
269+
// Consume stream
270+
}
271+
272+
// Verify Gemini includes urlContext
273+
const geminiCall = mockGenerateContentStream.mock.calls[mockGenerateContentStream.mock.calls.length - 1][0]
274+
expect(geminiCall.config.tools).toContainEqual({ urlContext: {} })
275+
276+
// Clear mock calls
277+
mockGenerateContentStream.mockClear()
278+
279+
// Test with Vertex handler using same options
280+
const vertexHandler = new VertexHandler({
281+
vertexProjectId: "test-project",
282+
vertexRegion: "us-central1",
283+
apiModelId: "gemini-1.5-pro-001",
284+
enableUrlContext: true, // Same setting, but should be ignored
285+
})
286+
vertexHandler["client"] = {
287+
models: { generateContentStream: mockGenerateContentStream },
288+
} as any
289+
290+
const vertexStream = vertexHandler.createMessage("System prompt", messages)
291+
for await (const chunk of vertexStream) {
292+
// Consume stream
293+
}
294+
295+
// Verify Vertex does NOT include urlContext
296+
const vertexCall = mockGenerateContentStream.mock.calls[0][0]
297+
if (vertexCall.config.tools) {
298+
expect(vertexCall.config.tools).not.toContainEqual({ urlContext: {} })
299+
}
300+
})
301+
})
302+
})

‎src/api/providers/gemini.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,15 @@ type GeminiHandlerOptions = ApiHandlerOptions & {
2727

2828
export class GeminiHandler extends BaseProvider implements SingleCompletionHandler {
2929
protected options: ApiHandlerOptions
30+
private isVertex: boolean
3031

3132
private client: GoogleGenAI
3233

3334
constructor({ isVertex, ...options }: GeminiHandlerOptions) {
3435
super()
3536

3637
this.options = options
38+
this.isVertex = isVertex ?? false
3739

3840
const project = this.options.vertexProjectId ?? "not-provided"
3941
const location = this.options.vertexRegion ?? "not-provided"
@@ -70,7 +72,8 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
7072
const contents = messages.map(convertAnthropicMessageToGemini)
7173

7274
const tools: GenerateContentConfig["tools"] = []
73-
if (this.options.enableUrlContext) {
75+
// urlContext is only supported in regular Gemini, not Vertex AI
76+
if (this.options.enableUrlContext && !this.isVertex) {
7477
tools.push({ urlContext: {} })
7578
}
7679

@@ -214,7 +217,8 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
214217
const { id: model } = this.getModel()
215218

216219
const tools: GenerateContentConfig["tools"] = []
217-
if (this.options.enableUrlContext) {
220+
// urlContext is only supported in regular Gemini, not Vertex AI
221+
if (this.options.enableUrlContext && !this.isVertex) {
218222
tools.push({ urlContext: {} })
219223
}
220224
if (this.options.enableGrounding) {

0 commit comments

Comments
 (0)