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:
Roo Code 2025-10-10 02:29:57 +00:00
parent bdc91b2aef
commit 768e5532bd
12 changed files with 574 additions and 4 deletions

View file

@ -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>

View file

@ -199,6 +199,7 @@ export const SECRET_STATE_KEYS = [
"codebaseIndexGeminiApiKey",
"codebaseIndexMistralApiKey",
"codebaseIndexVercelAiGatewayApiKey",
"codebaseIndexNebiusApiKey",
"huggingFaceApiKey",
"sambaNovaApiKey",
"zaiApiKey",

View file

@ -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

View file

@ -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,

View 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()
})
})
})

View 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
}
}

View file

@ -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
}

View file

@ -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

View file

@ -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

View file

@ -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(

View file

@ -302,6 +302,7 @@ export interface WebviewMessage {
codebaseIndexGeminiApiKey?: string
codebaseIndexMistralApiKey?: string
codebaseIndexVercelAiGatewayApiKey?: string
codebaseIndexNebiusApiKey?: string
}
}

View file

@ -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.`)