mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-05 08:10:14 +00:00
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
This commit is contained in:
parent
bdc91b2aef
commit
768e5532bd
12 changed files with 574 additions and 4 deletions
|
|
@ -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<typeof codebaseIndexModelsSchema>
|
||||
|
|
@ -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<typeof codebaseIndexProviderSchema>
|
||||
|
|
|
|||
|
|
@ -199,6 +199,7 @@ export const SECRET_STATE_KEYS = [
|
|||
"codebaseIndexGeminiApiKey",
|
||||
"codebaseIndexMistralApiKey",
|
||||
"codebaseIndexVercelAiGatewayApiKey",
|
||||
"codebaseIndexNebiusApiKey",
|
||||
"huggingFaceApiKey",
|
||||
"sambaNovaApiKey",
|
||||
"zaiApiKey",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
279
src/services/code-index/embedders/__tests__/nebius.spec.ts
Normal file
279
src/services/code-index/embedders/__tests__/nebius.spec.ts
Normal file
|
|
@ -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()
|
||||
})
|
||||
})
|
||||
})
|
||||
224
src/services/code-index/embedders/nebius.ts
Normal file
224
src/services/code-index/embedders/nebius.ts
Normal file
|
|
@ -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<EmbeddingResponse> {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -302,6 +302,7 @@ export interface WebviewMessage {
|
|||
codebaseIndexGeminiApiKey?: string
|
||||
codebaseIndexMistralApiKey?: string
|
||||
codebaseIndexVercelAiGatewayApiKey?: string
|
||||
codebaseIndexNebiusApiKey?: string
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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.`)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue