fix: improve type safety in VertexAI custom base URL implementation

- Replace 'any' type with proper VertexOptions type in anthropic-vertex.ts
- Use isVertex property instead of fragile constructor.name comparison in gemini.ts
- Improve test assertions clarity by checking undefined explicitly
This commit is contained in:
daniel-lxs 2025-09-12 16:22:19 -05:00
parent a8a83113ff
commit 1698100ea2
No known key found for this signature in database
GPG key ID: 21C74479048B3AA6
4 changed files with 16 additions and 19 deletions

View file

@ -893,11 +893,8 @@ describe("VertexHandler", () => {
})
// Verify that AnthropicVertex was called without baseURL
expect(AnthropicVertex).toHaveBeenCalledWith(
expect.not.objectContaining({
baseURL: expect.anything(),
}),
)
const callArgs = (AnthropicVertex as any).mock.calls[0][0]
expect(callArgs.baseURL).toBeUndefined()
})
})
})

View file

@ -184,13 +184,8 @@ describe("VertexHandler", () => {
await handler.completePrompt("Test prompt")
// Verify that httpOptions is undefined when no custom URL
expect(mockGenerateContent).toHaveBeenCalledWith(
expect.objectContaining({
config: expect.objectContaining({
httpOptions: undefined,
}),
}),
)
const callArgs = mockGenerateContent.mock.calls[0][0]
expect(callArgs.config.httpOptions).toBeUndefined()
})
})
})

View file

@ -34,7 +34,14 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple
const projectId = this.options.vertexProjectId ?? "not-provided"
const region = this.options.vertexRegion ?? "us-east5"
const baseOptions: any = {
type VertexOptions = {
projectId: string
region: string
baseURL?: string
googleAuth?: GoogleAuth
}
const baseOptions: VertexOptions = {
projectId,
region,
}

View file

@ -27,6 +27,7 @@ type GeminiHandlerOptions = ApiHandlerOptions & {
export class GeminiHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private isVertex: boolean
private client: GoogleGenAI
@ -34,6 +35,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
super()
this.options = options
this.isVertex = isVertex ?? false
const project = this.options.vertexProjectId ?? "not-provided"
const location = this.options.vertexRegion ?? "not-provided"
@ -79,8 +81,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
}
// Use vertexBaseUrl if this is a Vertex handler, otherwise use googleGeminiBaseUrl
const baseUrl =
this.constructor.name === "VertexHandler" ? this.options.vertexBaseUrl : this.options.googleGeminiBaseUrl
const baseUrl = this.isVertex ? this.options.vertexBaseUrl : this.options.googleGeminiBaseUrl
const config: GenerateContentConfig = {
systemInstruction,
@ -225,10 +226,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
tools.push({ googleSearch: {} })
}
// Use vertexBaseUrl if this is a Vertex handler, otherwise use googleGeminiBaseUrl
const baseUrl =
this.constructor.name === "VertexHandler"
? this.options.vertexBaseUrl
: this.options.googleGeminiBaseUrl
const baseUrl = this.isVertex ? this.options.vertexBaseUrl : this.options.googleGeminiBaseUrl
const promptConfig: GenerateContentConfig = {
httpOptions: baseUrl ? { baseUrl } : undefined,