Skip to content

Commit 5220b36

Browse files
roomoteroomote
committed
feat: enable grounding features for Vertex AI
- Add enableUrlContext and enableGrounding options to vertexSchema in provider-settings.ts - Update Vertex.tsx UI component to include URL context and grounding checkboxes - Add comprehensive tests for the new grounding features in Vertex.spec.tsx - Reuse existing localization strings from Gemini provider Fixes #6776
1 parent c52fdc4 commit 5220b36

File tree

3 files changed

+244
-51
lines changed

3 files changed

+244
-51
lines changed

‎packages/types/src/provider-settings.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,8 @@ const vertexSchema = apiModelIdProviderModelSchema.extend({
132132
vertexJsonCredentials: z.string().optional(),
133133
vertexProjectId: z.string().optional(),
134134
vertexRegion: z.string().optional(),
135+
enableUrlContext: z.boolean().optional(),
136+
enableGrounding: z.boolean().optional(),
135137
})
136138

137139
const openAiSchema = baseProviderSettingsSchema.extend({

‎webview-ui/src/components/settings/providers/Vertex.tsx‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { useCallback } from "react"
2+
import { Checkbox } from "vscrui"
23
import { VSCodeLink, VSCodeTextField } from "@vscode/webview-ui-toolkit/react"
34

45
import { type ProviderSettings, VERTEX_REGIONS } from "@roo-code/types"
@@ -91,6 +92,28 @@ export const Vertex = ({ apiConfiguration, setApiConfigurationField }: VertexPro
9192
</SelectContent>
9293
</Select>
9394
</div>
95+
96+
<div className="mt-6">
97+
<Checkbox
98+
data-testid="checkbox-url-context"
99+
checked={!!apiConfiguration.enableUrlContext}
100+
onChange={(checked: boolean) => setApiConfigurationField("enableUrlContext", checked)}>
101+
{t("settings:providers.geminiParameters.urlContext.title")}
102+
</Checkbox>
103+
<div className="text-sm text-vscode-descriptionForeground mb-3 mt-1.5">
104+
{t("settings:providers.geminiParameters.urlContext.description")}
105+
</div>
106+
107+
<Checkbox
108+
data-testid="checkbox-grounding-search"
109+
checked={!!apiConfiguration.enableGrounding}
110+
onChange={(checked: boolean) => setApiConfigurationField("enableGrounding", checked)}>
111+
{t("settings:providers.geminiParameters.groundingSearch.title")}
112+
</Checkbox>
113+
<div className="text-sm text-vscode-descriptionForeground mb-3 mt-1.5">
114+
{t("settings:providers.geminiParameters.groundingSearch.description")}
115+
</div>
116+
</div>
94117
</>
95118
)
96119
}

‎webview-ui/src/components/settings/providers/__tests__/Vertex.spec.tsx‎

Lines changed: 219 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,61 +1,229 @@
1-
// Tests for VERTEX_REGIONS "global" region handling
1+
import { render, screen } from "@testing-library/react"
2+
import userEvent from "@testing-library/user-event"
3+
import { Vertex } from "../Vertex"
4+
import type { ProviderSettings } from "@roo-code/types"
5+
import { VERTEX_REGIONS } from "@roo-code/types"
26

3-
import { describe, it, expect } from "vitest"
4-
import { VERTEX_REGIONS } from "../../../../../../packages/types/src/providers/vertex"
7+
vi.mock("@vscode/webview-ui-toolkit/react", () => ({
8+
VSCodeTextField: ({ children, value, onInput, type }: any) => (
9+
<div>
10+
{children}
11+
<input type={type} value={value} onChange={(e) => onInput(e)} />
12+
</div>
13+
),
14+
VSCodeLink: ({ children, href }: any) => <a href={href}>{children}</a>,
15+
}))
516

6-
describe("VERTEX_REGIONS", () => {
7-
it('should include the "global" region as the first entry', () => {
8-
expect(VERTEX_REGIONS[0]).toEqual({ value: "global", label: "global" })
17+
vi.mock("vscrui", () => ({
18+
Checkbox: ({ children, checked, onChange, "data-testid": testId }: any) => (
19+
<label data-testid={testId}>
20+
<input type="checkbox" checked={checked} onChange={(e) => onChange(e.target.checked)} />
21+
{children}
22+
</label>
23+
),
24+
}))
25+
26+
vi.mock("@src/i18n/TranslationContext", () => ({
27+
useAppTranslation: () => ({ t: (key: string) => key }),
28+
}))
29+
30+
vi.mock("@src/components/ui", () => ({
31+
Select: ({ children, value, onValueChange }: any) => (
32+
<div data-value={value} data-onvaluechange={onValueChange}>
33+
{children}
34+
</div>
35+
),
36+
SelectContent: ({ children }: any) => <div>{children}</div>,
37+
SelectItem: ({ children, value }: any) => <div data-value={value}>{children}</div>,
38+
SelectTrigger: ({ children }: any) => <div>{children}</div>,
39+
SelectValue: ({ placeholder }: any) => <div>{placeholder}</div>,
40+
}))
41+
42+
describe("Vertex", () => {
43+
const defaultApiConfiguration: ProviderSettings = {
44+
vertexKeyFile: "",
45+
vertexJsonCredentials: "",
46+
vertexProjectId: "",
47+
vertexRegion: "",
48+
enableUrlContext: false,
49+
enableGrounding: false,
50+
}
51+
52+
const mockSetApiConfigurationField = vi.fn()
53+
54+
beforeEach(() => {
55+
vi.clearAllMocks()
956
})
1057

11-
it('should contain "global" region exactly once', () => {
12-
const globalRegions = VERTEX_REGIONS.filter((r: { value: string; label: string }) => r.value === "global")
13-
expect(globalRegions).toHaveLength(1)
58+
describe("VERTEX_REGIONS", () => {
59+
it('should include the "global" region as the first entry', () => {
60+
expect(VERTEX_REGIONS[0]).toEqual({ value: "global", label: "global" })
61+
})
62+
63+
it('should contain "global" region exactly once', () => {
64+
const globalRegions = VERTEX_REGIONS.filter((r: { value: string; label: string }) => r.value === "global")
65+
expect(globalRegions).toHaveLength(1)
66+
})
67+
68+
it('should contain all expected regions including "global"', () => {
69+
// The expected list is the imported VERTEX_REGIONS itself
70+
expect(VERTEX_REGIONS).toEqual([
71+
{ value: "global", label: "global" },
72+
{ value: "us-central1", label: "us-central1" },
73+
{ value: "us-east1", label: "us-east1" },
74+
{ value: "us-east4", label: "us-east4" },
75+
{ value: "us-east5", label: "us-east5" },
76+
{ value: "us-west1", label: "us-west1" },
77+
{ value: "us-west2", label: "us-west2" },
78+
{ value: "us-west3", label: "us-west3" },
79+
{ value: "us-west4", label: "us-west4" },
80+
{ value: "northamerica-northeast1", label: "northamerica-northeast1" },
81+
{ value: "northamerica-northeast2", label: "northamerica-northeast2" },
82+
{ value: "southamerica-east1", label: "southamerica-east1" },
83+
{ value: "europe-west1", label: "europe-west1" },
84+
{ value: "europe-west2", label: "europe-west2" },
85+
{ value: "europe-west3", label: "europe-west3" },
86+
{ value: "europe-west4", label: "europe-west4" },
87+
{ value: "europe-west6", label: "europe-west6" },
88+
{ value: "europe-central2", label: "europe-central2" },
89+
{ value: "asia-east1", label: "asia-east1" },
90+
{ value: "asia-east2", label: "asia-east2" },
91+
{ value: "asia-northeast1", label: "asia-northeast1" },
92+
{ value: "asia-northeast2", label: "asia-northeast2" },
93+
{ value: "asia-northeast3", label: "asia-northeast3" },
94+
{ value: "asia-south1", label: "asia-south1" },
95+
{ value: "asia-south2", label: "asia-south2" },
96+
{ value: "asia-southeast1", label: "asia-southeast1" },
97+
{ value: "asia-southeast2", label: "asia-southeast2" },
98+
{ value: "australia-southeast1", label: "australia-southeast1" },
99+
{ value: "australia-southeast2", label: "australia-southeast2" },
100+
{ value: "me-west1", label: "me-west1" },
101+
{ value: "me-central1", label: "me-central1" },
102+
{ value: "africa-south1", label: "africa-south1" },
103+
])
104+
})
105+
106+
it('should contain "asia-east1" region exactly once', () => {
107+
const asiaEast1Regions = VERTEX_REGIONS.filter(
108+
(r: { value: string; label: string }) => r.value === "asia-east1" && r.label === "asia-east1",
109+
)
110+
expect(asiaEast1Regions).toHaveLength(1)
111+
expect(asiaEast1Regions[0]).toEqual({ value: "asia-east1", label: "asia-east1" })
112+
})
14113
})
15114

16-
it('should contain all expected regions including "global"', () => {
17-
// The expected list is the imported VERTEX_REGIONS itself
18-
expect(VERTEX_REGIONS).toEqual([
19-
{ value: "global", label: "global" },
20-
{ value: "us-central1", label: "us-central1" },
21-
{ value: "us-east1", label: "us-east1" },
22-
{ value: "us-east4", label: "us-east4" },
23-
{ value: "us-east5", label: "us-east5" },
24-
{ value: "us-west1", label: "us-west1" },
25-
{ value: "us-west2", label: "us-west2" },
26-
{ value: "us-west3", label: "us-west3" },
27-
{ value: "us-west4", label: "us-west4" },
28-
{ value: "northamerica-northeast1", label: "northamerica-northeast1" },
29-
{ value: "northamerica-northeast2", label: "northamerica-northeast2" },
30-
{ value: "southamerica-east1", label: "southamerica-east1" },
31-
{ value: "europe-west1", label: "europe-west1" },
32-
{ value: "europe-west2", label: "europe-west2" },
33-
{ value: "europe-west3", label: "europe-west3" },
34-
{ value: "europe-west4", label: "europe-west4" },
35-
{ value: "europe-west6", label: "europe-west6" },
36-
{ value: "europe-central2", label: "europe-central2" },
37-
{ value: "asia-east1", label: "asia-east1" },
38-
{ value: "asia-east2", label: "asia-east2" },
39-
{ value: "asia-northeast1", label: "asia-northeast1" },
40-
{ value: "asia-northeast2", label: "asia-northeast2" },
41-
{ value: "asia-northeast3", label: "asia-northeast3" },
42-
{ value: "asia-south1", label: "asia-south1" },
43-
{ value: "asia-south2", label: "asia-south2" },
44-
{ value: "asia-southeast1", label: "asia-southeast1" },
45-
{ value: "asia-southeast2", label: "asia-southeast2" },
46-
{ value: "australia-southeast1", label: "australia-southeast1" },
47-
{ value: "australia-southeast2", label: "australia-southeast2" },
48-
{ value: "me-west1", label: "me-west1" },
49-
{ value: "me-central1", label: "me-central1" },
50-
{ value: "africa-south1", label: "africa-south1" },
51-
])
115+
describe("URL Context Checkbox", () => {
116+
it("should render URL context checkbox unchecked by default", () => {
117+
render(
118+
<Vertex
119+
apiConfiguration={defaultApiConfiguration}
120+
setApiConfigurationField={mockSetApiConfigurationField}
121+
/>,
122+
)
123+
124+
const urlContextCheckbox = screen.getByTestId("checkbox-url-context")
125+
const checkbox = urlContextCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
126+
expect(checkbox.checked).toBe(false)
127+
})
128+
129+
it("should render URL context checkbox checked when enableUrlContext is true", () => {
130+
const apiConfiguration = { ...defaultApiConfiguration, enableUrlContext: true }
131+
render(
132+
<Vertex apiConfiguration={apiConfiguration} setApiConfigurationField={mockSetApiConfigurationField} />,
133+
)
134+
135+
const urlContextCheckbox = screen.getByTestId("checkbox-url-context")
136+
const checkbox = urlContextCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
137+
expect(checkbox.checked).toBe(true)
138+
})
139+
140+
it("should call setApiConfigurationField with correct parameters when URL context checkbox is toggled", async () => {
141+
const user = userEvent.setup()
142+
render(
143+
<Vertex
144+
apiConfiguration={defaultApiConfiguration}
145+
setApiConfigurationField={mockSetApiConfigurationField}
146+
/>,
147+
)
148+
149+
const urlContextCheckbox = screen.getByTestId("checkbox-url-context")
150+
const checkbox = urlContextCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
151+
152+
await user.click(checkbox)
153+
154+
expect(mockSetApiConfigurationField).toHaveBeenCalledWith("enableUrlContext", true)
155+
})
52156
})
53157

54-
it('should contain "asia-east1" region exactly once', () => {
55-
const asiaEast1Regions = VERTEX_REGIONS.filter(
56-
(r: { value: string; label: string }) => r.value === "asia-east1" && r.label === "asia-east1",
57-
)
58-
expect(asiaEast1Regions).toHaveLength(1)
59-
expect(asiaEast1Regions[0]).toEqual({ value: "asia-east1", label: "asia-east1" })
158+
describe("Grounding with Google Search Checkbox", () => {
159+
it("should render grounding search checkbox unchecked by default", () => {
160+
render(
161+
<Vertex
162+
apiConfiguration={defaultApiConfiguration}
163+
setApiConfigurationField={mockSetApiConfigurationField}
164+
/>,
165+
)
166+
167+
const groundingCheckbox = screen.getByTestId("checkbox-grounding-search")
168+
const checkbox = groundingCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
169+
expect(checkbox.checked).toBe(false)
170+
})
171+
172+
it("should render grounding search checkbox checked when enableGrounding is true", () => {
173+
const apiConfiguration = { ...defaultApiConfiguration, enableGrounding: true }
174+
render(
175+
<Vertex apiConfiguration={apiConfiguration} setApiConfigurationField={mockSetApiConfigurationField} />,
176+
)
177+
178+
const groundingCheckbox = screen.getByTestId("checkbox-grounding-search")
179+
const checkbox = groundingCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
180+
expect(checkbox.checked).toBe(true)
181+
})
182+
183+
it("should call setApiConfigurationField with correct parameters when grounding search checkbox is toggled", async () => {
184+
const user = userEvent.setup()
185+
render(
186+
<Vertex
187+
apiConfiguration={defaultApiConfiguration}
188+
setApiConfigurationField={mockSetApiConfigurationField}
189+
/>,
190+
)
191+
192+
const groundingCheckbox = screen.getByTestId("checkbox-grounding-search")
193+
const checkbox = groundingCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
194+
195+
await user.click(checkbox)
196+
197+
expect(mockSetApiConfigurationField).toHaveBeenCalledWith("enableGrounding", true)
198+
})
199+
})
200+
201+
describe("Both checkboxes interaction", () => {
202+
it("should be able to toggle both checkboxes independently", async () => {
203+
const user = userEvent.setup()
204+
render(
205+
<Vertex
206+
apiConfiguration={defaultApiConfiguration}
207+
setApiConfigurationField={mockSetApiConfigurationField}
208+
/>,
209+
)
210+
211+
const urlContextCheckbox = screen.getByTestId("checkbox-url-context")
212+
const urlCheckbox = urlContextCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
213+
214+
const groundingCheckbox = screen.getByTestId("checkbox-grounding-search")
215+
const groundCheckbox = groundingCheckbox.querySelector("input[type='checkbox']") as HTMLInputElement
216+
217+
// Toggle URL context
218+
await user.click(urlCheckbox)
219+
expect(mockSetApiConfigurationField).toHaveBeenCalledWith("enableUrlContext", true)
220+
221+
// Toggle grounding
222+
await user.click(groundCheckbox)
223+
expect(mockSetApiConfigurationField).toHaveBeenCalledWith("enableGrounding", true)
224+
225+
// Both should have been called
226+
expect(mockSetApiConfigurationField).toHaveBeenCalledTimes(2)
227+
})
60228
})
61229
})

0 commit comments

Comments
 (0)