Skip to content

Commit 1819bd1

Browse files
roomoteroomote
committed
feat: add Vertex AI as embedder provider for codebase indexing
- Add "vertex" to EmbedderProvider type - Add Vertex AI embedding models to EMBEDDING_MODEL_PROFILES - Create VertexEmbedder implementation using OpenAI-compatible approach - Update service factory to handle vertex provider - Add vertexOptions to CodeIndexConfig interface - Update CodeIndexPopover UI to include Vertex AI section - Add translation keys for Vertex AI - Add VERTEX_MAX_ITEM_TOKENS constant - Add comprehensive tests for VertexEmbedder Closes #6300
1 parent 342ee70 commit 1819bd1

File tree

10 files changed

+398
-3
lines changed

10 files changed

+398
-3
lines changed

‎src/i18n/locales/en/embeddings.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
"openAiCompatibleConfigMissing": "OpenAI Compatible configuration missing for embedder creation",
4848
"geminiConfigMissing": "Gemini configuration missing for embedder creation",
4949
"mistralConfigMissing": "Mistral configuration missing for embedder creation",
50+
"vertexConfigMissing": "Vertex AI configuration missing for embedder creation",
5051
"invalidEmbedderType": "Invalid embedder type configured: {{embedderProvider}}",
5152
"vectorDimensionNotDeterminedOpenAiCompatible": "Could not determine vector dimension for model '{{modelId}}' with provider '{{provider}}'. Please ensure the 'Embedding Dimension' is correctly set in the OpenAI-Compatible provider settings.",
5253
"vectorDimensionNotDetermined": "Could not determine vector dimension for model '{{modelId}}' with provider '{{provider}}'. Check model profiles or configuration.",

‎src/services/code-index/constants/index.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,3 +29,6 @@ export const BATCH_PROCESSING_CONCURRENCY = 10
2929

3030
/**Gemini Embedder */
3131
export const GEMINI_MAX_ITEM_TOKENS = 2048
32+
33+
/**Vertex AI Embedder */
34+
export const VERTEX_MAX_ITEM_TOKENS = 2048

