From 768e5532bd496b233c6f568c1109a6ce9a10bc18 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Fri, 10 Oct 2025 02:29:57 +0000 Subject: [PATCH] feat: add Nebius AI as codebase indexing provider - Add NebiusEmbedder class with rate limiting (600k TPM, 10k RPM) - Support Qwen/Qwen3-Embedding-8B model with 4096 dimensions - Implement proper rate limiting for tokens and requests per minute - Add comprehensive test coverage for the new embedder - Update types, config manager, and service factory to support Nebius - Cost-effective option at /bin/sh.01 per 1M tokens vs /bin/sh.13-0.18 for others Addresses #8589 --- packages/types/src/codebase-index.ts | 4 +- packages/types/src/global-settings.ts | 1 + src/core/webview/webviewMessageHandler.ts | 8 + src/services/code-index/config-manager.ts | 20 ++ .../embedders/__tests__/nebius.spec.ts | 279 ++++++++++++++++++ src/services/code-index/embedders/nebius.ts | 224 ++++++++++++++ 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/WebviewMessage.ts | 1 + src/shared/embeddingModels.ts | 15 +- 12 files changed, 574 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..c6387d3067 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", @@ -2640,6 +2647,7 @@ export const webviewMessageHandler = async ( hasGeminiApiKey, hasMistralApiKey, hasVercelAiGatewayApiKey, + hasNebiusApiKey, }, }) break 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..e765eabda9 --- /dev/null +++ b/src/services/code-index/embedders/__tests__/nebius.spec.ts @@ -0,0 +1,279 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" +import { NebiusEmbedder } from "../nebius" +import { OpenAICompatibleEmbedder } from "../openai-compatible" +import { TelemetryService } from "@roo-code/telemetry" +import { TelemetryEventName } from "@roo-code/types" + +// Mock the OpenAICompatibleEmbedder +vi.mock("../openai-compatible", () => ({ + OpenAICompatibleEmbedder: vi.fn().mockImplementation(() => ({ + createEmbeddings: vi.fn(), + validateConfiguration: vi.fn(), + })), +})) + +// Mock TelemetryService +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + instance: { + captureEvent: vi.fn(), + }, + }, +})) + +// Mock i18n +vi.mock("../../../../i18n", () => ({ + t: (key: string, params?: any) => { + if (params) { + return `${key} ${JSON.stringify(params)}` + } + return key + }, +})) + +describe("NebiusEmbedder", () => { + let embedder: NebiusEmbedder + const mockApiKey = "test-nebius-api-key" + const mockOpenAICompatibleEmbedder = { + createEmbeddings: vi.fn(), + validateConfiguration: vi.fn(), + } + + beforeEach(() => { + vi.clearAllMocks() + vi.useFakeTimers() + // Reset the mock implementation + ;(OpenAICompatibleEmbedder as any).mockImplementation(() => mockOpenAICompatibleEmbedder) + }) + + afterEach(() => { + vi.useRealTimers() + }) + + describe("constructor", () => { + it("should create an instance with default model", () => { + embedder = new NebiusEmbedder(mockApiKey) + expect(OpenAICompatibleEmbedder).toHaveBeenCalledWith( + "https://api.studio.nebius.com/v1", + mockApiKey, + "Qwen/Qwen3-Embedding-8B", + 8191, // MAX_ITEM_TOKENS + ) + }) + + it("should create an instance with custom model", () => { + const customModel = "custom-model-id" + embedder = new NebiusEmbedder(mockApiKey, customModel) + expect(OpenAICompatibleEmbedder).toHaveBeenCalledWith( + "https://api.studio.nebius.com/v1", + mockApiKey, + customModel, + 8191, + ) + }) + + it("should throw error if API key is not provided", () => { + expect(() => new NebiusEmbedder("")).toThrow("embeddings:validation.apiKeyRequired") + }) + }) + + describe("createEmbeddings", () => { + beforeEach(() => { + embedder = new NebiusEmbedder(mockApiKey) + }) + + it("should delegate to OpenAICompatibleEmbedder with rate limiting", async () => { + const texts = ["test text 1", "test text 2"] + const mockResponse = { + embeddings: [ + [0.1, 0.2], + [0.3, 0.4], + ], + usage: { promptTokens: 10, totalTokens: 10 }, + } + mockOpenAICompatibleEmbedder.createEmbeddings.mockResolvedValue(mockResponse) + + const result = await embedder.createEmbeddings(texts) + + expect(mockOpenAICompatibleEmbedder.createEmbeddings).toHaveBeenCalledWith(texts, "Qwen/Qwen3-Embedding-8B") + expect(result).toEqual(mockResponse) + }) + + it("should use custom model if provided", async () => { + const customModel = "custom-model" + const texts = ["test text"] + const mockResponse = { + embeddings: [[0.1, 0.2]], + usage: { promptTokens: 5, totalTokens: 5 }, + } + mockOpenAICompatibleEmbedder.createEmbeddings.mockResolvedValue(mockResponse) + + const result = await embedder.createEmbeddings(texts, customModel) + + expect(mockOpenAICompatibleEmbedder.createEmbeddings).toHaveBeenCalledWith(texts, customModel) + expect(result).toEqual(mockResponse) + }) + + it("should handle rate limiting when exceeding requests per minute", async () => { + const texts = ["test"] + const mockResponse = { + embeddings: [[0.1, 0.2]], + usage: { promptTokens: 5, totalTokens: 5 }, + } + mockOpenAICompatibleEmbedder.createEmbeddings.mockResolvedValue(mockResponse) + + // Make 10,000 requests to hit the RPM limit + const promises = [] + for (let i = 0; i < 10000; i++) { + promises.push(embedder.createEmbeddings(texts)) + } + await Promise.all(promises) + + // The next request should be rate limited + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + const consoleLogSpy = vi.spyOn(console, "log").mockImplementation(() => {}) + + // This request should trigger rate limiting + const rateLimitedPromise = embedder.createEmbeddings(texts) + + // Should log rate limit warning + expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining("embeddings:nebius.rateLimitExceeded")) + + // Fast-forward time to allow rate limit to reset + vi.advanceTimersByTime(60000) + + await rateLimitedPromise + + consoleWarnSpy.mockRestore() + consoleLogSpy.mockRestore() + }) + + it("should handle rate limiting when exceeding tokens per minute", async () => { + // Create very large texts that will exceed 600,000 TPM + const largeText = "a".repeat(150000) // ~37,500 tokens per text + const texts = Array(20).fill(largeText) // ~750,000 tokens total + const mockResponse = { + embeddings: Array(20).fill([0.1, 0.2]), + usage: { promptTokens: 750000, totalTokens: 750000 }, + } + mockOpenAICompatibleEmbedder.createEmbeddings.mockResolvedValue(mockResponse) + + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + const consoleLogSpy = vi.spyOn(console, "log").mockImplementation(() => {}) + + // This should trigger token rate limiting + const promise = embedder.createEmbeddings(texts) + + // Should log rate limit warning + expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining("embeddings:nebius.rateLimitExceeded")) + + // Fast-forward time to allow rate limit to reset + vi.advanceTimersByTime(60000) + + await promise + + consoleWarnSpy.mockRestore() + consoleLogSpy.mockRestore() + }) + + it("should capture telemetry on error", async () => { + const texts = ["test text"] + const error = new Error("Test error") + mockOpenAICompatibleEmbedder.createEmbeddings.mockRejectedValue(error) + + await expect(embedder.createEmbeddings(texts)).rejects.toThrow(error) + + expect(TelemetryService.instance.captureEvent).toHaveBeenCalledWith(TelemetryEventName.CODE_INDEX_ERROR, { + error: error.message, + stack: error.stack, + location: "NebiusEmbedder:createEmbeddings", + }) + }) + }) + + describe("validateConfiguration", () => { + beforeEach(() => { + embedder = new NebiusEmbedder(mockApiKey) + }) + + it("should delegate validation to OpenAICompatibleEmbedder", async () => { + const mockValidationResult = { valid: true } + mockOpenAICompatibleEmbedder.validateConfiguration.mockResolvedValue(mockValidationResult) + + const result = await embedder.validateConfiguration() + + expect(mockOpenAICompatibleEmbedder.validateConfiguration).toHaveBeenCalled() + expect(result).toEqual(mockValidationResult) + }) + + it("should handle validation errors", async () => { + const error = new Error("Validation failed") + mockOpenAICompatibleEmbedder.validateConfiguration.mockRejectedValue(error) + + await expect(embedder.validateConfiguration()).rejects.toThrow(error) + + expect(TelemetryService.instance.captureEvent).toHaveBeenCalledWith(TelemetryEventName.CODE_INDEX_ERROR, { + error: error.message, + stack: error.stack, + location: "NebiusEmbedder:validateConfiguration", + }) + }) + }) + + describe("embedderInfo", () => { + it("should return correct embedder info", () => { + embedder = new NebiusEmbedder(mockApiKey) + expect(embedder.embedderInfo).toEqual({ name: "nebius" }) + }) + }) + + describe("modelDimension", () => { + it("should return the correct model dimension", () => { + expect(NebiusEmbedder.modelDimension).toBe(4096) + }) + }) + + describe("rate limiting", () => { + beforeEach(() => { + embedder = new NebiusEmbedder(mockApiKey) + mockOpenAICompatibleEmbedder.createEmbeddings.mockResolvedValue({ + embeddings: [[0.1, 0.2]], + usage: { promptTokens: 10, totalTokens: 10 }, + }) + }) + + it("should reset rate limits after window expires", async () => { + const texts = ["test"] + + // Make requests up to the limit + for (let i = 0; i < 9999; i++) { + await embedder.createEmbeddings(texts) + } + + // Advance time to reset the window + vi.advanceTimersByTime(60001) + + // Should be able to make requests again + const result = await embedder.createEmbeddings(texts) + expect(result).toBeDefined() + expect(mockOpenAICompatibleEmbedder.createEmbeddings).toHaveBeenCalled() + }) + + it("should log debug information about rate limit usage", async () => { + const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + + const texts = ["test text"] + const mockResponse = { + embeddings: [[0.1, 0.2]], + usage: { promptTokens: 100, totalTokens: 100 }, + } + mockOpenAICompatibleEmbedder.createEmbeddings.mockResolvedValue(mockResponse) + + await embedder.createEmbeddings(texts) + + expect(consoleDebugSpy).toHaveBeenCalledWith(expect.stringContaining("Nebius AI embedding usage")) + + consoleDebugSpy.mockRestore() + }) + }) +}) diff --git a/src/services/code-index/embedders/nebius.ts b/src/services/code-index/embedders/nebius.ts new file mode 100644 index 0000000000..33031e4ff0 --- /dev/null +++ b/src/services/code-index/embedders/nebius.ts @@ -0,0 +1,224 @@ +import { OpenAICompatibleEmbedder } from "./openai-compatible" +import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder" +import { MAX_ITEM_TOKENS } from "../constants" +import { t } from "../../../i18n" +import { TelemetryEventName } from "@roo-code/types" +import { TelemetryService } from "@roo-code/telemetry" + +/** + * Rate limiting configuration for Nebius AI + * Based on the documented limits: + * - 600,000 TPM (tokens per minute) + * - 10,000 RPM (requests per minute) + */ +interface RateLimitState { + tokensUsed: number + requestsCount: number + windowStart: number +} + +/** + * Nebius AI embedder implementation that wraps the OpenAI Compatible embedder + * with configuration for Nebius AI's embedding API and rate limiting. + * + * Supported model: + * - Qwen/Qwen3-Embedding-8B (dimension: 4096) + * + * Rate limits: + * - 600,000 tokens per minute + * - 10,000 requests per minute + * + * Pricing: $0.01 per 1M tokens + */ +export class NebiusEmbedder implements IEmbedder { + private readonly openAICompatibleEmbedder: OpenAICompatibleEmbedder + private static readonly NEBIUS_BASE_URL = "https://api.studio.nebius.com/v1" + private static readonly DEFAULT_MODEL = "Qwen/Qwen3-Embedding-8B" + private static readonly MODEL_DIMENSION = 4096 + private static readonly MAX_TOKENS_PER_MINUTE = 600000 + private static readonly MAX_REQUESTS_PER_MINUTE = 10000 + private static readonly RATE_LIMIT_WINDOW_MS = 60000 // 1 minute in milliseconds + + private readonly modelId: string + private rateLimitState: RateLimitState = { + tokensUsed: 0, + requestsCount: 0, + windowStart: Date.now(), + } + + /** + * Creates a new Nebius AI embedder + * @param apiKey The Nebius AI API key for authentication + * @param modelId The model ID to use (defaults to Qwen/Qwen3-Embedding-8B) + */ + constructor(apiKey: string, modelId?: string) { + if (!apiKey) { + throw new Error(t("embeddings:validation.apiKeyRequired")) + } + + // Use provided model or default + this.modelId = modelId || NebiusEmbedder.DEFAULT_MODEL + + // Create an OpenAI Compatible embedder with Nebius's configuration + // Note: MAX_ITEM_TOKENS is the token limit per item, not the embedding dimension + this.openAICompatibleEmbedder = new OpenAICompatibleEmbedder( + NebiusEmbedder.NEBIUS_BASE_URL, + apiKey, + this.modelId, + MAX_ITEM_TOKENS, + ) + } + + /** + * Checks and updates rate limit state, implementing a sliding window approach + * @returns true if the request can proceed, false if rate limited + */ + private checkAndUpdateRateLimit(estimatedTokens: number): boolean { + const now = Date.now() + const windowElapsed = now - this.rateLimitState.windowStart + + // Reset the window if a minute has passed + if (windowElapsed >= NebiusEmbedder.RATE_LIMIT_WINDOW_MS) { + this.rateLimitState = { + tokensUsed: 0, + requestsCount: 0, + windowStart: now, + } + } + + // Check if we would exceed rate limits + if (this.rateLimitState.requestsCount >= NebiusEmbedder.MAX_REQUESTS_PER_MINUTE) { + console.warn( + t("embeddings:nebius.rateLimitExceeded", { + type: "requests", + limit: NebiusEmbedder.MAX_REQUESTS_PER_MINUTE, + window: "minute", + }), + ) + return false + } + + if (this.rateLimitState.tokensUsed + estimatedTokens > NebiusEmbedder.MAX_TOKENS_PER_MINUTE) { + console.warn( + t("embeddings:nebius.rateLimitExceeded", { + type: "tokens", + limit: NebiusEmbedder.MAX_TOKENS_PER_MINUTE, + window: "minute", + }), + ) + return false + } + + // Update the state + this.rateLimitState.tokensUsed += estimatedTokens + this.rateLimitState.requestsCount += 1 + + return true + } + + /** + * Calculates the wait time until rate limits reset + * @returns milliseconds to wait, or 0 if no wait needed + */ + private getWaitTimeMs(): number { + const now = Date.now() + const windowElapsed = now - this.rateLimitState.windowStart + const remainingTime = NebiusEmbedder.RATE_LIMIT_WINDOW_MS - windowElapsed + + return remainingTime > 0 ? remainingTime : 0 + } + + /** + * Creates embeddings for the given texts using Nebius AI's embedding API + * with built-in rate limiting for 600k TPM and 10k RPM + * @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 + + // Estimate tokens for rate limiting (rough estimate: 1 token ≈ 4 characters) + const estimatedTokens = texts.reduce((sum, text) => sum + Math.ceil(text.length / 4), 0) + + // Check rate limits + if (!this.checkAndUpdateRateLimit(estimatedTokens)) { + // Wait for the rate limit window to reset + const waitTime = this.getWaitTimeMs() + if (waitTime > 0) { + console.log( + t("embeddings:nebius.waitingForRateLimit", { + waitTimeMs: waitTime, + }), + ) + await new Promise((resolve) => setTimeout(resolve, waitTime)) + // After waiting, reset the window and try again + this.rateLimitState = { + tokensUsed: estimatedTokens, + requestsCount: 1, + windowStart: Date.now(), + } + } + } + + // Delegate to the OpenAI-compatible embedder + const result = await this.openAICompatibleEmbedder.createEmbeddings(texts, modelToUse) + + // Log usage for monitoring (optional) + if (result.usage) { + console.debug( + `Nebius AI embedding usage - Tokens: ${result.usage.totalTokens}, ` + + `Rate limit status: ${this.rateLimitState.tokensUsed}/${NebiusEmbedder.MAX_TOKENS_PER_MINUTE} TPM, ` + + `${this.rateLimitState.requestsCount}/${NebiusEmbedder.MAX_REQUESTS_PER_MINUTE} RPM`, + ) + } + + return result + } 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: "NebiusEmbedder:createEmbeddings", + }) + throw error + } + } + + /** + * Validates the Nebius AI 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 Nebius since we're using Nebius'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: "NebiusEmbedder:validateConfiguration", + }) + throw error + } + } + + /** + * Returns information about this embedder + */ + get embedderInfo(): EmbedderInfo { + return { + name: "nebius", + } + } + + /** + * Gets the model dimension for the Nebius AI model + * @returns The embedding dimension (4096 for Qwen/Qwen3-Embedding-8B) + */ + static get modelDimension(): number { + return NebiusEmbedder.MODEL_DIMENSION + } +} 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/WebviewMessage.ts b/src/shared/WebviewMessage.ts index d43a2fce04..d66b1eda00 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -302,6 +302,7 @@ export interface WebviewMessage { codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string codebaseIndexVercelAiGatewayApiKey?: string + codebaseIndexNebiusApiKey?: string } } 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.`)