From 8bf8326a105abaca7996aedfde57672d32708060 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Mon, 28 Jul 2025 18:56:55 +0000 Subject: [PATCH] feat: rewrite VertexEmbedder to use @google/genai library with multiple auth options - Replace OpenAI-compatible approach with native @google/genai SDK - Add support for multiple authentication methods: - API key (uses regular Gemini API endpoint) - JSON credentials (service account) - Key file path - Application default credentials - Add projectId and location fields for Vertex AI configuration - Update UI to show all authentication options for Vertex - Update tests to reflect new implementation - Update all related type definitions and interfaces --- packages/types/src/codebase-index.ts | 5 + packages/types/src/global-settings.ts | 2 + src/core/webview/ClineProvider.ts | 4 + src/core/webview/webviewMessageHandler.ts | 28 ++ .../embedders/__tests__/vertex.spec.ts | 377 +++++++++++++----- src/services/code-index/embedders/vertex.ts | 139 +++++-- src/services/code-index/interfaces/config.ts | 12 +- src/services/code-index/service-factory.ts | 23 +- src/shared/WebviewMessage.ts | 7 +- .../src/components/chat/CodeIndexPopover.tsx | 158 +++++++- webview-ui/src/i18n/locales/en/settings.json | 13 + 11 files changed, 629 insertions(+), 139 deletions(-) diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index 8c7986e51f..5102af8444 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -36,6 +36,9 @@ export const codebaseIndexConfigSchema = z.object({ // OpenAI Compatible specific fields codebaseIndexOpenAiCompatibleBaseUrl: z.string().optional(), codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(), + // Vertex AI specific fields + codebaseIndexVertexProjectId: z.string().optional(), + codebaseIndexVertexLocation: z.string().optional(), }) export type CodebaseIndexConfig = z.infer @@ -68,6 +71,8 @@ export const codebaseIndexProviderSchema = z.object({ codebaseIndexGeminiApiKey: z.string().optional(), codebaseIndexMistralApiKey: z.string().optional(), codebaseIndexVertexApiKey: z.string().optional(), + codebaseIndexVertexJsonCredentials: z.string().optional(), + codebaseIndexVertexKeyFile: z.string().optional(), }) export type CodebaseIndexProvider = z.infer diff --git a/packages/types/src/global-settings.ts b/packages/types/src/global-settings.ts index d84c99b84a..96dfbb88d8 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -185,6 +185,8 @@ export const SECRET_STATE_KEYS = [ "codebaseIndexGeminiApiKey", "codebaseIndexMistralApiKey", "codebaseIndexVertexApiKey", + "codebaseIndexVertexJsonCredentials", + "codebaseIndexVertexKeyFile", "huggingFaceApiKey", ] as const satisfies readonly (keyof ProviderSettings)[] export type SecretState = Pick diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 6bcb85e337..8e562d1ffe 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -1554,6 +1554,8 @@ export class ClineProvider codebaseIndexOpenAiCompatibleBaseUrl: codebaseIndexConfig?.codebaseIndexOpenAiCompatibleBaseUrl, codebaseIndexSearchMaxResults: codebaseIndexConfig?.codebaseIndexSearchMaxResults, codebaseIndexSearchMinScore: codebaseIndexConfig?.codebaseIndexSearchMinScore, + codebaseIndexVertexProjectId: codebaseIndexConfig?.codebaseIndexVertexProjectId, + codebaseIndexVertexLocation: codebaseIndexConfig?.codebaseIndexVertexLocation, }, mdmCompliant: this.checkMdmCompliance(), profileThresholds: profileThresholds ?? {}, @@ -1726,6 +1728,8 @@ export class ClineProvider stateValues.codebaseIndexConfig?.codebaseIndexOpenAiCompatibleBaseUrl, codebaseIndexSearchMaxResults: stateValues.codebaseIndexConfig?.codebaseIndexSearchMaxResults, codebaseIndexSearchMinScore: stateValues.codebaseIndexConfig?.codebaseIndexSearchMinScore, + codebaseIndexVertexProjectId: stateValues.codebaseIndexConfig?.codebaseIndexVertexProjectId, + codebaseIndexVertexLocation: stateValues.codebaseIndexConfig?.codebaseIndexVertexLocation, }, profileThresholds: stateValues.profileThresholds ?? {}, // Add diagnostic message settings diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index da73c56920..947c1c5218 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -1998,6 +1998,8 @@ export const webviewMessageHandler = async ( codebaseIndexOpenAiCompatibleBaseUrl: settings.codebaseIndexOpenAiCompatibleBaseUrl, codebaseIndexSearchMaxResults: settings.codebaseIndexSearchMaxResults, codebaseIndexSearchMinScore: settings.codebaseIndexSearchMinScore, + codebaseIndexVertexProjectId: settings.codebaseIndexVertexProjectId, + codebaseIndexVertexLocation: settings.codebaseIndexVertexLocation, } // Save global state first @@ -2028,6 +2030,24 @@ export const webviewMessageHandler = async ( settings.codebaseIndexMistralApiKey, ) } + if (settings.codebaseIndexVertexApiKey !== undefined) { + await provider.contextProxy.storeSecret( + "codebaseIndexVertexApiKey", + settings.codebaseIndexVertexApiKey, + ) + } + if (settings.codebaseIndexVertexJsonCredentials !== undefined) { + await provider.contextProxy.storeSecret( + "codebaseIndexVertexJsonCredentials", + settings.codebaseIndexVertexJsonCredentials, + ) + } + if (settings.codebaseIndexVertexKeyFile !== undefined) { + await provider.contextProxy.storeSecret( + "codebaseIndexVertexKeyFile", + settings.codebaseIndexVertexKeyFile, + ) + } // Send success response first - settings are saved regardless of validation await provider.postMessageToWebview({ @@ -2149,6 +2169,11 @@ export const webviewMessageHandler = async ( )) const hasGeminiApiKey = !!(await provider.context.secrets.get("codebaseIndexGeminiApiKey")) const hasMistralApiKey = !!(await provider.context.secrets.get("codebaseIndexMistralApiKey")) + const hasVertexApiKey = !!(await provider.context.secrets.get("codebaseIndexVertexApiKey")) + const hasVertexJsonCredentials = !!(await provider.context.secrets.get( + "codebaseIndexVertexJsonCredentials", + )) + const hasVertexKeyFile = !!(await provider.context.secrets.get("codebaseIndexVertexKeyFile")) provider.postMessageToWebview({ type: "codeIndexSecretStatus", @@ -2158,6 +2183,9 @@ export const webviewMessageHandler = async ( hasOpenAiCompatibleApiKey, hasGeminiApiKey, hasMistralApiKey, + hasVertexApiKey, + hasVertexJsonCredentials, + hasVertexKeyFile, }, }) break diff --git a/src/services/code-index/embedders/__tests__/vertex.spec.ts b/src/services/code-index/embedders/__tests__/vertex.spec.ts index 743bd3695a..aa91ebeac4 100644 --- a/src/services/code-index/embedders/__tests__/vertex.spec.ts +++ b/src/services/code-index/embedders/__tests__/vertex.spec.ts @@ -1,10 +1,9 @@ import { vitest, describe, it, expect, beforeEach } from "vitest" -import type { MockedClass } from "vitest" import { VertexEmbedder } from "../vertex" -import { OpenAICompatibleEmbedder } from "../openai-compatible" +import { GoogleGenAI } from "@google/genai" -// Mock the OpenAICompatibleEmbedder -vitest.mock("../openai-compatible") +// Mock the @google/genai library +vitest.mock("@google/genai") // Mock TelemetryService vitest.mock("@roo-code/telemetry", () => ({ @@ -15,61 +14,161 @@ vitest.mock("@roo-code/telemetry", () => ({ }, })) -const MockedOpenAICompatibleEmbedder = OpenAICompatibleEmbedder as MockedClass +// Mock i18n +vitest.mock("../../../../i18n", () => ({ + t: (key: string, params?: Record) => { + const translations: Record = { + "validation.apiKeyRequired": "API key is required", + "embeddings:validation.authenticationFailed": "Authentication failed", + "embeddings:validation.connectionFailed": "Connection failed", + "embeddings:validation.modelNotAvailable": "Model not available", + "embeddings:validation.unexpectedError": "Unexpected error", + "embeddings:validation.vertexAuthRequired": "At least one authentication method is required for Vertex AI", + "embeddings:validation.noEmbeddingsReturned": "No embeddings returned", + "embeddings:validation.configurationError": "Configuration error", + } + return translations[key] || key + }, +})) + +// Mock safeJsonParse +vitest.mock("../../../shared/safeJsonParse", () => ({ + safeJsonParse: (json: string, defaultValue: any) => { + try { + return JSON.parse(json) + } catch { + return defaultValue + } + }, +})) describe("VertexEmbedder", () => { let embedder: VertexEmbedder + let mockClient: any + let mockModel: any + let mockEmbedContent: any beforeEach(() => { vitest.clearAllMocks() + + // Setup mock for embedContent + mockEmbedContent = vitest.fn() + mockModel = { + embedContent: mockEmbedContent, + } + mockClient = { + models: { + embedContent: mockEmbedContent, + }, + } + + // Mock GoogleGenAI constructor + ;(GoogleGenAI as any).mockImplementation(() => mockClient) }) describe("constructor", () => { - it("should create an instance with default model when no model specified", () => { - // Arrange - const apiKey = "test-vertex-api-key" - + it("should create an instance with API key authentication", () => { // Act - embedder = new VertexEmbedder(apiKey) + embedder = new VertexEmbedder({ + apiKey: "test-api-key", + projectId: "test-project", + location: "us-central1", + }) // Assert - expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith( - "https://generativelanguage.googleapis.com/v1beta/openai/", - apiKey, - "text-embedding-004", - 2048, - ) + expect(GoogleGenAI).toHaveBeenCalledWith({ apiKey: "test-api-key" }) + expect(embedder.embedderInfo.name).toBe("vertex") }) - it("should create an instance with specified model", () => { - // Arrange - const apiKey = "test-vertex-api-key" - const modelId = "text-multilingual-embedding-002" - + it("should create an instance with JSON credentials authentication", () => { // Act - embedder = new VertexEmbedder(apiKey, modelId) + embedder = new VertexEmbedder({ + jsonCredentials: '{"type": "service_account"}', + projectId: "test-project", + location: "us-central1", + }) // Assert - expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith( - "https://generativelanguage.googleapis.com/v1beta/openai/", - apiKey, - "text-multilingual-embedding-002", - 2048, - ) + expect(GoogleGenAI).toHaveBeenCalledWith({ + vertexai: true, + project: "test-project", + location: "us-central1", + googleAuthOptions: { + credentials: { type: "service_account" }, + }, + }) + expect(embedder.embedderInfo.name).toBe("vertex") }) - it("should throw error when API key is not provided", () => { - // Act & Assert - expect(() => new VertexEmbedder("")).toThrow("validation.apiKeyRequired") - expect(() => new VertexEmbedder(null as any)).toThrow("validation.apiKeyRequired") - expect(() => new VertexEmbedder(undefined as any)).toThrow("validation.apiKeyRequired") + it("should create an instance with key file authentication", () => { + // Act + embedder = new VertexEmbedder({ + keyFile: "/path/to/keyfile.json", + projectId: "test-project", + location: "us-central1", + }) + + // Assert + expect(GoogleGenAI).toHaveBeenCalledWith({ + vertexai: true, + project: "test-project", + location: "us-central1", + googleAuthOptions: { keyFile: "/path/to/keyfile.json" }, + }) + expect(embedder.embedderInfo.name).toBe("vertex") + }) + + it("should create an instance with application default credentials", () => { + // Act + embedder = new VertexEmbedder({ + apiKey: "", // Empty string to trigger ADC path + projectId: "test-project", + location: "us-central1", + }) + + // Assert + expect(GoogleGenAI).toHaveBeenCalledWith({ + vertexai: true, + project: "test-project", + location: "us-central1", + }) + expect(embedder.embedderInfo.name).toBe("vertex") + }) + + it("should use default model when not specified", () => { + // Act + embedder = new VertexEmbedder({ + apiKey: "test-api-key", + projectId: "test-project", + location: "us-central1", + }) + + // Assert + expect(embedder["modelId"]).toBe("text-embedding-004") + }) + + it("should use specified model", () => { + // Act + embedder = new VertexEmbedder({ + apiKey: "test-api-key", + modelId: "text-multilingual-embedding-002", + projectId: "test-project", + location: "us-central1", + }) + + // Assert + expect(embedder["modelId"]).toBe("text-multilingual-embedding-002") }) }) describe("embedderInfo", () => { it("should return correct embedder info", () => { // Arrange - embedder = new VertexEmbedder("test-api-key") + embedder = new VertexEmbedder({ + apiKey: "test-api-key", + projectId: "test-project", + location: "us-central1", + }) // Act const info = embedder.embedderInfo @@ -79,115 +178,185 @@ describe("VertexEmbedder", () => { name: "vertex", }) }) + }) - describe("createEmbeddings", () => { - let mockCreateEmbeddings: any - - beforeEach(() => { - mockCreateEmbeddings = vitest.fn() - MockedOpenAICompatibleEmbedder.prototype.createEmbeddings = mockCreateEmbeddings + describe("createEmbeddings", () => { + beforeEach(() => { + embedder = new VertexEmbedder({ + apiKey: "test-api-key", + projectId: "test-project", + location: "us-central1", }) + }) - it("should use instance model when no model parameter provided", async () => { - // Arrange - embedder = new VertexEmbedder("test-api-key") - const texts = ["test text 1", "test text 2"] - const mockResponse = { - embeddings: [ - [0.1, 0.2], - [0.3, 0.4], - ], - } - mockCreateEmbeddings.mockResolvedValue(mockResponse) + it("should create embeddings for single text", async () => { + // Arrange + const texts = ["test text"] + const mockResponse = { + embeddings: [{ values: [0.1, 0.2, 0.3] }], + } + mockEmbedContent.mockResolvedValue(mockResponse) - // Act - const result = await embedder.createEmbeddings(texts) + // Act + const result = await embedder.createEmbeddings(texts) - // Assert - expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "text-embedding-004") - expect(result).toEqual(mockResponse) + // Assert + expect(mockEmbedContent).toHaveBeenCalledWith({ + model: "text-embedding-004", + contents: [{ parts: [{ text: "test text" }] }], }) - - it("should use provided model parameter when specified", async () => { - // Arrange - embedder = new VertexEmbedder("test-api-key", "textembedding-gecko@003") - const texts = ["test text 1", "test text 2"] - const mockResponse = { - embeddings: [ - [0.1, 0.2], - [0.3, 0.4], - ], - } - mockCreateEmbeddings.mockResolvedValue(mockResponse) - - // Act - const result = await embedder.createEmbeddings(texts, "text-multilingual-embedding-002") - - // Assert - expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "text-multilingual-embedding-002") - expect(result).toEqual(mockResponse) + expect(result).toEqual({ + embeddings: [[0.1, 0.2, 0.3]], }) + }) - it("should handle errors from OpenAICompatibleEmbedder", async () => { - // Arrange - embedder = new VertexEmbedder("test-api-key") - const texts = ["test text"] - const error = new Error("Embedding failed") - mockCreateEmbeddings.mockRejectedValue(error) + it("should create embeddings for multiple texts in batches", async () => { + // Arrange + const texts = ["text1", "text2", "text3"] + const mockResponse = { + embeddings: [{ values: [0.1, 0.2] }, { values: [0.3, 0.4] }, { values: [0.5, 0.6] }], + } + mockEmbedContent.mockResolvedValue(mockResponse) - // Act & Assert - await expect(embedder.createEmbeddings(texts)).rejects.toThrow("Embedding failed") + // Act + const result = await embedder.createEmbeddings(texts) + + // Assert + expect(mockEmbedContent).toHaveBeenCalledTimes(1) + expect(mockEmbedContent).toHaveBeenCalledWith({ + model: "text-embedding-004", + contents: [ + { parts: [{ text: "text1" }] }, + { parts: [{ text: "text2" }] }, + { parts: [{ text: "text3" }] }, + ], }) + expect(result).toEqual({ + embeddings: [ + [0.1, 0.2], + [0.3, 0.4], + [0.5, 0.6], + ], + }) + }) + + it("should use custom model when provided", async () => { + // Arrange + const texts = ["test text"] + const mockResponse = { + embeddings: [{ values: [0.1, 0.2] }], + } + mockEmbedContent.mockResolvedValue(mockResponse) + + // Act + await embedder.createEmbeddings(texts, "text-multilingual-embedding-002") + + // Assert + expect(mockEmbedContent).toHaveBeenCalledWith({ + model: "text-multilingual-embedding-002", + contents: [{ parts: [{ text: "test text" }] }], + }) + }) + + it("should handle empty text array", async () => { + // Act + const result = await embedder.createEmbeddings([]) + + // Assert + expect(mockEmbedContent).not.toHaveBeenCalled() + expect(result).toEqual({ embeddings: [] }) + }) + + it("should handle API errors", async () => { + // Arrange + const texts = ["test text"] + const error = new Error("API Error") + mockEmbedContent.mockRejectedValue(error) + + // Act & Assert + await expect(embedder.createEmbeddings(texts)).rejects.toThrow("API Error") }) }) describe("validateConfiguration", () => { - let mockValidateConfiguration: any - beforeEach(() => { - mockValidateConfiguration = vitest.fn() - MockedOpenAICompatibleEmbedder.prototype.validateConfiguration = mockValidateConfiguration + embedder = new VertexEmbedder({ + apiKey: "test-api-key", + projectId: "test-project", + location: "us-central1", + }) }) - it("should delegate validation to OpenAICompatibleEmbedder", async () => { + it("should validate configuration successfully", async () => { // Arrange - embedder = new VertexEmbedder("test-api-key") - mockValidateConfiguration.mockResolvedValue({ valid: true }) + const mockResponse = { + embeddings: [{ values: [0.1, 0.2] }], + } + mockEmbedContent.mockResolvedValue(mockResponse) // Act const result = await embedder.validateConfiguration() // Assert - expect(mockValidateConfiguration).toHaveBeenCalled() + expect(mockEmbedContent).toHaveBeenCalledWith({ + model: "text-embedding-004", + contents: [{ parts: [{ text: "test" }] }], + }) expect(result).toEqual({ valid: true }) }) - it("should pass through validation errors from OpenAICompatibleEmbedder", async () => { + it("should handle unexpected errors", async () => { // Arrange - embedder = new VertexEmbedder("test-api-key") - mockValidateConfiguration.mockResolvedValue({ - valid: false, - error: "embeddings:validation.authenticationFailed", - }) + const error = new Error("Something went wrong") + mockEmbedContent.mockRejectedValue(error) // Act const result = await embedder.validateConfiguration() // Assert - expect(mockValidateConfiguration).toHaveBeenCalled() expect(result).toEqual({ valid: false, - error: "embeddings:validation.authenticationFailed", + error: "Something went wrong", + }) + }) + }) + + describe("createBatches", () => { + beforeEach(() => { + embedder = new VertexEmbedder({ + apiKey: "test-api-key", + projectId: "test-project", + location: "us-central1", }) }) - it("should handle validation exceptions", async () => { + it("should create batches respecting token limits", () => { // Arrange - embedder = new VertexEmbedder("test-api-key") - mockValidateConfiguration.mockRejectedValue(new Error("Validation failed")) + const texts = [ + "short text", + "another short text", + "a".repeat(5000), // Long text + "more text", + ] - // Act & Assert - await expect(embedder.validateConfiguration()).rejects.toThrow("Validation failed") + // Act + const batches = embedder["createBatches"](texts) + + // Assert + expect(batches.length).toBe(1) + expect(batches[0].length).toBe(4) // All texts in one batch (under 100 limit) + }) + + it("should handle all oversized texts", () => { + // Arrange + const texts = ["a".repeat(10000), "b".repeat(10000)] + + // Act + const batches = embedder["createBatches"](texts) + + // Assert + expect(batches.length).toBe(1) + expect(batches[0].length).toBe(2) // Both texts in one batch }) }) }) diff --git a/src/services/code-index/embedders/vertex.ts b/src/services/code-index/embedders/vertex.ts index adc58d5da6..11e2e22cde 100644 --- a/src/services/code-index/embedders/vertex.ts +++ b/src/services/code-index/embedders/vertex.ts @@ -1,13 +1,15 @@ -import { OpenAICompatibleEmbedder } from "./openai-compatible" +import { GoogleGenAI } from "@google/genai" +import type { JWTInput } from "google-auth-library" import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder" import { VERTEX_MAX_ITEM_TOKENS } from "../constants" import { t } from "../../../i18n" import { TelemetryEventName } from "@roo-code/types" import { TelemetryService } from "@roo-code/telemetry" +import { safeJsonParse } from "../../../shared/safeJsonParse" /** - * Vertex AI embedder implementation that wraps the OpenAI Compatible embedder - * with configuration for Google's Vertex AI embedding API. + * Vertex AI embedder implementation using the @google/genai library + * with support for multiple authentication methods. * * Supported models: * - text-embedding-004 (dimension: 768) @@ -16,31 +18,65 @@ import { TelemetryService } from "@roo-code/telemetry" * - textembedding-gecko-multilingual@001 (dimension: 768) */ export class VertexEmbedder implements IEmbedder { - private readonly openAICompatibleEmbedder: OpenAICompatibleEmbedder - private static readonly VERTEX_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/" + private readonly client: GoogleGenAI private static readonly DEFAULT_MODEL = "text-embedding-004" private readonly modelId: string + private readonly maxItemTokens: number /** * Creates a new Vertex AI embedder - * @param apiKey The Google AI API key for authentication - * @param modelId The model ID to use (defaults to text-embedding-004) + * @param options Configuration options including authentication methods */ - constructor(apiKey: string, modelId?: string) { - if (!apiKey) { - throw new Error(t("embeddings:validation.apiKeyRequired")) + constructor(options: { + apiKey?: string + jsonCredentials?: string + keyFile?: string + projectId: string + location: string + modelId?: string + }) { + const { apiKey, jsonCredentials, keyFile, projectId, location, modelId } = options + + // Validate required fields + if (!projectId) { + throw new Error("Project ID is required for Vertex AI") + } + if (!location) { + throw new Error("Location is required for Vertex AI") } // Use provided model or default this.modelId = modelId || VertexEmbedder.DEFAULT_MODEL + this.maxItemTokens = VERTEX_MAX_ITEM_TOKENS - // Create an OpenAI Compatible embedder with Vertex AI's configuration - this.openAICompatibleEmbedder = new OpenAICompatibleEmbedder( - VertexEmbedder.VERTEX_BASE_URL, - apiKey, - this.modelId, - VERTEX_MAX_ITEM_TOKENS, - ) + // Create the GoogleGenAI client with appropriate auth + if (jsonCredentials) { + this.client = new GoogleGenAI({ + vertexai: true, + project: projectId, + location, + googleAuthOptions: { + credentials: safeJsonParse(jsonCredentials, undefined), + }, + }) + } else if (keyFile) { + this.client = new GoogleGenAI({ + vertexai: true, + project: projectId, + location, + googleAuthOptions: { keyFile }, + }) + } else if (apiKey && apiKey.trim() !== "") { + // For API key auth, we use the regular Gemini API endpoint + this.client = new GoogleGenAI({ apiKey }) + } else { + // Default to application default credentials + this.client = new GoogleGenAI({ + vertexai: true, + project: projectId, + location, + }) + } } /** @@ -51,9 +87,33 @@ export class VertexEmbedder implements IEmbedder { */ async createEmbeddings(texts: string[], model?: string): Promise { try { - // Use the provided model or fall back to the instance's model const modelToUse = model || this.modelId - return await this.openAICompatibleEmbedder.createEmbeddings(texts, modelToUse) + + // Batch texts if they exceed token limits + const batches = this.createBatches(texts) + const allEmbeddings: number[][] = [] + + for (const batch of batches) { + const result = await this.client.models.embedContent({ + model: modelToUse, + contents: batch.map((text) => ({ parts: [{ text }] })), + }) + + if (!result.embeddings || result.embeddings.length === 0) { + throw new Error(t("embeddings:validation.noEmbeddingsReturned")) + } + + // Filter out any embeddings without values + const validEmbeddings = result.embeddings + .filter((e) => e.values !== undefined) + .map((e) => e.values as number[]) + + allEmbeddings.push(...validEmbeddings) + } + + return { + embeddings: allEmbeddings, + } } catch (error) { TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { error: error instanceof Error ? error.message : String(error), @@ -65,21 +125,52 @@ export class VertexEmbedder implements IEmbedder { } /** - * Validates the Vertex AI embedder configuration by delegating to the underlying OpenAI-compatible embedder + * Creates batches of texts that respect token limits + */ + private createBatches(texts: string[]): string[][] { + // Simple batching - in production, you'd want to estimate tokens + const batchSize = 100 // Vertex AI typically supports up to 100 texts per batch + const batches: string[][] = [] + + for (let i = 0; i < texts.length; i += batchSize) { + batches.push(texts.slice(i, i + batchSize)) + } + + return batches + } + + /** + * Validates the Vertex AI embedder configuration * @returns Promise resolving to validation result with success status and optional error message */ async validateConfiguration(): Promise<{ valid: boolean; error?: string }> { try { - // Delegate validation to the OpenAI-compatible embedder - // The error messages will be specific to Vertex AI since we're using Vertex AI's base URL - return await this.openAICompatibleEmbedder.validateConfiguration() + // Test with a simple embedding request + const testText = "test" + const result = await this.client.models.embedContent({ + model: this.modelId, + contents: [{ parts: [{ text: testText }] }], + }) + + if (!result.embeddings || result.embeddings.length === 0) { + return { + valid: false, + error: t("embeddings:validation.noEmbeddingsReturned"), + } + } + + return { valid: true } } catch (error) { TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { error: error instanceof Error ? error.message : String(error), stack: error instanceof Error ? error.stack : undefined, location: "VertexEmbedder:validateConfiguration", }) - throw error + + return { + valid: false, + error: error instanceof Error ? error.message : t("embeddings:validation.configurationError"), + } } } diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts index 32a0085d22..69dff5699a 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -14,7 +14,13 @@ export interface CodeIndexConfig { openAiCompatibleOptions?: { baseUrl: string; apiKey: string } geminiOptions?: { apiKey: string } mistralOptions?: { apiKey: string } - vertexOptions?: { apiKey: string } + vertexOptions?: { + apiKey?: string + jsonCredentials?: string + keyFile?: string + projectId?: string + location?: string + } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -37,6 +43,10 @@ export type PreviousConfigSnapshot = { geminiApiKey?: string mistralApiKey?: string vertexApiKey?: string + vertexJsonCredentials?: string + vertexKeyFile?: string + vertexProjectId?: string + vertexLocation?: string qdrantUrl?: string qdrantApiKey?: string } diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts index d0a60c44fd..b9bdb74873 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -72,10 +72,29 @@ export class CodeIndexServiceFactory { } return new MistralEmbedder(config.mistralOptions.apiKey, config.modelId) } else if (provider === "vertex") { - if (!config.vertexOptions?.apiKey) { + const vertexOptions = config.vertexOptions + if (!vertexOptions) { throw new Error(t("embeddings:serviceFactory.vertexConfigMissing")) } - return new VertexEmbedder(config.vertexOptions.apiKey, config.modelId) + + // Validate that at least one auth method is provided + if (!vertexOptions.apiKey && !vertexOptions.jsonCredentials && !vertexOptions.keyFile) { + throw new Error(t("embeddings:serviceFactory.vertexAuthRequired")) + } + + // Validate required fields for Vertex AI + if (!vertexOptions.projectId || !vertexOptions.location) { + throw new Error(t("embeddings:serviceFactory.vertexProjectLocationRequired")) + } + + return new VertexEmbedder({ + apiKey: vertexOptions.apiKey, + jsonCredentials: vertexOptions.jsonCredentials, + keyFile: vertexOptions.keyFile, + projectId: vertexOptions.projectId, + location: vertexOptions.location, + modelId: config.modelId, + }) } throw new Error( diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index 1304e4c7d5..540e802037 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -249,13 +249,15 @@ export interface WebviewMessage { // Global state settings codebaseIndexEnabled: boolean codebaseIndexQdrantUrl: string - codebaseIndexEmbedderProvider: "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" + codebaseIndexEmbedderProvider: "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vertex" codebaseIndexEmbedderBaseUrl?: string codebaseIndexEmbedderModelId: string codebaseIndexEmbedderModelDimension?: number // Generic dimension for all providers codebaseIndexOpenAiCompatibleBaseUrl?: string codebaseIndexSearchMaxResults?: number codebaseIndexSearchMinScore?: number + codebaseIndexVertexProjectId?: string + codebaseIndexVertexLocation?: string // Secret settings codeIndexOpenAiKey?: string @@ -263,6 +265,9 @@ export interface WebviewMessage { codebaseIndexOpenAiCompatibleApiKey?: string codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string + codebaseIndexVertexApiKey?: string + codebaseIndexVertexJsonCredentials?: string + codebaseIndexVertexKeyFile?: string } } diff --git a/webview-ui/src/components/chat/CodeIndexPopover.tsx b/webview-ui/src/components/chat/CodeIndexPopover.tsx index 8ee7fd8823..8ad621043e 100644 --- a/webview-ui/src/components/chat/CodeIndexPopover.tsx +++ b/webview-ui/src/components/chat/CodeIndexPopover.tsx @@ -70,6 +70,10 @@ interface LocalCodeIndexSettings { codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string codebaseIndexVertexApiKey?: string + codebaseIndexVertexJsonCredentials?: string + codebaseIndexVertexKeyFile?: string + codebaseIndexVertexProjectId?: string + codebaseIndexVertexLocation?: string } // Validation schema for codebase index settings @@ -137,12 +141,32 @@ const createValidationSchema = (provider: EmbedderProvider, t: any) => { }) case "vertex": - return baseSchema.extend({ - codebaseIndexVertexApiKey: z.string().min(1, t("settings:codeIndex.validation.vertexApiKeyRequired")), - codebaseIndexEmbedderModelId: z - .string() - .min(1, t("settings:codeIndex.validation.modelSelectionRequired")), - }) + return baseSchema + .extend({ + // At least one auth method is required + codebaseIndexVertexApiKey: z.string().optional(), + codebaseIndexVertexJsonCredentials: z.string().optional(), + codebaseIndexVertexKeyFile: z.string().optional(), + codebaseIndexVertexProjectId: z + .string() + .min(1, t("settings:codeIndex.validation.vertexProjectIdRequired")), + codebaseIndexVertexLocation: z + .string() + .min(1, t("settings:codeIndex.validation.vertexLocationRequired")), + codebaseIndexEmbedderModelId: z + .string() + .min(1, t("settings:codeIndex.validation.modelSelectionRequired")), + }) + .refine( + (data) => + data.codebaseIndexVertexApiKey || + data.codebaseIndexVertexJsonCredentials || + data.codebaseIndexVertexKeyFile, + { + message: t("settings:codeIndex.validation.vertexAuthRequired"), + path: ["codebaseIndexVertexApiKey"], + }, + ) default: return baseSchema @@ -189,6 +213,10 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexGeminiApiKey: "", codebaseIndexMistralApiKey: "", codebaseIndexVertexApiKey: "", + codebaseIndexVertexJsonCredentials: "", + codebaseIndexVertexKeyFile: "", + codebaseIndexVertexProjectId: "", + codebaseIndexVertexLocation: "us-central1", }) // Initial settings state - stores the settings when popover opens @@ -224,6 +252,10 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexGeminiApiKey: "", codebaseIndexMistralApiKey: "", codebaseIndexVertexApiKey: "", + codebaseIndexVertexJsonCredentials: "", + codebaseIndexVertexKeyFile: "", + codebaseIndexVertexProjectId: codebaseIndexConfig.codebaseIndexVertexProjectId || "", + codebaseIndexVertexLocation: codebaseIndexConfig.codebaseIndexVertexLocation || "us-central1", } setInitialSettings(settings) setCurrentSettings(settings) @@ -321,6 +353,17 @@ export const CodeIndexPopover: React.FC = ({ if (!prev.codebaseIndexVertexApiKey || prev.codebaseIndexVertexApiKey === SECRET_PLACEHOLDER) { updated.codebaseIndexVertexApiKey = secretStatus.hasVertexApiKey ? SECRET_PLACEHOLDER : "" } + if ( + !prev.codebaseIndexVertexJsonCredentials || + prev.codebaseIndexVertexJsonCredentials === SECRET_PLACEHOLDER + ) { + updated.codebaseIndexVertexJsonCredentials = secretStatus.hasVertexJsonCredentials + ? SECRET_PLACEHOLDER + : "" + } + if (!prev.codebaseIndexVertexKeyFile || prev.codebaseIndexVertexKeyFile === SECRET_PLACEHOLDER) { + updated.codebaseIndexVertexKeyFile = secretStatus.hasVertexKeyFile ? SECRET_PLACEHOLDER : "" + } return updated } @@ -394,7 +437,9 @@ export const CodeIndexPopover: React.FC = ({ key === "codebaseIndexOpenAiCompatibleApiKey" || key === "codebaseIndexGeminiApiKey" || key === "codebaseIndexMistralApiKey" || - key === "codebaseIndexVertexApiKey" + key === "codebaseIndexVertexApiKey" || + key === "codebaseIndexVertexJsonCredentials" || + key === "codebaseIndexVertexKeyFile" ) { dataToValidate[key] = "placeholder-valid" } @@ -1036,6 +1081,15 @@ export const CodeIndexPopover: React.FC = ({ {currentSettings.codebaseIndexEmbedderProvider === "vertex" && ( <> +
+ +

+ {t("settings:codeIndex.vertexAuthMethodDescription")} +

+
+
+
+ + + updateSetting( + "codebaseIndexVertexJsonCredentials", + e.target.value, + ) + } + placeholder={t( + "settings:codeIndex.vertexJsonCredentialsPlaceholder", + )} + className={cn("w-full", { + "border-red-500": formErrors.codebaseIndexVertexJsonCredentials, + })} + /> + {formErrors.codebaseIndexVertexJsonCredentials && ( +

+ {formErrors.codebaseIndexVertexJsonCredentials} +

+ )} +
+ +
+ + + updateSetting("codebaseIndexVertexKeyFile", e.target.value) + } + placeholder={t("settings:codeIndex.vertexKeyFilePlaceholder")} + className={cn("w-full", { + "border-red-500": formErrors.codebaseIndexVertexKeyFile, + })} + /> + {formErrors.codebaseIndexVertexKeyFile && ( +

+ {formErrors.codebaseIndexVertexKeyFile} +

+ )} +
+ +
+ + + updateSetting("codebaseIndexVertexProjectId", e.target.value) + } + placeholder={t("settings:codeIndex.vertexProjectIdPlaceholder")} + className={cn("w-full", { + "border-red-500": formErrors.codebaseIndexVertexProjectId, + })} + /> + {formErrors.codebaseIndexVertexProjectId && ( +

+ {formErrors.codebaseIndexVertexProjectId} +

+ )} +
+ +
+ + + updateSetting("codebaseIndexVertexLocation", e.target.value) + } + placeholder={t("settings:codeIndex.vertexLocationPlaceholder")} + className={cn("w-full", { + "border-red-500": formErrors.codebaseIndexVertexLocation, + })} + /> + {formErrors.codebaseIndexVertexLocation && ( +

+ {formErrors.codebaseIndexVertexLocation} +

+ )} +
+