‎src/services/code-index/embedders/__tests__/vertex.spec.ts‎

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
import { vitest, describe, it, expect, beforeEach } from "vitest"
2+
import type { MockedClass } from "vitest"
3+
import { VertexEmbedder } from "../vertex"
4+
import { OpenAICompatibleEmbedder } from "../openai-compatible"
5+
6+
// Mock the OpenAICompatibleEmbedder
7+
vitest.mock("../openai-compatible")
8+
9+
// Mock TelemetryService
10+
vitest.mock("@roo-code/telemetry", () => ({
11+
TelemetryService: {
12+
instance: {
13+
captureEvent: vitest.fn(),
14+
},
15+
},
16+
}))
17+
18+
const MockedOpenAICompatibleEmbedder = OpenAICompatibleEmbedder as MockedClass<typeof OpenAICompatibleEmbedder>
19+
20+
describe("VertexEmbedder", () => {
21+
let embedder: VertexEmbedder
22+
23+
beforeEach(() => {
24+
vitest.clearAllMocks()
25+
})
26+
27+
describe("constructor", () => {
28+
it("should create an instance with default model when no model specified", () => {
29+
// Arrange
30+
const apiKey = "test-vertex-api-key"
31+
32+
// Act
33+
embedder = new VertexEmbedder(apiKey)
34+
35+
// Assert
36+
expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith(
37+
"https://generativelanguage.googleapis.com/v1beta/openai/",
38+
apiKey,
39+
"text-embedding-004",
40+
2048,
41+
)
42+
})
43+
44+
it("should create an instance with specified model", () => {
45+
// Arrange
46+
const apiKey = "test-vertex-api-key"
47+
const modelId = "text-multilingual-embedding-002"
48+
49+
// Act
50+
embedder = new VertexEmbedder(apiKey, modelId)
51+
52+
// Assert
53+
expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith(
54+
"https://generativelanguage.googleapis.com/v1beta/openai/",
55+
apiKey,
56+
"text-multilingual-embedding-002",
57+
2048,
58+
)
59+
})
60+
61+
it("should throw error when API key is not provided", () => {
62+
// Act & Assert
63+
expect(() => new VertexEmbedder("")).toThrow("validation.apiKeyRequired")
64+
expect(() => new VertexEmbedder(null as any)).toThrow("validation.apiKeyRequired")
65+
expect(() => new VertexEmbedder(undefined as any)).toThrow("validation.apiKeyRequired")
66+
})
67+
})
68+
69+
describe("embedderInfo", () => {
70+
it("should return correct embedder info", () => {
71+
// Arrange
72+
embedder = new VertexEmbedder("test-api-key")
73+
74+
// Act
75+
const info = embedder.embedderInfo
76+
77+
// Assert
78+
expect(info).toEqual({
79+
name: "vertex",
80+
})
81+
})
82+
83+
describe("createEmbeddings", () => {
84+
let mockCreateEmbeddings: any
85+
86+
beforeEach(() => {
87+
mockCreateEmbeddings = vitest.fn()
88+
MockedOpenAICompatibleEmbedder.prototype.createEmbeddings = mockCreateEmbeddings
89+
})
90+
91+
it("should use instance model when no model parameter provided", async () => {
92+
// Arrange
93+
embedder = new VertexEmbedder("test-api-key")
94+
const texts = ["test text 1", "test text 2"]
95+
const mockResponse = {
96+
embeddings: [
97+
[0.1, 0.2],
98+
[0.3, 0.4],
99+
],
100+
}
101+
mockCreateEmbeddings.mockResolvedValue(mockResponse)
102+
103+
// Act
104+
const result = await embedder.createEmbeddings(texts)
105+
106+
// Assert
107+
expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "text-embedding-004")
108+
expect(result).toEqual(mockResponse)
109+
})
110+
111+
it("should use provided model parameter when specified", async () => {
112+
// Arrange
113+
embedder = new VertexEmbedder("test-api-key", "textembedding-gecko@003")
114+
const texts = ["test text 1", "test text 2"]
115+
const mockResponse = {
116+
embeddings: [
117+
[0.1, 0.2],
118+
[0.3, 0.4],
119+
],
120+
}
121+
mockCreateEmbeddings.mockResolvedValue(mockResponse)
122+
123+
// Act
124+
const result = await embedder.createEmbeddings(texts, "text-multilingual-embedding-002")
125+
126+
// Assert
127+
expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "text-multilingual-embedding-002")
128+
expect(result).toEqual(mockResponse)
129+
})
130+
131+
it("should handle errors from OpenAICompatibleEmbedder", async () => {
132+
// Arrange
133+
embedder = new VertexEmbedder("test-api-key")
134+
const texts = ["test text"]
135+
const error = new Error("Embedding failed")
136+
mockCreateEmbeddings.mockRejectedValue(error)
137+
138+
// Act & Assert
139+
await expect(embedder.createEmbeddings(texts)).rejects.toThrow("Embedding failed")
140+
})
141+
})
142+
})
143+
144+
describe("validateConfiguration", () => {
145+
let mockValidateConfiguration: any
146+
147+
beforeEach(() => {
148+
mockValidateConfiguration = vitest.fn()
149+
MockedOpenAICompatibleEmbedder.prototype.validateConfiguration = mockValidateConfiguration
150+
})
151+
152+
it("should delegate validation to OpenAICompatibleEmbedder", async () => {
153+
// Arrange
154+
embedder = new VertexEmbedder("test-api-key")
155+
mockValidateConfiguration.mockResolvedValue({ valid: true })
156+
157+
// Act
158+
const result = await embedder.validateConfiguration()
159+
160+
// Assert
161+
expect(mockValidateConfiguration).toHaveBeenCalled()
162+
expect(result).toEqual({ valid: true })
163+
})
164+
165+
it("should pass through validation errors from OpenAICompatibleEmbedder", async () => {
166+
// Arrange
167+
embedder = new VertexEmbedder("test-api-key")
168+
mockValidateConfiguration.mockResolvedValue({
169+
valid: false,
170+
error: "embeddings:validation.authenticationFailed",
171+
})
172+
173+
// Act
174+
const result = await embedder.validateConfiguration()
175+
176+
// Assert
177+
expect(mockValidateConfiguration).toHaveBeenCalled()
178+
expect(result).toEqual({
179+
valid: false,
180+
error: "embeddings:validation.authenticationFailed",
181+
})
182+
})
183+
184+
it("should handle validation exceptions", async () => {
185+
// Arrange
186+
embedder = new VertexEmbedder("test-api-key")
187+
mockValidateConfiguration.mockRejectedValue(new Error("Validation failed"))
188+
189+
// Act & Assert
190+
await expect(embedder.validateConfiguration()).rejects.toThrow("Validation failed")
191+
})
192+
})
193+
})

