mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-06 08:18:39 +00:00
feat: add custom base URL support for VertexAI provider
- Add vertexBaseUrl field to provider settings schema - Update VertexHandler to use custom base URL via GeminiHandler - Update AnthropicVertexHandler to use custom base URL - Add UI checkbox and input field for custom base URL in Vertex.tsx - Add comprehensive tests for the new functionality Fixes #7899
This commit is contained in:
parent
8fee3127ff
commit
a8a83113ff
6 changed files with 197 additions and 10 deletions
|
|
@ -168,6 +168,7 @@ const vertexSchema = apiModelIdProviderModelSchema.extend({
|
|||
vertexJsonCredentials: z.string().optional(),
|
||||
vertexProjectId: z.string().optional(),
|
||||
vertexRegion: z.string().optional(),
|
||||
vertexBaseUrl: z.string().optional(),
|
||||
enableUrlContext: z.boolean().optional(),
|
||||
enableGrounding: z.boolean().optional(),
|
||||
})
|
||||
|
|
|
|||
|
|
@ -809,4 +809,95 @@ describe("VertexHandler", () => {
|
|||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("custom base URL", () => {
|
||||
it("should use custom base URL when provided with JSON credentials", () => {
|
||||
const customBaseUrl = "https://custom-vertex-endpoint.example.com"
|
||||
|
||||
const handler = new AnthropicVertexHandler({
|
||||
apiModelId: "claude-3-5-sonnet-v2@20241022",
|
||||
vertexProjectId: "test-project",
|
||||
vertexRegion: "us-central1",
|
||||
vertexBaseUrl: customBaseUrl,
|
||||
vertexJsonCredentials: JSON.stringify({
|
||||
type: "service_account",
|
||||
project_id: "test-project",
|
||||
private_key_id: "key-id",
|
||||
private_key: "-----BEGIN PRIVATE KEY-----\ntest\n-----END PRIVATE KEY-----\n",
|
||||
client_email: "test@test.iam.gserviceaccount.com",
|
||||
client_id: "123456789",
|
||||
auth_uri: "https://accounts.google.com/o/oauth2/auth",
|
||||
token_uri: "https://oauth2.googleapis.com/token",
|
||||
auth_provider_x509_cert_url: "https://www.googleapis.com/oauth2/v1/certs",
|
||||
client_x509_cert_url:
|
||||
"https://www.googleapis.com/robot/v1/metadata/x509/test%40test.iam.gserviceaccount.com",
|
||||
}),
|
||||
})
|
||||
|
||||
// Verify that AnthropicVertex was called with baseURL
|
||||
expect(AnthropicVertex).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: customBaseUrl,
|
||||
projectId: "test-project",
|
||||
region: "us-central1",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should use custom base URL when provided with key file", () => {
|
||||
const customBaseUrl = "https://custom-vertex-endpoint.example.com"
|
||||
|
||||
const handler = new AnthropicVertexHandler({
|
||||
apiModelId: "claude-3-5-sonnet-v2@20241022",
|
||||
vertexProjectId: "test-project",
|
||||
vertexRegion: "us-central1",
|
||||
vertexBaseUrl: customBaseUrl,
|
||||
vertexKeyFile: "/path/to/keyfile.json",
|
||||
})
|
||||
|
||||
// Verify that AnthropicVertex was called with baseURL
|
||||
expect(AnthropicVertex).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: customBaseUrl,
|
||||
projectId: "test-project",
|
||||
region: "us-central1",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should use custom base URL when provided without credentials", () => {
|
||||
const customBaseUrl = "https://custom-vertex-endpoint.example.com"
|
||||
|
||||
const handler = new AnthropicVertexHandler({
|
||||
apiModelId: "claude-3-5-sonnet-v2@20241022",
|
||||
vertexProjectId: "test-project",
|
||||
vertexRegion: "us-central1",
|
||||
vertexBaseUrl: customBaseUrl,
|
||||
})
|
||||
|
||||
// Verify that AnthropicVertex was called with baseURL
|
||||
expect(AnthropicVertex).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: customBaseUrl,
|
||||
projectId: "test-project",
|
||||
region: "us-central1",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should not include baseURL when no custom URL is provided", () => {
|
||||
const handler = new AnthropicVertexHandler({
|
||||
apiModelId: "claude-3-5-sonnet-v2@20241022",
|
||||
vertexProjectId: "test-project",
|
||||
vertexRegion: "us-central1",
|
||||
})
|
||||
|
||||
// Verify that AnthropicVertex was called without baseURL
|
||||
expect(AnthropicVertex).toHaveBeenCalledWith(
|
||||
expect.not.objectContaining({
|
||||
baseURL: expect.anything(),
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -138,4 +138,59 @@ describe("VertexHandler", () => {
|
|||
expect(modelInfo.info.contextWindow).toBe(1048576)
|
||||
})
|
||||
})
|
||||
|
||||
describe("custom base URL", () => {
|
||||
it("should use custom base URL when provided", async () => {
|
||||
const customBaseUrl = "https://custom-vertex-endpoint.example.com"
|
||||
|
||||
handler = new VertexHandler({
|
||||
apiModelId: "gemini-1.5-pro-001",
|
||||
vertexProjectId: "test-project",
|
||||
vertexRegion: "us-central1",
|
||||
vertexBaseUrl: customBaseUrl,
|
||||
})
|
||||
|
||||
// Mock the generateContent method
|
||||
const mockGenerateContent = vitest.fn().mockResolvedValue({
|
||||
text: "Test response with custom URL",
|
||||
})
|
||||
handler["client"].models.generateContent = mockGenerateContent
|
||||
|
||||
await handler.completePrompt("Test prompt")
|
||||
|
||||
// Verify that the custom base URL was passed in the config
|
||||
expect(mockGenerateContent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
config: expect.objectContaining({
|
||||
httpOptions: { baseUrl: customBaseUrl },
|
||||
}),
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should not include httpOptions when no custom base URL is provided", async () => {
|
||||
handler = new VertexHandler({
|
||||
apiModelId: "gemini-1.5-pro-001",
|
||||
vertexProjectId: "test-project",
|
||||
vertexRegion: "us-central1",
|
||||
})
|
||||
|
||||
// Mock the generateContent method
|
||||
const mockGenerateContent = vitest.fn().mockResolvedValue({
|
||||
text: "Test response without custom URL",
|
||||
})
|
||||
handler["client"].models.generateContent = mockGenerateContent
|
||||
|
||||
await handler.completePrompt("Test prompt")
|
||||
|
||||
// Verify that httpOptions is undefined when no custom URL
|
||||
expect(mockGenerateContent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
config: expect.objectContaining({
|
||||
httpOptions: undefined,
|
||||
}),
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -34,10 +34,19 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple
|
|||
const projectId = this.options.vertexProjectId ?? "not-provided"
|
||||
const region = this.options.vertexRegion ?? "us-east5"
|
||||
|
||||
const baseOptions: any = {
|
||||
projectId,
|
||||
region,
|
||||
}
|
||||
|
||||
// Add custom base URL if provided
|
||||
if (this.options.vertexBaseUrl) {
|
||||
baseOptions.baseURL = this.options.vertexBaseUrl
|
||||
}
|
||||
|
||||
if (this.options.vertexJsonCredentials) {
|
||||
this.client = new AnthropicVertex({
|
||||
projectId,
|
||||
region,
|
||||
...baseOptions,
|
||||
googleAuth: new GoogleAuth({
|
||||
scopes: ["https://www.googleapis.com/auth/cloud-platform"],
|
||||
credentials: safeJsonParse<JWTInput>(this.options.vertexJsonCredentials, undefined),
|
||||
|
|
@ -45,15 +54,14 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple
|
|||
})
|
||||
} else if (this.options.vertexKeyFile) {
|
||||
this.client = new AnthropicVertex({
|
||||
projectId,
|
||||
region,
|
||||
...baseOptions,
|
||||
googleAuth: new GoogleAuth({
|
||||
scopes: ["https://www.googleapis.com/auth/cloud-platform"],
|
||||
keyFile: this.options.vertexKeyFile,
|
||||
}),
|
||||
})
|
||||
} else {
|
||||
this.client = new AnthropicVertex({ projectId, region })
|
||||
this.client = new AnthropicVertex(baseOptions)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -78,9 +78,13 @@ 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 config: GenerateContentConfig = {
|
||||
systemInstruction,
|
||||
httpOptions: this.options.googleGeminiBaseUrl ? { baseUrl: this.options.googleGeminiBaseUrl } : undefined,
|
||||
httpOptions: baseUrl ? { baseUrl } : undefined,
|
||||
thinkingConfig,
|
||||
maxOutputTokens: this.options.modelMaxTokens ?? maxTokens ?? undefined,
|
||||
temperature: this.options.modelTemperature ?? 0,
|
||||
|
|
@ -220,10 +224,14 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
|
|||
if (this.options.enableGrounding) {
|
||||
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 promptConfig: GenerateContentConfig = {
|
||||
httpOptions: this.options.googleGeminiBaseUrl
|
||||
? { baseUrl: this.options.googleGeminiBaseUrl }
|
||||
: undefined,
|
||||
httpOptions: baseUrl ? { baseUrl } : undefined,
|
||||
temperature: this.options.modelTemperature ?? 0,
|
||||
...(tools.length > 0 ? { tools } : {}),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
import { useCallback } from "react"
|
||||
import { useCallback, useState } from "react"
|
||||
import { Checkbox } from "vscrui"
|
||||
import { VSCodeLink, VSCodeTextField } from "@vscode/webview-ui-toolkit/react"
|
||||
|
||||
|
|
@ -18,6 +18,8 @@ type VertexProps = {
|
|||
export const Vertex = ({ apiConfiguration, setApiConfigurationField, fromWelcomeView }: VertexProps) => {
|
||||
const { t } = useAppTranslation()
|
||||
|
||||
const [vertexBaseUrlSelected, setVertexBaseUrlSelected] = useState(!!apiConfiguration?.vertexBaseUrl)
|
||||
|
||||
const handleInputChange = useCallback(
|
||||
<K extends keyof ProviderSettings, E>(
|
||||
field: K,
|
||||
|
|
@ -94,6 +96,28 @@ export const Vertex = ({ apiConfiguration, setApiConfigurationField, fromWelcome
|
|||
</Select>
|
||||
</div>
|
||||
|
||||
<div className="mt-4">
|
||||
<Checkbox
|
||||
checked={vertexBaseUrlSelected}
|
||||
onChange={(checked: boolean) => {
|
||||
setVertexBaseUrlSelected(checked)
|
||||
if (!checked) {
|
||||
setApiConfigurationField("vertexBaseUrl", "")
|
||||
}
|
||||
}}>
|
||||
{t("settings:providers.useCustomBaseUrl")}
|
||||
</Checkbox>
|
||||
{vertexBaseUrlSelected && (
|
||||
<VSCodeTextField
|
||||
value={apiConfiguration?.vertexBaseUrl || ""}
|
||||
type="url"
|
||||
onInput={handleInputChange("vertexBaseUrl")}
|
||||
placeholder="https://us-central1-aiplatform.googleapis.com"
|
||||
className="w-full mt-1"
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{!fromWelcomeView && apiConfiguration.apiModelId?.startsWith("gemini") && (
|
||||
<div className="mt-6">
|
||||
<Checkbox
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue