Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion packages/types/src/codebase-index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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", "openrouter"])
.optional(),
codebaseIndexEmbedderBaseUrl: z.string().optional(),
codebaseIndexEmbedderModelId: z.string().optional(),
Expand Down Expand Up @@ -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(),
openrouter: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
})

export type CodebaseIndexModels = z.infer<typeof codebaseIndexModelsSchema>
Expand All @@ -68,6 +69,7 @@ export const codebaseIndexProviderSchema = z.object({
codebaseIndexGeminiApiKey: z.string().optional(),
codebaseIndexMistralApiKey: z.string().optional(),
codebaseIndexVercelAiGatewayApiKey: z.string().optional(),
codebaseIndexOpenRouterApiKey: z.string().optional(),
})

export type CodebaseIndexProvider = z.infer<typeof codebaseIndexProviderSchema>
1 change: 1 addition & 0 deletions packages/types/src/global-settings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,7 @@ export const SECRET_STATE_KEYS = [
"codebaseIndexGeminiApiKey",
"codebaseIndexMistralApiKey",
"codebaseIndexVercelAiGatewayApiKey",
"codebaseIndexOpenRouterApiKey",
"huggingFaceApiKey",
"sambaNovaApiKey",
"zaiApiKey",
Expand Down
8 changes: 8 additions & 0 deletions src/core/webview/webviewMessageHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2550,6 +2550,12 @@ export const webviewMessageHandler = async (
settings.codebaseIndexVercelAiGatewayApiKey,
)
}
if (settings.codebaseIndexOpenRouterApiKey !== undefined) {
await provider.contextProxy.storeSecret(
"codebaseIndexOpenRouterApiKey",
settings.codebaseIndexOpenRouterApiKey,
)
}

// Send success response first - settings are saved regardless of validation
await provider.postMessageToWebview({
Expand Down Expand Up @@ -2687,6 +2693,7 @@ export const webviewMessageHandler = async (
const hasVercelAiGatewayApiKey = !!(await provider.context.secrets.get(
"codebaseIndexVercelAiGatewayApiKey",
))
const hasOpenRouterApiKey = !!(await provider.context.secrets.get("codebaseIndexOpenRouterApiKey"))

provider.postMessageToWebview({
type: "codeIndexSecretStatus",
Expand All @@ -2697,6 +2704,7 @@ export const webviewMessageHandler = async (
hasGeminiApiKey,
hasMistralApiKey,
hasVercelAiGatewayApiKey,
hasOpenRouterApiKey,
},
})
break
Expand Down
20 changes: 20 additions & 0 deletions src/services/code-index/config-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export class CodeIndexConfigManager {
private geminiOptions?: { apiKey: string }
private mistralOptions?: { apiKey: string }
private vercelAiGatewayOptions?: { apiKey: string }
private openRouterOptions?: { apiKey: string }
private qdrantUrl?: string = "http://localhost:6333"
private qdrantApiKey?: string
private searchMinScore?: number
Expand Down Expand Up @@ -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 openRouterApiKey = this.contextProxy?.getSecret("codebaseIndexOpenRouterApiKey") ?? ""

// Update instance variables with configuration
this.codebaseIndexEnabled = codebaseIndexEnabled ?? true
Expand Down Expand Up @@ -108,6 +110,8 @@ export class CodeIndexConfigManager {
this.embedderProvider = "mistral"
} else if (codebaseIndexEmbedderProvider === "vercel-ai-gateway") {
this.embedderProvider = "vercel-ai-gateway"
} else if (codebaseIndexEmbedderProvider === "openrouter") {
this.embedderProvider = "openrouter"
} else {
this.embedderProvider = "openai"
}
Expand All @@ -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.openRouterOptions = openRouterApiKey ? { apiKey: openRouterApiKey } : undefined
}

/**
Expand All @@ -147,6 +152,7 @@ export class CodeIndexConfigManager {
geminiOptions?: { apiKey: string }
mistralOptions?: { apiKey: string }
vercelAiGatewayOptions?: { apiKey: string }
openRouterOptions?: { apiKey: string }
qdrantUrl?: string
qdrantApiKey?: string
searchMinScore?: number
Expand All @@ -167,6 +173,7 @@ export class CodeIndexConfigManager {
geminiApiKey: this.geminiOptions?.apiKey ?? "",
mistralApiKey: this.mistralOptions?.apiKey ?? "",
vercelAiGatewayApiKey: this.vercelAiGatewayOptions?.apiKey ?? "",
openRouterApiKey: this.openRouterOptions?.apiKey ?? "",
qdrantUrl: this.qdrantUrl ?? "",
qdrantApiKey: this.qdrantApiKey ?? "",
}
Expand All @@ -192,6 +199,7 @@ export class CodeIndexConfigManager {
geminiOptions: this.geminiOptions,
mistralOptions: this.mistralOptions,
vercelAiGatewayOptions: this.vercelAiGatewayOptions,
openRouterOptions: this.openRouterOptions,
qdrantUrl: this.qdrantUrl,
qdrantApiKey: this.qdrantApiKey,
searchMinScore: this.currentSearchMinScore,
Expand Down Expand Up @@ -234,6 +242,11 @@ export class CodeIndexConfigManager {
const qdrantUrl = this.qdrantUrl
const isConfigured = !!(apiKey && qdrantUrl)
return isConfigured
} else if (this.embedderProvider === "openrouter") {
const apiKey = this.openRouterOptions?.apiKey
const qdrantUrl = this.qdrantUrl
const isConfigured = !!(apiKey && qdrantUrl)
return isConfigured
}
return false // Should not happen if embedderProvider is always set correctly
}
Expand Down Expand Up @@ -269,6 +282,7 @@ export class CodeIndexConfigManager {
const prevGeminiApiKey = prev?.geminiApiKey ?? ""
const prevMistralApiKey = prev?.mistralApiKey ?? ""
const prevVercelAiGatewayApiKey = prev?.vercelAiGatewayApiKey ?? ""
const prevOpenRouterApiKey = prev?.openRouterApiKey ?? ""
const prevQdrantUrl = prev?.qdrantUrl ?? ""
const prevQdrantApiKey = prev?.qdrantApiKey ?? ""

Expand Down Expand Up @@ -307,6 +321,7 @@ export class CodeIndexConfigManager {
const currentGeminiApiKey = this.geminiOptions?.apiKey ?? ""
const currentMistralApiKey = this.mistralOptions?.apiKey ?? ""
const currentVercelAiGatewayApiKey = this.vercelAiGatewayOptions?.apiKey ?? ""
const currentOpenRouterApiKey = this.openRouterOptions?.apiKey ?? ""
const currentQdrantUrl = this.qdrantUrl ?? ""
const currentQdrantApiKey = this.qdrantApiKey ?? ""

Expand Down Expand Up @@ -337,6 +352,10 @@ export class CodeIndexConfigManager {
return true
}

if (prevOpenRouterApiKey !== currentOpenRouterApiKey) {
return true
}

// Check for model dimension changes (generic for all providers)
if (prevModelDimension !== currentModelDimension) {
return true
Expand Down Expand Up @@ -395,6 +414,7 @@ export class CodeIndexConfigManager {
geminiOptions: this.geminiOptions,
mistralOptions: this.mistralOptions,
vercelAiGatewayOptions: this.vercelAiGatewayOptions,
openRouterOptions: this.openRouterOptions,
qdrantUrl: this.qdrantUrl,
qdrantApiKey: this.qdrantApiKey,
searchMinScore: this.currentSearchMinScore,
Expand Down
213 changes: 213 additions & 0 deletions src/services/code-index/embedders/__tests__/openrouter.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,213 @@
import { describe, it, expect, beforeEach, vi } from "vitest"
import { OpenRouterEmbedder } from "../openrouter"
import { getModelDimension, getDefaultModelId } from "../../../../shared/embeddingModels"

// Mock global fetch
const mockFetch = vi.fn()
global.fetch = mockFetch

describe("OpenRouterEmbedder", () => {
const mockApiKey = "test-api-key"

describe("constructor", () => {
it("should create an instance with valid API key", () => {
const embedder = new OpenRouterEmbedder(mockApiKey)
expect(embedder).toBeInstanceOf(OpenRouterEmbedder)
})

it("should throw error with empty API key", () => {
expect(() => new OpenRouterEmbedder("")).toThrow("API key is required")
})

it("should use default model when none specified", () => {
const embedder = new OpenRouterEmbedder(mockApiKey)
const expectedDefault = getDefaultModelId("openrouter")
expect(embedder.embedderInfo.name).toBe("openrouter")
})

it("should use custom model when specified", () => {
const customModel = "openai/text-embedding-3-small"
const embedder = new OpenRouterEmbedder(mockApiKey, customModel)
expect(embedder.embedderInfo.name).toBe("openrouter")
})
})

describe("embedderInfo", () => {
it("should return correct embedder info", () => {
const embedder = new OpenRouterEmbedder(mockApiKey)
expect(embedder.embedderInfo).toEqual({
name: "openrouter",
})
})
})

describe("createEmbeddings", () => {
let embedder: OpenRouterEmbedder

beforeEach(() => {
embedder = new OpenRouterEmbedder(mockApiKey)
mockFetch.mockClear()
})

it("should create embeddings successfully", async () => {
const mockResponse = {
ok: true,
json: vi.fn().mockResolvedValue({
data: [
{
embedding: Buffer.from(new Float32Array([0.1, 0.2, 0.3]).buffer).toString("base64"),
},
],
usage: {
prompt_tokens: 5,
total_tokens: 5,
},
}),
}

mockFetch.mockResolvedValue(mockResponse)

const result = await embedder.createEmbeddings(["test text"])

expect(result.embeddings).toHaveLength(1)
expect(result.embeddings[0]).toEqual([0.1, 0.2, 0.3])
expect(result.usage?.promptTokens).toBe(5)
expect(result.usage?.totalTokens).toBe(5)
})

it("should handle multiple texts", async () => {
const mockResponse = {
ok: true,
json: vi.fn().mockResolvedValue({
data: [
{
embedding: Buffer.from(new Float32Array([0.1, 0.2]).buffer).toString("base64"),
},
{
embedding: Buffer.from(new Float32Array([0.3, 0.4]).buffer).toString("base64"),
},
],
usage: {
prompt_tokens: 10,
total_tokens: 10,
},
}),
}

mockFetch.mockResolvedValue(mockResponse)

const result = await embedder.createEmbeddings(["text1", "text2"])

expect(result.embeddings).toHaveLength(2)
expect(result.embeddings[0]).toEqual([0.1, 0.2])
expect(result.embeddings[1]).toEqual([0.3, 0.4])
})

it("should use custom model when provided", async () => {
const customModel = "mistralai/mistral-embed-2312"
const embedderWithCustomModel = new OpenRouterEmbedder(mockApiKey, customModel)

const mockResponse = {
ok: true,
json: vi.fn().mockResolvedValue({
data: [
{
embedding: Buffer.from(new Float32Array([0.1, 0.2]).buffer).toString("base64"),
},
],
usage: {
prompt_tokens: 5,
total_tokens: 5,
},
}),
}

mockFetch.mockResolvedValue(mockResponse)

await embedderWithCustomModel.createEmbeddings(["test"])

// Verify the fetch was called with the custom model
expect(mockFetch).toHaveBeenCalledWith(
expect.stringContaining("openrouter.ai/api/v1/embeddings"),
expect.objectContaining({
body: expect.stringContaining(`"model":"${customModel}"`),
}),
)
})
})

describe("validateConfiguration", () => {
let embedder: OpenRouterEmbedder

beforeEach(() => {
embedder = new OpenRouterEmbedder(mockApiKey)
mockFetch.mockClear()
})

it("should validate configuration successfully", async () => {
const mockResponse = {
ok: true,
json: vi.fn().mockResolvedValue({
data: [
{
embedding: Buffer.from(new Float32Array([0.1, 0.2]).buffer).toString("base64"),
},
],
}),
}

mockFetch.mockResolvedValue(mockResponse)

const result = await embedder.validateConfiguration()

expect(result.valid).toBe(true)
expect(result.error).toBeUndefined()
})

it("should handle validation failure", async () => {
const mockResponse = {
ok: false,
status: 401,
text: vi.fn().mockResolvedValue("Unauthorized"),
}

mockFetch.mockResolvedValue(mockResponse)

const result = await embedder.validateConfiguration()

expect(result.valid).toBe(false)
expect(result.error).toBeDefined()
})
})

describe("integration with shared models", () => {
it("should work with defined OpenRouter models", () => {
const openRouterModels = [
"openai/text-embedding-3-small",
"openai/text-embedding-3-large",
"openai/text-embedding-ada-002",
"google/gemini-embedding-001",
"mistralai/mistral-embed-2312",
"mistralai/codestral-embed-2505",
"qwen/qwen3-embedding-8b",
]

openRouterModels.forEach((model) => {
const dimension = getModelDimension("openrouter", model)
expect(dimension).toBeDefined()
expect(dimension).toBeGreaterThan(0)

const embedder = new OpenRouterEmbedder(mockApiKey, model)
expect(embedder.embedderInfo.name).toBe("openrouter")
})
})

it("should use correct default model", () => {
const defaultModel = getDefaultModelId("openrouter")
expect(defaultModel).toBe("openai/text-embedding-3-large")

const dimension = getModelDimension("openrouter", defaultModel)
expect(dimension).toBe(3072)
})
})
})
Loading
Loading