‎src/services/code-index/embedders/vertex.ts‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
import { OpenAICompatibleEmbedder } from "./openai-compatible"
2+
import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder"
3+
import { VERTEX_MAX_ITEM_TOKENS } from "../constants"
4+
import { t } from "../../../i18n"
5+
import { TelemetryEventName } from "@roo-code/types"
6+
import { TelemetryService } from "@roo-code/telemetry"
7+
8+
/**
9+
* Vertex AI embedder implementation that wraps the OpenAI Compatible embedder
10+
* with configuration for Google's Vertex AI embedding API.
11+
*
12+
* Supported models:
13+
* - text-embedding-004 (dimension: 768)
14+
* - text-multilingual-embedding-002 (dimension: 768)
15+
* - textembedding-gecko@003 (dimension: 768)
16+
* - textembedding-gecko-multilingual@001 (dimension: 768)
17+
*/
18+
export class VertexEmbedder implements IEmbedder {
19+
private readonly openAICompatibleEmbedder: OpenAICompatibleEmbedder
20+
private static readonly VERTEX_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
21+
private static readonly DEFAULT_MODEL = "text-embedding-004"
22+
private readonly modelId: string
23+
24+
/**
25+
* Creates a new Vertex AI embedder
26+
* @param apiKey The Google AI API key for authentication
27+
* @param modelId The model ID to use (defaults to text-embedding-004)
28+
*/
29+
constructor(apiKey: string, modelId?: string) {
30+
if (!apiKey) {
31+
throw new Error(t("embeddings:validation.apiKeyRequired"))
32+
}
33+
34+
// Use provided model or default
35+
this.modelId = modelId || VertexEmbedder.DEFAULT_MODEL
36+
37+
// Create an OpenAI Compatible embedder with Vertex AI's configuration
38+
this.openAICompatibleEmbedder = new OpenAICompatibleEmbedder(
39+
VertexEmbedder.VERTEX_BASE_URL,
40+
apiKey,
41+
this.modelId,
42+
VERTEX_MAX_ITEM_TOKENS,
43+
)
44+
}
45+
46+
/**
47+
* Creates embeddings for the given texts using Vertex AI's embedding API
48+
* @param texts Array of text strings to embed
49+
* @param model Optional model identifier (uses constructor model if not provided)
50+
* @returns Promise resolving to embedding response
51+
*/
52+
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
53+
try {
54+
// Use the provided model or fall back to the instance's model
55+
const modelToUse = model || this.modelId
56+
return await this.openAICompatibleEmbedder.createEmbeddings(texts, modelToUse)
57+
} catch (error) {
58+
TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
59+
error: error instanceof Error ? error.message : String(error),
60+
stack: error instanceof Error ? error.stack : undefined,
61+
location: "VertexEmbedder:createEmbeddings",
62+
})
63+
throw error
64+
}
65+
}
66+
67+
/**
68+
* Validates the Vertex AI embedder configuration by delegating to the underlying OpenAI-compatible embedder
69+
* @returns Promise resolving to validation result with success status and optional error message
70+
*/
71+
async validateConfiguration(): Promise<{ valid: boolean; error?: string }> {
72+
try {
73+
// Delegate validation to the OpenAI-compatible embedder
74+
// The error messages will be specific to Vertex AI since we're using Vertex AI's base URL
75+
return await this.openAICompatibleEmbedder.validateConfiguration()
76+
} catch (error) {
77+
TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
78+
error: error instanceof Error ? error.message : String(error),
79+
stack: error instanceof Error ? error.stack : undefined,
80+
location: "VertexEmbedder:validateConfiguration",
81+
})
82+
throw error
83+
}
84+
}
85+
86+
/**
87+
* Returns information about this embedder
88+
*/
89+
get embedderInfo(): EmbedderInfo {
90+
return {
91+
name: "vertex",
92+
}
93+
}
94+
}

