From 8ba958123be00485d5e702f60b7e5573f2f6938f Mon Sep 17 00:00:00 2001 From: Roo Code Date: Fri, 18 Jul 2025 23:45:31 +0000 Subject: [PATCH] feat: add Mistral embedding provider for codebase indexing - Add Mistral as a new embedding provider option - Support for codestral-embed model with 1024 dimensions - OpenAI-compatible API integration at https://api.mistral.ai/v1 - Free tier availability for cost-effective embedding generation - Comprehensive test coverage for MistralEmbedder class - Configuration support in service factory and config manager Resolves #5932 --- src/services/code-index/config-manager.ts | 24 +++ src/services/code-index/constants/index.ts | 3 + .../embedders/__tests__/mistral.spec.ts | 191 ++++++++++++++++++ src/services/code-index/embedders/mistral.ts | 91 +++++++++ src/services/code-index/interfaces/config.ts | 2 + .../code-index/interfaces/embedder.ts | 2 +- src/services/code-index/interfaces/manager.ts | 2 +- src/services/code-index/service-factory.ts | 6 + src/shared/embeddingModels.ts | 8 +- 9 files changed, 326 insertions(+), 3 deletions(-) create mode 100644 src/services/code-index/embedders/__tests__/mistral.spec.ts create mode 100644 src/services/code-index/embedders/mistral.ts diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index 9958f456c3..843759d793 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -18,6 +18,7 @@ export class CodeIndexConfigManager { private ollamaOptions?: ApiHandlerOptions private openAiCompatibleOptions?: { baseUrl: string; apiKey: string } private geminiOptions?: { apiKey: string } + private mistralOptions?: { apiKey: string } private qdrantUrl?: string = "http://localhost:6333" private qdrantApiKey?: string private searchMinScore?: number @@ -67,6 +68,7 @@ export class CodeIndexConfigManager { const openAiCompatibleBaseUrl = codebaseIndexConfig.codebaseIndexOpenAiCompatibleBaseUrl ?? "" const openAiCompatibleApiKey = this.contextProxy?.getSecret("codebaseIndexOpenAiCompatibleApiKey") ?? "" const geminiApiKey = this.contextProxy?.getSecret("codebaseIndexGeminiApiKey") ?? "" + const mistralApiKey = this.contextProxy?.getSecret("mistralApiKey") ?? "" // Update instance variables with configuration this.codebaseIndexEnabled = codebaseIndexEnabled ?? true @@ -100,6 +102,8 @@ export class CodeIndexConfigManager { this.embedderProvider = "openai-compatible" } else if (codebaseIndexEmbedderProvider === "gemini") { this.embedderProvider = "gemini" + } else if (codebaseIndexEmbedderProvider === "mistral") { + this.embedderProvider = "mistral" } else { this.embedderProvider = "openai" } @@ -119,6 +123,7 @@ export class CodeIndexConfigManager { : undefined this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey } : undefined + this.mistralOptions = mistralApiKey ? { apiKey: mistralApiKey } : undefined } /** @@ -135,6 +140,7 @@ export class CodeIndexConfigManager { ollamaOptions?: ApiHandlerOptions openAiCompatibleOptions?: { baseUrl: string; apiKey: string } geminiOptions?: { apiKey: string } + mistralOptions?: { apiKey: string } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -153,6 +159,7 @@ export class CodeIndexConfigManager { openAiCompatibleBaseUrl: this.openAiCompatibleOptions?.baseUrl ?? "", openAiCompatibleApiKey: this.openAiCompatibleOptions?.apiKey ?? "", geminiApiKey: this.geminiOptions?.apiKey ?? "", + mistralApiKey: this.mistralOptions?.apiKey ?? "", qdrantUrl: this.qdrantUrl ?? "", qdrantApiKey: this.qdrantApiKey ?? "", } @@ -176,6 +183,7 @@ export class CodeIndexConfigManager { ollamaOptions: this.ollamaOptions, openAiCompatibleOptions: this.openAiCompatibleOptions, geminiOptions: this.geminiOptions, + mistralOptions: this.mistralOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, @@ -208,6 +216,11 @@ export class CodeIndexConfigManager { const qdrantUrl = this.qdrantUrl const isConfigured = !!(apiKey && qdrantUrl) return isConfigured + } else if (this.embedderProvider === "mistral") { + const apiKey = this.mistralOptions?.apiKey + const qdrantUrl = this.qdrantUrl + const isConfigured = !!(apiKey && qdrantUrl) + return isConfigured } return false // Should not happen if embedderProvider is always set correctly } @@ -241,6 +254,7 @@ export class CodeIndexConfigManager { const prevOpenAiCompatibleApiKey = prev?.openAiCompatibleApiKey ?? "" const prevModelDimension = prev?.modelDimension const prevGeminiApiKey = prev?.geminiApiKey ?? "" + const prevMistralApiKey = prev?.mistralApiKey ?? "" const prevQdrantUrl = prev?.qdrantUrl ?? "" const prevQdrantApiKey = prev?.qdrantApiKey ?? "" @@ -277,6 +291,7 @@ export class CodeIndexConfigManager { const currentOpenAiCompatibleApiKey = this.openAiCompatibleOptions?.apiKey ?? "" const currentModelDimension = this.modelDimension const currentGeminiApiKey = this.geminiOptions?.apiKey ?? "" + const currentMistralApiKey = this.mistralOptions?.apiKey ?? "" const currentQdrantUrl = this.qdrantUrl ?? "" const currentQdrantApiKey = this.qdrantApiKey ?? "" @@ -295,6 +310,14 @@ export class CodeIndexConfigManager { return true } + if (prevGeminiApiKey !== currentGeminiApiKey) { + return true + } + + if (prevMistralApiKey !== currentMistralApiKey) { + return true + } + // Check for model dimension changes (generic for all providers) if (prevModelDimension !== currentModelDimension) { return true @@ -351,6 +374,7 @@ export class CodeIndexConfigManager { ollamaOptions: this.ollamaOptions, openAiCompatibleOptions: this.openAiCompatibleOptions, geminiOptions: this.geminiOptions, + mistralOptions: this.mistralOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, diff --git a/src/services/code-index/constants/index.ts b/src/services/code-index/constants/index.ts index 706a73935a..b0bf206fe3 100644 --- a/src/services/code-index/constants/index.ts +++ b/src/services/code-index/constants/index.ts @@ -28,3 +28,6 @@ export const BATCH_PROCESSING_CONCURRENCY = 10 /**Gemini Embedder */ export const GEMINI_MAX_ITEM_TOKENS = 2048 + +/**Mistral Embedder */ +export const MISTRAL_MAX_ITEM_TOKENS = 8192 diff --git a/src/services/code-index/embedders/__tests__/mistral.spec.ts b/src/services/code-index/embedders/__tests__/mistral.spec.ts new file mode 100644 index 0000000000..1ba2405a27 --- /dev/null +++ b/src/services/code-index/embedders/__tests__/mistral.spec.ts @@ -0,0 +1,191 @@ +import { describe, it, expect, vi, beforeEach, MockedFunction } from "vitest" +import { MistralEmbedder } from "../mistral" +import { OpenAICompatibleEmbedder } from "../openai-compatible" + +// Mock the OpenAICompatibleEmbedder +vi.mock("../openai-compatible") +const MockedOpenAICompatibleEmbedder = vi.mocked(OpenAICompatibleEmbedder) + +// Mock telemetry +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + instance: { + captureEvent: vi.fn(), + }, + }, +})) + +// Mock i18n +vi.mock("../../../i18n", () => ({ + t: vi.fn((key: string) => key), +})) + +describe("MistralEmbedder", () => { + let embedder: MistralEmbedder + let mockCreateEmbeddings: MockedFunction + let mockValidateConfiguration: MockedFunction + + beforeEach(() => { + vi.clearAllMocks() + + // Setup mocks for OpenAICompatibleEmbedder + mockCreateEmbeddings = vi.fn() + mockValidateConfiguration = vi.fn() + + MockedOpenAICompatibleEmbedder.prototype.createEmbeddings = mockCreateEmbeddings + MockedOpenAICompatibleEmbedder.prototype.validateConfiguration = mockValidateConfiguration + }) + + describe("constructor", () => { + it("should initialize with provided API key and default model", () => { + // Arrange + const apiKey = "test-mistral-api-key" + + // Act + embedder = new MistralEmbedder(apiKey) + + // Assert + expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith( + "https://api.mistral.ai/v1", + apiKey, + "codestral-embed", + 8192, + ) + }) + + it("should initialize with provided API key and custom model", () => { + // Arrange + const apiKey = "test-mistral-api-key" + const modelId = "custom-mistral-model" + + // Act + embedder = new MistralEmbedder(apiKey, modelId) + + // Assert + expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith( + "https://api.mistral.ai/v1", + apiKey, + "custom-mistral-model", + 8192, + ) + }) + + it("should throw error when API key is not provided", () => { + // Act & Assert + expect(() => new MistralEmbedder("")).toThrow("validation.apiKeyRequired") + }) + }) + + describe("embedderInfo", () => { + it("should return correct embedder info", () => { + // Arrange + embedder = new MistralEmbedder("test-api-key") + + // Act + const info = embedder.embedderInfo + + // Assert + expect(info).toEqual({ + name: "mistral", + }) + }) + }) + + describe("createEmbeddings", () => { + describe("success cases", () => { + it("should delegate to OpenAI Compatible embedder with default model", async () => { + // Arrange + embedder = new MistralEmbedder("test-api-key") + const texts = ["test text 1", "test text 2"] + const mockResponse = { + embeddings: [ + [0.1, 0.2], + [0.3, 0.4], + ], + usage: { promptTokens: 10, totalTokens: 15 }, + } + mockCreateEmbeddings.mockResolvedValue(mockResponse) + + // Act + const result = await embedder.createEmbeddings(texts) + + // Assert + expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "codestral-embed") + expect(result).toEqual(mockResponse) + }) + + it("should delegate to OpenAI Compatible embedder with custom model", async () => { + // Arrange + embedder = new MistralEmbedder("test-api-key", "custom-model") + const texts = ["test text 1", "test text 2"] + const mockResponse = { + embeddings: [ + [0.1, 0.2], + [0.3, 0.4], + ], + usage: { promptTokens: 10, totalTokens: 15 }, + } + mockCreateEmbeddings.mockResolvedValue(mockResponse) + + // Act + const result = await embedder.createEmbeddings(texts, "codestral-embed") + + // Assert + expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "codestral-embed") + expect(result).toEqual(mockResponse) + }) + + it("should handle errors from OpenAI Compatible embedder", async () => { + // Arrange + embedder = new MistralEmbedder("test-api-key") + const texts = ["test text"] + mockCreateEmbeddings.mockRejectedValue(new Error("Embedding failed")) + + // Act & Assert + await expect(embedder.createEmbeddings(texts)).rejects.toThrow("Embedding failed") + }) + }) + }) + + describe("validateConfiguration", () => { + it("should delegate to OpenAI Compatible embedder and return success", async () => { + // Arrange + embedder = new MistralEmbedder("test-api-key") + mockValidateConfiguration.mockResolvedValue({ valid: true }) + + // Act + const result = await embedder.validateConfiguration() + + // Assert + expect(mockValidateConfiguration).toHaveBeenCalled() + expect(result).toEqual({ valid: true }) + }) + + it("should delegate to OpenAI Compatible embedder and return error", async () => { + // Arrange + embedder = new MistralEmbedder("test-api-key") + mockValidateConfiguration.mockResolvedValue({ + valid: false, + error: "embeddings:validation.authenticationFailed", + }) + + // Act + const result = await embedder.validateConfiguration() + + // Assert + expect(result).toEqual({ + valid: false, + error: "embeddings:validation.authenticationFailed", + }) + }) + + it("should handle validation errors", async () => { + // Arrange + embedder = new MistralEmbedder("test-api-key") + mockValidateConfiguration.mockRejectedValue(new Error("Validation failed")) + + // Act & Assert + await expect(embedder.validateConfiguration()).rejects.toThrow("Validation failed") + }) + }) +}) diff --git a/src/services/code-index/embedders/mistral.ts b/src/services/code-index/embedders/mistral.ts new file mode 100644 index 0000000000..3189e16c46 --- /dev/null +++ b/src/services/code-index/embedders/mistral.ts @@ -0,0 +1,91 @@ +import { OpenAICompatibleEmbedder } from "./openai-compatible" +import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder" +import { MISTRAL_MAX_ITEM_TOKENS } from "../constants" +import { t } from "../../../i18n" +import { TelemetryEventName } from "@roo-code/types" +import { TelemetryService } from "@roo-code/telemetry" + +/** + * Mistral embedder implementation that wraps the OpenAI Compatible embedder + * with configuration for Mistral's embedding API. + * + * Supported models: + * - codestral-embed (dimension: 1024) + */ +export class MistralEmbedder implements IEmbedder { + private readonly openAICompatibleEmbedder: OpenAICompatibleEmbedder + private static readonly MISTRAL_BASE_URL = "https://api.mistral.ai/v1" + private static readonly DEFAULT_MODEL = "codestral-embed" + private readonly modelId: string + + /** + * Creates a new Mistral embedder + * @param apiKey The Mistral API key for authentication + * @param modelId The model ID to use (defaults to codestral-embed) + */ + constructor(apiKey: string, modelId?: string) { + if (!apiKey) { + throw new Error(t("embeddings:validation.apiKeyRequired")) + } + + // Use provided model or default + this.modelId = modelId || MistralEmbedder.DEFAULT_MODEL + + // Create an OpenAI Compatible embedder with Mistral's configuration + this.openAICompatibleEmbedder = new OpenAICompatibleEmbedder( + MistralEmbedder.MISTRAL_BASE_URL, + apiKey, + this.modelId, + MISTRAL_MAX_ITEM_TOKENS, + ) + } + + /** + * Creates embeddings for the given texts using Mistral's embedding API + * @param texts Array of text strings to embed + * @param model Optional model identifier (uses constructor model if not provided) + * @returns Promise resolving to embedding response + */ + 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) + } 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: "MistralEmbedder:createEmbeddings", + }) + throw error + } + } + + /** + * Validates the Mistral embedder configuration by delegating to the underlying OpenAI-compatible embedder + * @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 Mistral since we're using Mistral's base URL + return await this.openAICompatibleEmbedder.validateConfiguration() + } 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: "MistralEmbedder:validateConfiguration", + }) + throw error + } + } + + /** + * Returns information about this embedder + */ + get embedderInfo(): EmbedderInfo { + return { + name: "mistral", + } + } +} diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts index 190a23e2a3..9098a60091 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -13,6 +13,7 @@ export interface CodeIndexConfig { ollamaOptions?: ApiHandlerOptions openAiCompatibleOptions?: { baseUrl: string; apiKey: string } geminiOptions?: { apiKey: string } + mistralOptions?: { apiKey: string } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -33,6 +34,7 @@ export type PreviousConfigSnapshot = { openAiCompatibleBaseUrl?: string openAiCompatibleApiKey?: string geminiApiKey?: string + mistralApiKey?: string qdrantUrl?: string qdrantApiKey?: string } diff --git a/src/services/code-index/interfaces/embedder.ts b/src/services/code-index/interfaces/embedder.ts index 0a74446d5e..c5653ea2b7 100644 --- a/src/services/code-index/interfaces/embedder.ts +++ b/src/services/code-index/interfaces/embedder.ts @@ -28,7 +28,7 @@ export interface EmbeddingResponse { } } -export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" +export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" export interface EmbedderInfo { name: AvailableEmbedders diff --git a/src/services/code-index/interfaces/manager.ts b/src/services/code-index/interfaces/manager.ts index 70e3fd9765..fd3b2bfdda 100644 --- a/src/services/code-index/interfaces/manager.ts +++ b/src/services/code-index/interfaces/manager.ts @@ -70,7 +70,7 @@ export interface ICodeIndexManager { } export type IndexingState = "Standby" | "Indexing" | "Indexed" | "Error" -export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" +export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" export interface IndexProgressUpdate { systemStatus: IndexingState diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts index b7951db7ac..68b0f5c0bc 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -3,6 +3,7 @@ import { OpenAiEmbedder } from "./embedders/openai" import { CodeIndexOllamaEmbedder } from "./embedders/ollama" import { OpenAICompatibleEmbedder } from "./embedders/openai-compatible" import { GeminiEmbedder } from "./embedders/gemini" +import { MistralEmbedder } from "./embedders/mistral" import { EmbedderProvider, getDefaultModelId, getModelDimension } from "../../shared/embeddingModels" import { QdrantVectorStore } from "./vector-store/qdrant-client" import { codeParser, DirectoryScanner, FileWatcher } from "./processors" @@ -64,6 +65,11 @@ export class CodeIndexServiceFactory { throw new Error(t("embeddings:serviceFactory.geminiConfigMissing")) } return new GeminiEmbedder(config.geminiOptions.apiKey, config.modelId) + } else if (provider === "mistral") { + if (!config.mistralOptions?.apiKey) { + throw new Error(t("embeddings:serviceFactory.mistralConfigMissing")) + } + return new MistralEmbedder(config.mistralOptions.apiKey, config.modelId) } throw new Error( diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index f387480c65..e9d3871a84 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -2,7 +2,7 @@ * Defines profiles for different embedding models, including their dimensions. */ -export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" // Add other providers as needed +export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" // Add other providers as needed export interface EmbeddingModelProfile { dimension: number @@ -50,6 +50,9 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { "text-embedding-004": { dimension: 768 }, "gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 }, }, + mistral: { + "codestral-embed": { dimension: 1024, scoreThreshold: 0.4 }, + }, } /** @@ -137,6 +140,9 @@ export function getDefaultModelId(provider: EmbedderProvider): string { case "gemini": return "gemini-embedding-001" + case "mistral": + return "codestral-embed" + default: // Fallback for unknown providers console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`)