From b574f5e7b582a66f05a2ecf9d54fcc8adba3ae7d Mon Sep 17 00:00:00 2001 From: Roo Code Date: Fri, 10 Oct 2025 02:26:10 +0000 Subject: [PATCH] feat: add Nebius AI as codebase indexing provider for cost-effective embeddings --- packages/types/src/codebase-index.ts | 4 +- packages/types/src/global-settings.ts | 1 + src/core/webview/webviewMessageHandler.ts | 7 + src/i18n/locales/en/embeddings.json | 4 + src/services/code-index/config-manager.ts | 20 ++ .../embedders/__tests__/nebius.spec.ts | 255 ++++++++++++++++++ src/services/code-index/embedders/nebius.ts | 250 +++++++++++++++++ src/services/code-index/interfaces/config.ts | 2 + .../code-index/interfaces/embedder.ts | 9 +- src/services/code-index/interfaces/manager.ts | 9 +- src/services/code-index/service-factory.ts | 6 + src/shared/embeddingModels.ts | 15 +- 12 files changed, 578 insertions(+), 4 deletions(-) create mode 100644 src/services/code-index/embedders/__tests__/nebius.spec.ts create mode 100644 src/services/code-index/embedders/nebius.ts diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index be7778f538..14831c7ded 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -22,7 +22,7 @@ export const codebaseIndexConfigSchema = z.object({ codebaseIndexEnabled: z.boolean().optional(), codebaseIndexQdrantUrl: z.string().optional(), codebaseIndexEmbedderProvider: z - .enum(["openai", "ollama", "openai-compatible", "gemini", "mistral", "vercel-ai-gateway"]) + .enum(["openai", "ollama", "openai-compatible", "gemini", "mistral", "vercel-ai-gateway", "nebius"]) .optional(), codebaseIndexEmbedderBaseUrl: z.string().optional(), codebaseIndexEmbedderModelId: z.string().optional(), @@ -51,6 +51,7 @@ export const codebaseIndexModelsSchema = z.object({ gemini: z.record(z.string(), z.object({ dimension: z.number() })).optional(), mistral: z.record(z.string(), z.object({ dimension: z.number() })).optional(), "vercel-ai-gateway": z.record(z.string(), z.object({ dimension: z.number() })).optional(), + nebius: z.record(z.string(), z.object({ dimension: z.number() })).optional(), }) export type CodebaseIndexModels = z.infer @@ -68,6 +69,7 @@ export const codebaseIndexProviderSchema = z.object({ codebaseIndexGeminiApiKey: z.string().optional(), codebaseIndexMistralApiKey: z.string().optional(), codebaseIndexVercelAiGatewayApiKey: z.string().optional(), + codebaseIndexNebiusApiKey: 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 a56a00fc35..e81a12cf23 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -199,6 +199,7 @@ export const SECRET_STATE_KEYS = [ "codebaseIndexGeminiApiKey", "codebaseIndexMistralApiKey", "codebaseIndexVercelAiGatewayApiKey", + "codebaseIndexNebiusApiKey", "huggingFaceApiKey", "sambaNovaApiKey", "zaiApiKey", diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index af5f9925c3..0a8b6a168a 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -2493,6 +2493,12 @@ export const webviewMessageHandler = async ( settings.codebaseIndexVercelAiGatewayApiKey, ) } + if (settings.codebaseIndexNebiusApiKey !== undefined) { + await provider.contextProxy.storeSecret( + "codebaseIndexNebiusApiKey", + settings.codebaseIndexNebiusApiKey, + ) + } // Send success response first - settings are saved regardless of validation await provider.postMessageToWebview({ @@ -2630,6 +2636,7 @@ export const webviewMessageHandler = async ( const hasVercelAiGatewayApiKey = !!(await provider.context.secrets.get( "codebaseIndexVercelAiGatewayApiKey", )) + const hasNebiusApiKey = !!(await provider.context.secrets.get("codebaseIndexNebiusApiKey")) provider.postMessageToWebview({ type: "codeIndexSecretStatus", diff --git a/src/i18n/locales/en/embeddings.json b/src/i18n/locales/en/embeddings.json index fc902cadc1..cf70b28e43 100644 --- a/src/i18n/locales/en/embeddings.json +++ b/src/i18n/locales/en/embeddings.json @@ -17,6 +17,9 @@ "modelNotEmbeddingCapable": "Ollama model is not embedding capable: {{modelId}}", "hostNotFound": "Ollama host not found: {{baseUrl}}" }, + "nebius": { + "invalidResponseFormat": "Invalid response format from Nebius AI API" + }, "scanner": { "unknownErrorProcessingFile": "Unknown error processing file {{filePath}}", "unknownErrorDeletingPoints": "Unknown error deleting points for {{filePath}}", @@ -48,6 +51,7 @@ "geminiConfigMissing": "Gemini configuration missing for embedder creation", "mistralConfigMissing": "Mistral configuration missing for embedder creation", "vercelAiGatewayConfigMissing": "Vercel AI Gateway configuration missing for embedder creation", + "nebiusConfigMissing": "Nebius AI configuration missing for embedder creation", "invalidEmbedderType": "Invalid embedder type configured: {{embedderProvider}}", "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.", "vectorDimensionNotDetermined": "Could not determine vector dimension for model '{{modelId}}' with provider '{{provider}}'. Check model profiles or configuration.", diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index 2c0e8bb5c9..e05a341032 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -20,6 +20,7 @@ export class CodeIndexConfigManager { private geminiOptions?: { apiKey: string } private mistralOptions?: { apiKey: string } private vercelAiGatewayOptions?: { apiKey: string } + private nebiusOptions?: { apiKey: string } private qdrantUrl?: string = "http://localhost:6333" private qdrantApiKey?: string private searchMinScore?: number @@ -71,6 +72,7 @@ export class CodeIndexConfigManager { const geminiApiKey = this.contextProxy?.getSecret("codebaseIndexGeminiApiKey") ?? "" const mistralApiKey = this.contextProxy?.getSecret("codebaseIndexMistralApiKey") ?? "" const vercelAiGatewayApiKey = this.contextProxy?.getSecret("codebaseIndexVercelAiGatewayApiKey") ?? "" + const nebiusApiKey = this.contextProxy?.getSecret("codebaseIndexNebiusApiKey") ?? "" // Update instance variables with configuration this.codebaseIndexEnabled = codebaseIndexEnabled ?? true @@ -108,6 +110,8 @@ export class CodeIndexConfigManager { this.embedderProvider = "mistral" } else if (codebaseIndexEmbedderProvider === "vercel-ai-gateway") { this.embedderProvider = "vercel-ai-gateway" + } else if (codebaseIndexEmbedderProvider === "nebius") { + this.embedderProvider = "nebius" } else { this.embedderProvider = "openai" } @@ -129,6 +133,7 @@ export class CodeIndexConfigManager { this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey } : undefined this.mistralOptions = mistralApiKey ? { apiKey: mistralApiKey } : undefined this.vercelAiGatewayOptions = vercelAiGatewayApiKey ? { apiKey: vercelAiGatewayApiKey } : undefined + this.nebiusOptions = nebiusApiKey ? { apiKey: nebiusApiKey } : undefined } /** @@ -147,6 +152,7 @@ export class CodeIndexConfigManager { geminiOptions?: { apiKey: string } mistralOptions?: { apiKey: string } vercelAiGatewayOptions?: { apiKey: string } + nebiusOptions?: { apiKey: string } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -167,6 +173,7 @@ export class CodeIndexConfigManager { geminiApiKey: this.geminiOptions?.apiKey ?? "", mistralApiKey: this.mistralOptions?.apiKey ?? "", vercelAiGatewayApiKey: this.vercelAiGatewayOptions?.apiKey ?? "", + nebiusApiKey: this.nebiusOptions?.apiKey ?? "", qdrantUrl: this.qdrantUrl ?? "", qdrantApiKey: this.qdrantApiKey ?? "", } @@ -192,6 +199,7 @@ export class CodeIndexConfigManager { geminiOptions: this.geminiOptions, mistralOptions: this.mistralOptions, vercelAiGatewayOptions: this.vercelAiGatewayOptions, + nebiusOptions: this.nebiusOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, @@ -234,6 +242,11 @@ export class CodeIndexConfigManager { const qdrantUrl = this.qdrantUrl const isConfigured = !!(apiKey && qdrantUrl) return isConfigured + } else if (this.embedderProvider === "nebius") { + const apiKey = this.nebiusOptions?.apiKey + const qdrantUrl = this.qdrantUrl + const isConfigured = !!(apiKey && qdrantUrl) + return isConfigured } return false // Should not happen if embedderProvider is always set correctly } @@ -269,6 +282,7 @@ export class CodeIndexConfigManager { const prevGeminiApiKey = prev?.geminiApiKey ?? "" const prevMistralApiKey = prev?.mistralApiKey ?? "" const prevVercelAiGatewayApiKey = prev?.vercelAiGatewayApiKey ?? "" + const prevNebiusApiKey = prev?.nebiusApiKey ?? "" const prevQdrantUrl = prev?.qdrantUrl ?? "" const prevQdrantApiKey = prev?.qdrantApiKey ?? "" @@ -307,6 +321,7 @@ export class CodeIndexConfigManager { const currentGeminiApiKey = this.geminiOptions?.apiKey ?? "" const currentMistralApiKey = this.mistralOptions?.apiKey ?? "" const currentVercelAiGatewayApiKey = this.vercelAiGatewayOptions?.apiKey ?? "" + const currentNebiusApiKey = this.nebiusOptions?.apiKey ?? "" const currentQdrantUrl = this.qdrantUrl ?? "" const currentQdrantApiKey = this.qdrantApiKey ?? "" @@ -337,6 +352,10 @@ export class CodeIndexConfigManager { return true } + if (prevNebiusApiKey !== currentNebiusApiKey) { + return true + } + // Check for model dimension changes (generic for all providers) if (prevModelDimension !== currentModelDimension) { return true @@ -395,6 +414,7 @@ export class CodeIndexConfigManager { geminiOptions: this.geminiOptions, mistralOptions: this.mistralOptions, vercelAiGatewayOptions: this.vercelAiGatewayOptions, + nebiusOptions: this.nebiusOptions, qdrantUrl: this.qdrantUrl, qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, diff --git a/src/services/code-index/embedders/__tests__/nebius.spec.ts b/src/services/code-index/embedders/__tests__/nebius.spec.ts new file mode 100644 index 0000000000..1f5781dcfb --- /dev/null +++ b/src/services/code-index/embedders/__tests__/nebius.spec.ts @@ -0,0 +1,255 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" +import { NebiusEmbedder } from "../nebius" +import { OpenAI } from "openai" +import { t } from "../../../../i18n" + +// Mock dependencies +vi.mock("openai") +vi.mock("../../../../i18n", () => ({ + t: vi.fn((key: string, params?: any) => { + if (params) { + return `${key} ${JSON.stringify(params)}` + } + return key + }), +})) + +// Mock the validation helpers +vi.mock("../../shared/validation-helpers", () => ({ + withValidationErrorHandling: vi.fn(async (fn, provider) => { + try { + return await fn() + } catch (error) { + return { + valid: false, + error: error instanceof Error ? error.message : String(error), + } + } + }), + formatEmbeddingError: vi.fn((error, maxRetries) => { + if (error instanceof Error) { + return error + } + return new Error(`Failed after ${maxRetries} attempts`) + }), + HttpError: class HttpError extends Error { + status?: number + constructor(message: string, status?: number) { + super(message) + this.status = status + } + }, +})) + +describe("NebiusEmbedder", () => { + let embedder: NebiusEmbedder + let mockCreate: ReturnType + + beforeEach(() => { + mockCreate = vi.fn() + ;(OpenAI as any).mockImplementation(() => ({ + embeddings: { + create: mockCreate, + }, + })) + }) + + afterEach(() => { + vi.clearAllMocks() + }) + + describe("constructor", () => { + it("should create embedder with API key", () => { + const apiKey = "test-api-key" + + // Act + embedder = new NebiusEmbedder(apiKey) + + // Assert + expect(OpenAI).toHaveBeenCalledWith({ + baseURL: "https://api.studio.nebius.com/v1/", + apiKey: apiKey, + }) + }) + + it("should create embedder with custom model ID", () => { + const apiKey = "test-api-key" + const modelId = "custom-model" + + // Act + embedder = new NebiusEmbedder(apiKey, modelId) + + // Assert + expect(OpenAI).toHaveBeenCalledWith({ + baseURL: "https://api.studio.nebius.com/v1/", + apiKey: apiKey, + }) + }) + + it("should throw error if API key is not provided", () => { + // Act & Assert + expect(() => new NebiusEmbedder("")).toThrow("validation.apiKeyRequired") + expect(() => new NebiusEmbedder(null as any)).toThrow("validation.apiKeyRequired") + expect(() => new NebiusEmbedder(undefined as any)).toThrow("validation.apiKeyRequired") + }) + }) + + describe("createEmbeddings", () => { + beforeEach(() => { + // Arrange + embedder = new NebiusEmbedder("test-api-key") + }) + + it("should create embeddings successfully", async () => { + // Arrange + const texts = ["test text 1", "test text 2"] + const mockResponse = { + data: [ + { embedding: btoa(new Float32Array(4096).buffer as any) }, + { embedding: btoa(new Float32Array(4096).buffer as any) }, + ], + usage: { + prompt_tokens: 10, + total_tokens: 20, + }, + } + mockCreate.mockResolvedValue(mockResponse) + + // Act + const result = await embedder.createEmbeddings(texts) + + // Assert + expect(mockCreate).toHaveBeenCalledWith({ + input: texts, + model: "Qwen/Qwen3-Embedding-8B", + encoding_format: "base64", + }) + expect(result.embeddings).toHaveLength(2) + expect(result.usage).toEqual({ + promptTokens: 10, + totalTokens: 20, + }) + }) + + it("should use custom model if provided", async () => { + // Arrange + embedder = new NebiusEmbedder("test-api-key", "custom-embed-model") + const texts = ["test text 1", "test text 2"] + const mockResponse = { + data: [ + { embedding: btoa(new Float32Array(4096).buffer as any) }, + { embedding: btoa(new Float32Array(4096).buffer as any) }, + ], + usage: { + prompt_tokens: 10, + total_tokens: 20, + }, + } + mockCreate.mockResolvedValue(mockResponse) + + // Act + const result = await embedder.createEmbeddings(texts, "override-model") + + // Assert + expect(mockCreate).toHaveBeenCalledWith({ + input: texts, + model: "override-model", + encoding_format: "base64", + }) + expect(result.embeddings).toHaveLength(2) + }) + + it("should handle rate limit errors with retry", async () => { + // Arrange + embedder = new NebiusEmbedder("test-api-key") + const texts = ["test text"] + const rateLimitError = new Error("Rate limit exceeded") as any + rateLimitError.status = 429 + + const mockResponse = { + data: [{ embedding: btoa(new Float32Array(4096).buffer as any) }], + usage: { + prompt_tokens: 5, + total_tokens: 10, + }, + } + + mockCreate.mockRejectedValueOnce(rateLimitError).mockResolvedValueOnce(mockResponse) + + // Act + const result = await embedder.createEmbeddings(texts) + + // Assert + expect(mockCreate).toHaveBeenCalledTimes(2) + expect(result.embeddings).toHaveLength(1) + }) + }) + + describe("validateConfiguration", () => { + it("should validate configuration successfully", async () => { + // Arrange + embedder = new NebiusEmbedder("test-api-key") + mockCreate.mockResolvedValue({ + data: [{ embedding: btoa(new Float32Array(4096).buffer as any) }], + }) + + // Act + const result = await embedder.validateConfiguration() + + // Assert + expect(result).toEqual({ valid: true }) + expect(mockCreate).toHaveBeenCalledWith({ + input: ["test"], + model: "Qwen/Qwen3-Embedding-8B", + encoding_format: "base64", + }) + }) + + it("should return invalid if response has no data", async () => { + // Arrange + embedder = new NebiusEmbedder("test-api-key") + mockCreate.mockResolvedValue({ + data: [], + }) + + // Act + const result = await embedder.validateConfiguration() + + // Assert + expect(result).toEqual({ + valid: false, + error: "embeddings:nebius.invalidResponseFormat", + }) + }) + + it("should handle validation errors", async () => { + // Arrange + embedder = new NebiusEmbedder("test-api-key") + mockCreate.mockRejectedValue(new Error("Validation failed")) + + // Act + const result = await embedder.validateConfiguration() + + // Assert + expect(result).toEqual({ + valid: false, + error: "Validation failed", + }) + }) + }) + + describe("embedderInfo", () => { + it("should return correct embedder info", () => { + // Arrange + embedder = new NebiusEmbedder("test-api-key") + + // Act + const info = embedder.embedderInfo + + // Assert + expect(info).toEqual({ + name: "nebius", + }) + }) + }) +}) diff --git a/src/services/code-index/embedders/nebius.ts b/src/services/code-index/embedders/nebius.ts new file mode 100644 index 0000000000..256b2f33d7 --- /dev/null +++ b/src/services/code-index/embedders/nebius.ts @@ -0,0 +1,250 @@ +import { OpenAI } from "openai" +import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder" +import { + MAX_BATCH_TOKENS, + MAX_ITEM_TOKENS, + MAX_BATCH_RETRIES as MAX_RETRIES, + INITIAL_RETRY_DELAY_MS as INITIAL_DELAY_MS, +} from "../constants" +import { getDefaultModelId, getModelQueryPrefix } from "../../../shared/embeddingModels" +import { t } from "../../../i18n" +import { withValidationErrorHandling, HttpError, formatEmbeddingError } from "../shared/validation-helpers" +import { TelemetryEventName } from "@roo-code/types" +import { TelemetryService } from "@roo-code/telemetry" +import { handleOpenAIError } from "../../../api/providers/utils/openai-error-handler" + +/** + * Nebius AI implementation of the embedder interface with batching and rate limiting. + * Uses the Qwen/Qwen3-Embedding-8B model for cost-effective embeddings. + */ +export class NebiusEmbedder implements IEmbedder { + private embeddingsClient: OpenAI + private readonly defaultModelId: string + private readonly baseUrl: string = "https://api.studio.nebius.com/v1/" + private readonly apiKey: string + + /** + * Creates a new Nebius AI embedder + * @param apiKey The API key for authentication + * @param modelId Optional model identifier (defaults to "Qwen/Qwen3-Embedding-8B") + */ + constructor(apiKey: string, modelId?: string) { + if (!apiKey) { + throw new Error(t("embeddings:validation.apiKeyRequired")) + } + + this.apiKey = apiKey + + // Wrap OpenAI client creation to handle invalid API key characters + try { + this.embeddingsClient = new OpenAI({ + baseURL: this.baseUrl, + apiKey: apiKey, + }) + } catch (error) { + // Use the error handler to transform ByteString conversion errors + throw handleOpenAIError(error, "Nebius AI") + } + + this.defaultModelId = modelId || getDefaultModelId("nebius") + } + + /** + * Creates embeddings for the given texts with batching and rate limiting + * @param texts Array of text strings to embed + * @param model Optional model identifier + * @returns Promise resolving to embedding response + */ + async createEmbeddings(texts: string[], model?: string): Promise { + const modelToUse = model || this.defaultModelId + + // Apply model-specific query prefix if required + const queryPrefix = getModelQueryPrefix("nebius", modelToUse) + const processedTexts = queryPrefix + ? texts.map((text, index) => { + // Prevent double-prefixing + if (text.startsWith(queryPrefix)) { + return text + } + const prefixedText = `${queryPrefix}${text}` + const estimatedTokens = Math.ceil(prefixedText.length / 4) + if (estimatedTokens > MAX_ITEM_TOKENS) { + console.warn( + t("embeddings:textWithPrefixExceedsTokenLimit", { + index, + estimatedTokens, + maxTokens: MAX_ITEM_TOKENS, + }), + ) + // Return original text if adding prefix would exceed limit + return text + } + return prefixedText + }) + : texts + + const allEmbeddings: number[][] = [] + const usage = { promptTokens: 0, totalTokens: 0 } + const remainingTexts = [...processedTexts] + + while (remainingTexts.length > 0) { + const currentBatch: string[] = [] + let currentBatchTokens = 0 + const processedIndices: number[] = [] + + for (let i = 0; i < remainingTexts.length; i++) { + const text = remainingTexts[i] + const itemTokens = Math.ceil(text.length / 4) + + if (itemTokens > MAX_ITEM_TOKENS) { + console.warn( + t("embeddings:textExceedsTokenLimit", { + index: i, + itemTokens, + maxTokens: MAX_ITEM_TOKENS, + }), + ) + processedIndices.push(i) + continue + } + + if (currentBatchTokens + itemTokens <= MAX_BATCH_TOKENS) { + currentBatch.push(text) + currentBatchTokens += itemTokens + processedIndices.push(i) + } else { + break + } + } + + // Remove processed items from remainingTexts (in reverse order to maintain correct indices) + for (let i = processedIndices.length - 1; i >= 0; i--) { + remainingTexts.splice(processedIndices[i], 1) + } + + if (currentBatch.length > 0) { + const batchResult = await this._embedBatchWithRetries(currentBatch, modelToUse) + allEmbeddings.push(...batchResult.embeddings) + usage.promptTokens += batchResult.usage.promptTokens + usage.totalTokens += batchResult.usage.totalTokens + } + } + + return { embeddings: allEmbeddings, usage } + } + + /** + * Helper method to handle batch embedding with retries and exponential backoff + * @param batchTexts Array of texts to embed in this batch + * @param model Model identifier to use + * @returns Promise resolving to embeddings and usage statistics + */ + private async _embedBatchWithRetries( + batchTexts: string[], + model: string, + ): Promise<{ embeddings: number[][]; usage: { promptTokens: number; totalTokens: number } }> { + for (let attempts = 0; attempts < MAX_RETRIES; attempts++) { + try { + const response = await this.embeddingsClient.embeddings.create({ + input: batchTexts, + model: model, + // Request base64 encoding to handle large dimension arrays properly + encoding_format: "base64", + }) + + // Convert base64 embeddings to float32 arrays if needed + const embeddings = response.data.map((item: any) => { + if (typeof item.embedding === "string") { + const buffer = Buffer.from(item.embedding, "base64") + // Create Float32Array view over the buffer + const float32Array = new Float32Array(buffer.buffer, buffer.byteOffset, buffer.byteLength / 4) + return Array.from(float32Array) + } + return item.embedding as number[] + }) + + return { + embeddings: embeddings, + usage: { + promptTokens: response.usage?.prompt_tokens || 0, + totalTokens: response.usage?.total_tokens || 0, + }, + } + } catch (error: any) { + const hasMoreAttempts = attempts < MAX_RETRIES - 1 + + // Check if it's a rate limit error + const httpError = error as HttpError + if (httpError?.status === 429 && hasMoreAttempts) { + const delayMs = INITIAL_DELAY_MS * Math.pow(2, attempts) + console.warn( + t("embeddings:rateLimitRetry", { + delayMs, + attempt: attempts + 1, + maxRetries: MAX_RETRIES, + }), + ) + await new Promise((resolve) => setTimeout(resolve, delayMs)) + continue + } + + // Capture telemetry before reformatting the error + TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { + error: error instanceof Error ? error.message : String(error), + stack: error instanceof Error ? error.stack : undefined, + location: "NebiusEmbedder:_embedBatchWithRetries", + attempt: attempts + 1, + }) + + // Log the error for debugging + console.error(`Nebius AI embedder error (attempt ${attempts + 1}/${MAX_RETRIES}):`, error) + + // Format and throw the error + throw formatEmbeddingError(error, MAX_RETRIES) + } + } + + throw new Error(t("embeddings:failedMaxAttempts", { attempts: MAX_RETRIES })) + } + + /** + * Validates the Nebius AI embedder configuration by attempting a minimal embedding request + * @returns Promise resolving to validation result with success status and optional error message + */ + async validateConfiguration(): Promise<{ valid: boolean; error?: string }> { + return withValidationErrorHandling(async () => { + try { + // Test with a minimal embedding request + const response = await this.embeddingsClient.embeddings.create({ + input: ["test"], + model: this.defaultModelId, + encoding_format: "base64", + }) + + // Check if we got a valid response + if (!response.data || response.data.length === 0) { + return { + valid: false, + error: t("embeddings:nebius.invalidResponseFormat"), + } + } + + return { valid: true } + } catch (error) { + // Capture telemetry for validation errors + TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { + error: error instanceof Error ? error.message : String(error), + stack: error instanceof Error ? error.stack : undefined, + location: "NebiusEmbedder:validateConfiguration", + }) + throw error + } + }, "nebius") + } + + get embedderInfo(): EmbedderInfo { + return { + name: "nebius", + } + } +} diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts index f168e26869..87f305d67e 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -15,6 +15,7 @@ export interface CodeIndexConfig { geminiOptions?: { apiKey: string } mistralOptions?: { apiKey: string } vercelAiGatewayOptions?: { apiKey: string } + nebiusOptions?: { apiKey: string } qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -37,6 +38,7 @@ export type PreviousConfigSnapshot = { geminiApiKey?: string mistralApiKey?: string vercelAiGatewayApiKey?: string + nebiusApiKey?: string qdrantUrl?: string qdrantApiKey?: string } diff --git a/src/services/code-index/interfaces/embedder.ts b/src/services/code-index/interfaces/embedder.ts index 1fcda3aca3..2f2ac0e939 100644 --- a/src/services/code-index/interfaces/embedder.ts +++ b/src/services/code-index/interfaces/embedder.ts @@ -28,7 +28,14 @@ export interface EmbeddingResponse { } } -export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vercel-ai-gateway" +export type AvailableEmbedders = + | "openai" + | "ollama" + | "openai-compatible" + | "gemini" + | "mistral" + | "vercel-ai-gateway" + | "nebius" export interface EmbedderInfo { name: AvailableEmbedders diff --git a/src/services/code-index/interfaces/manager.ts b/src/services/code-index/interfaces/manager.ts index 527900f6d1..219b24aa03 100644 --- a/src/services/code-index/interfaces/manager.ts +++ b/src/services/code-index/interfaces/manager.ts @@ -70,7 +70,14 @@ export interface ICodeIndexManager { } export type IndexingState = "Standby" | "Indexing" | "Indexed" | "Error" -export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vercel-ai-gateway" +export type EmbedderProvider = + | "openai" + | "ollama" + | "openai-compatible" + | "gemini" + | "mistral" + | "vercel-ai-gateway" + | "nebius" export interface IndexProgressUpdate { systemStatus: IndexingState diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts index 6d69e1f0b6..b366777a03 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -5,6 +5,7 @@ import { OpenAICompatibleEmbedder } from "./embedders/openai-compatible" import { GeminiEmbedder } from "./embedders/gemini" import { MistralEmbedder } from "./embedders/mistral" import { VercelAiGatewayEmbedder } from "./embedders/vercel-ai-gateway" +import { NebiusEmbedder } from "./embedders/nebius" import { EmbedderProvider, getDefaultModelId, getModelDimension } from "../../shared/embeddingModels" import { QdrantVectorStore } from "./vector-store/qdrant-client" import { codeParser, DirectoryScanner, FileWatcher } from "./processors" @@ -79,6 +80,11 @@ export class CodeIndexServiceFactory { throw new Error(t("embeddings:serviceFactory.vercelAiGatewayConfigMissing")) } return new VercelAiGatewayEmbedder(config.vercelAiGatewayOptions.apiKey, config.modelId) + } else if (provider === "nebius") { + if (!config.nebiusOptions?.apiKey) { + throw new Error(t("embeddings:serviceFactory.nebiusConfigMissing")) + } + return new NebiusEmbedder(config.nebiusOptions.apiKey, config.modelId) } throw new Error( diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index 80c51a6b45..8d2aac9092 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -2,7 +2,14 @@ * Defines profiles for different embedding models, including their dimensions. */ -export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" | "vercel-ai-gateway" // Add other providers as needed +export type EmbedderProvider = + | "openai" + | "ollama" + | "openai-compatible" + | "gemini" + | "mistral" + | "vercel-ai-gateway" + | "nebius" // Add other providers as needed export interface EmbeddingModelProfile { dimension: number @@ -70,6 +77,9 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { "mistral/codestral-embed": { dimension: 1536, scoreThreshold: 0.4 }, "mistral/mistral-embed": { dimension: 1024, scoreThreshold: 0.4 }, }, + nebius: { + "Qwen/Qwen3-Embedding-8B": { dimension: 4096, scoreThreshold: 0.4 }, + }, } /** @@ -163,6 +173,9 @@ export function getDefaultModelId(provider: EmbedderProvider): string { case "vercel-ai-gateway": return "openai/text-embedding-3-large" + case "nebius": + return "Qwen/Qwen3-Embedding-8B" + default: // Fallback for unknown providers console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`)