‎src/services/code-index/interfaces/config.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ export interface CodeIndexConfig {
1414
openAiCompatibleOptions?: { baseUrl: string; apiKey: string }
1515
geminiOptions?: { apiKey: string }
1616
mistralOptions?: { apiKey: string }
17+
vertexOptions?: { apiKey: string }
1718
qdrantUrl?: string
1819
qdrantApiKey?: string
1920
searchMinScore?: number
@@ -35,6 +36,7 @@ export type PreviousConfigSnapshot = {
3536
openAiCompatibleApiKey?: string
3637
geminiApiKey?: string
3738
mistralApiKey?: string
39+
vertexApiKey?: string
3840
qdrantUrl?: string
3941
qdrantApiKey?: string
4042
}

‎src/services/code-index/interfaces/embedder.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ export interface EmbeddingResponse {
2828
}
2929
}
3030

31-
export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral"
31+
export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vertex"
3232

3333
export interface EmbedderInfo {
3434
name: AvailableEmbedders

‎src/services/code-index/service-factory.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { CodeIndexOllamaEmbedder } from "./embedders/ollama"
44
import { OpenAICompatibleEmbedder } from "./embedders/openai-compatible"
55
import { GeminiEmbedder } from "./embedders/gemini"
66
import { MistralEmbedder } from "./embedders/mistral"
7+
import { VertexEmbedder } from "./embedders/vertex"
78
import { EmbedderProvider, getDefaultModelId, getModelDimension } from "../../shared/embeddingModels"
89
import { QdrantVectorStore } from "./vector-store/qdrant-client"
910
import { codeParser, DirectoryScanner, FileWatcher } from "./processors"
@@ -70,6 +71,11 @@ export class CodeIndexServiceFactory {
7071
throw new Error(t("embeddings:serviceFactory.mistralConfigMissing"))
7172
}
7273
return new MistralEmbedder(config.mistralOptions.apiKey, config.modelId)
74+
} else if (provider === "vertex") {
75+
if (!config.vertexOptions?.apiKey) {
76+
throw new Error(t("embeddings:serviceFactory.vertexConfigMissing"))
77+
}
78+
return new VertexEmbedder(config.vertexOptions.apiKey, config.modelId)
7379
}
7480

7581
throw new Error(

‎src/shared/embeddingModels.ts‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
* Defines profiles for different embedding models, including their dimensions.
33
*/
44

5-
export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" // Add other providers as needed
5+
export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vertex" // Add other providers as needed
66

77
export interface EmbeddingModelProfile {
88
dimension: number
@@ -53,6 +53,12 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
5353
mistral: {
5454
"codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 },
5555
},
56+
vertex: {
57+
"text-embedding-004": { dimension: 768, scoreThreshold: 0.4 },
58+
"text-multilingual-embedding-002": { dimension: 768, scoreThreshold: 0.4 },
59+
"textembedding-gecko@003": { dimension: 768, scoreThreshold: 0.4 },
60+
"textembedding-gecko-multilingual@001": { dimension: 768, scoreThreshold: 0.4 },
61+
},
5662
}
5763

5864
/**
@@ -143,6 +149,9 @@ export function getDefaultModelId(provider: EmbedderProvider): string {
143149
case "mistral":
144150
return "codestral-embed-2505"
145151

152+
case "vertex":
153+
return "text-embedding-004"
154+
146155
default:
147156
// Fallback for unknown providers
148157
console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`)

0 commit comments

Comments
 (0)