mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
Merge 859327122a into b867ec9145
This commit is contained in:
commit
48357b4a29
15 changed files with 865 additions and 2 deletions
|
|
@ -31,6 +31,7 @@ export const codebaseIndexConfigSchema = z.object({
|
|||
"vercel-ai-gateway",
|
||||
"bedrock",
|
||||
"openrouter",
|
||||
"roo",
|
||||
])
|
||||
.optional(),
|
||||
codebaseIndexEmbedderBaseUrl: z.string().optional(),
|
||||
|
|
@ -67,6 +68,7 @@ export const codebaseIndexModelsSchema = z.object({
|
|||
"vercel-ai-gateway": z.record(z.string(), z.object({ dimension: z.number() })).optional(),
|
||||
openrouter: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
|
||||
bedrock: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
|
||||
roo: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
|
||||
})
|
||||
|
||||
export type CodebaseIndexModels = z.infer<typeof codebaseIndexModelsSchema>
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ export type EmbedderProvider =
|
|||
| "mistral"
|
||||
| "vercel-ai-gateway"
|
||||
| "bedrock"
|
||||
| "openrouter" // Add other providers as needed.
|
||||
| "openrouter"
|
||||
| "roo" // Add other providers as needed.
|
||||
|
||||
export interface EmbeddingModelProfile {
|
||||
dimension: number
|
||||
|
|
|
|||
|
|
@ -596,6 +596,7 @@ export interface WebviewMessage {
|
|||
| "vercel-ai-gateway"
|
||||
| "bedrock"
|
||||
| "openrouter"
|
||||
| "roo"
|
||||
codebaseIndexEmbedderBaseUrl?: string
|
||||
codebaseIndexEmbedderModelId: string
|
||||
codebaseIndexEmbedderModelDimension?: number // Generic dimension for all providers
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@
|
|||
"geminiConfigMissing": "Gemini configuration missing for embedder creation",
|
||||
"mistralConfigMissing": "Mistral configuration missing for embedder creation",
|
||||
"openRouterConfigMissing": "OpenRouter configuration missing for embedder creation",
|
||||
"rooConfigMissing": "Roo Code Cloud authentication required for embedder creation. Please sign in to Roo Code Cloud.",
|
||||
"vercelAiGatewayConfigMissing": "Vercel AI Gateway configuration missing for embedder creation",
|
||||
"bedrockConfigMissing": "Amazon Bedrock configuration missing for embedder creation",
|
||||
"invalidEmbedderType": "Invalid embedder type configured: {{embedderProvider}}",
|
||||
|
|
|
|||
|
|
@ -6,6 +6,25 @@ import { PreviousConfigSnapshot } from "../interfaces/config"
|
|||
// Mock ContextProxy
|
||||
vi.mock("../../../core/config/ContextProxy")
|
||||
|
||||
// Mock CloudService - use vi.hoisted so variables are available when vi.mock runs
|
||||
const { mockIsAuthenticated, mockCloudHasInstance } = vi.hoisted(() => ({
|
||||
mockIsAuthenticated: vi.fn().mockReturnValue(false),
|
||||
mockCloudHasInstance: vi.fn().mockReturnValue(false),
|
||||
}))
|
||||
vi.mock("@roo-code/cloud", () => ({
|
||||
CloudService: {
|
||||
hasInstance: mockCloudHasInstance,
|
||||
get instance() {
|
||||
return {
|
||||
isAuthenticated: mockIsAuthenticated,
|
||||
authService: {
|
||||
getSessionToken: vi.fn().mockReturnValue("test-session-token"),
|
||||
},
|
||||
}
|
||||
},
|
||||
},
|
||||
}))
|
||||
|
||||
// Mock embeddingModels module
|
||||
vi.mock("../../../shared/embeddingModels")
|
||||
|
||||
|
|
@ -1684,6 +1703,57 @@ describe("CodeIndexConfigManager", () => {
|
|||
expect(configManager.isConfigured()).toBe(false)
|
||||
})
|
||||
|
||||
it("should return true when Roo provider is authenticated and Qdrant configured", () => {
|
||||
mockCloudHasInstance.mockReturnValue(true)
|
||||
mockIsAuthenticated.mockReturnValue(true)
|
||||
|
||||
mockContextProxy.getGlobalState.mockReturnValue({
|
||||
codebaseIndexEnabled: true,
|
||||
codebaseIndexEmbedderProvider: "roo",
|
||||
codebaseIndexQdrantUrl: "http://localhost:6333",
|
||||
})
|
||||
mockContextProxy.getSecret.mockReturnValue(undefined)
|
||||
|
||||
configManager = new CodeIndexConfigManager(mockContextProxy)
|
||||
expect(configManager.isConfigured()).toBe(true)
|
||||
|
||||
// Cleanup
|
||||
mockCloudHasInstance.mockReturnValue(false)
|
||||
mockIsAuthenticated.mockReturnValue(false)
|
||||
})
|
||||
|
||||
it("should return false when Roo provider is not authenticated", () => {
|
||||
mockCloudHasInstance.mockReturnValue(true)
|
||||
mockIsAuthenticated.mockReturnValue(false)
|
||||
|
||||
mockContextProxy.getGlobalState.mockReturnValue({
|
||||
codebaseIndexEnabled: true,
|
||||
codebaseIndexEmbedderProvider: "roo",
|
||||
codebaseIndexQdrantUrl: "http://localhost:6333",
|
||||
})
|
||||
mockContextProxy.getSecret.mockReturnValue(undefined)
|
||||
|
||||
configManager = new CodeIndexConfigManager(mockContextProxy)
|
||||
expect(configManager.isConfigured()).toBe(false)
|
||||
|
||||
// Cleanup
|
||||
mockCloudHasInstance.mockReturnValue(false)
|
||||
})
|
||||
|
||||
it("should return false when Roo provider has no CloudService instance", () => {
|
||||
mockCloudHasInstance.mockReturnValue(false)
|
||||
|
||||
mockContextProxy.getGlobalState.mockReturnValue({
|
||||
codebaseIndexEnabled: true,
|
||||
codebaseIndexEmbedderProvider: "roo",
|
||||
codebaseIndexQdrantUrl: "http://localhost:6333",
|
||||
})
|
||||
mockContextProxy.getSecret.mockReturnValue(undefined)
|
||||
|
||||
configManager = new CodeIndexConfigManager(mockContextProxy)
|
||||
expect(configManager.isConfigured()).toBe(false)
|
||||
})
|
||||
|
||||
describe("currentModelDimension", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
|
|
|
|||
|
|
@ -5,12 +5,14 @@ import { CodeIndexOllamaEmbedder } from "../embedders/ollama"
|
|||
import { OpenAICompatibleEmbedder } from "../embedders/openai-compatible"
|
||||
import { GeminiEmbedder } from "../embedders/gemini"
|
||||
import { QdrantVectorStore } from "../vector-store/qdrant-client"
|
||||
import { RooEmbedder } from "../embedders/roo"
|
||||
|
||||
// Mock the embedders and vector store
|
||||
vitest.mock("../embedders/openai")
|
||||
vitest.mock("../embedders/ollama")
|
||||
vitest.mock("../embedders/openai-compatible")
|
||||
vitest.mock("../embedders/gemini")
|
||||
vitest.mock("../embedders/roo")
|
||||
vitest.mock("../vector-store/qdrant-client")
|
||||
|
||||
// Mock the embedding models module
|
||||
|
|
@ -24,6 +26,7 @@ const MockedCodeIndexOllamaEmbedder = CodeIndexOllamaEmbedder as MockedClass<typ
|
|||
const MockedOpenAICompatibleEmbedder = OpenAICompatibleEmbedder as MockedClass<typeof OpenAICompatibleEmbedder>
|
||||
const MockedGeminiEmbedder = GeminiEmbedder as MockedClass<typeof GeminiEmbedder>
|
||||
const MockedQdrantVectorStore = QdrantVectorStore as MockedClass<typeof QdrantVectorStore>
|
||||
const MockedRooEmbedder = RooEmbedder as MockedClass<typeof RooEmbedder>
|
||||
|
||||
// Import the mocked functions
|
||||
import { getDefaultModelId, getModelDimension } from "../../../shared/embeddingModels"
|
||||
|
|
@ -336,6 +339,21 @@ describe("CodeIndexServiceFactory", () => {
|
|||
expect(() => factory.createEmbedder()).toThrow("serviceFactory.geminiConfigMissing")
|
||||
})
|
||||
|
||||
it("should create RooEmbedder when using Roo provider", () => {
|
||||
// Arrange
|
||||
const testConfig = {
|
||||
embedderProvider: "roo",
|
||||
modelId: "text-embedding-3-small",
|
||||
}
|
||||
mockConfigManager.getConfig.mockReturnValue(testConfig as any)
|
||||
|
||||
// Act
|
||||
factory.createEmbedder()
|
||||
|
||||
// Assert
|
||||
expect(MockedRooEmbedder).toHaveBeenCalledWith("text-embedding-3-small")
|
||||
})
|
||||
|
||||
it("should throw error for invalid embedder provider", () => {
|
||||
// Arrange
|
||||
const testConfig = {
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import { CloudService } from "@roo-code/cloud"
|
||||
|
||||
import { ApiHandlerOptions } from "../../shared/api"
|
||||
import { ContextProxy } from "../../core/config/ContextProxy"
|
||||
import { EmbedderProvider } from "./interfaces/manager"
|
||||
|
|
@ -120,6 +122,8 @@ export class CodeIndexConfigManager {
|
|||
this.embedderProvider = "bedrock"
|
||||
} else if (codebaseIndexEmbedderProvider === "openrouter") {
|
||||
this.embedderProvider = "openrouter"
|
||||
} else if (codebaseIndexEmbedderProvider === "roo") {
|
||||
this.embedderProvider = "roo"
|
||||
} else {
|
||||
this.embedderProvider = "openai"
|
||||
}
|
||||
|
|
@ -272,6 +276,12 @@ export class CodeIndexConfigManager {
|
|||
const qdrantUrl = this.qdrantUrl
|
||||
const isConfigured = !!(apiKey && qdrantUrl)
|
||||
return isConfigured
|
||||
} else if (this.embedderProvider === "roo") {
|
||||
// Roo Code Router uses CloudService session token for auth.
|
||||
// Use isAuthenticated() for a stable auth check (not raw session token parsing).
|
||||
const qdrantUrl = this.qdrantUrl
|
||||
const isAuthenticated = CloudService.hasInstance() && CloudService.instance.isAuthenticated()
|
||||
return !!(isAuthenticated && qdrantUrl)
|
||||
}
|
||||
return false // Should not happen if embedderProvider is always set correctly
|
||||
}
|
||||
|
|
|
|||
282
src/services/code-index/embedders/__tests__/roo.spec.ts
Normal file
282
src/services/code-index/embedders/__tests__/roo.spec.ts
Normal file
|
|
@ -0,0 +1,282 @@
|
|||
import type { MockedClass, MockedFunction } from "vitest"
|
||||
import { describe, it, expect, beforeEach, vi, afterEach } from "vitest"
|
||||
import { OpenAI } from "openai"
|
||||
import { RooEmbedder } from "../roo"
|
||||
import { getDefaultModelId } from "../../../../shared/embeddingModels"
|
||||
|
||||
// Mock the OpenAI SDK
|
||||
vi.mock("openai")
|
||||
|
||||
// Mock CloudService - use vi.hoisted so variables are available when vi.mock runs
|
||||
const { mockGetSessionToken, mockCloudIsAuthenticated, mockCloudHasInstance } = vi.hoisted(() => ({
|
||||
mockGetSessionToken: vi.fn().mockReturnValue("test-session-token"),
|
||||
mockCloudIsAuthenticated: vi.fn().mockReturnValue(true),
|
||||
mockCloudHasInstance: vi.fn().mockReturnValue(true),
|
||||
}))
|
||||
vi.mock("@roo-code/cloud", () => ({
|
||||
CloudService: {
|
||||
hasInstance: mockCloudHasInstance,
|
||||
get instance() {
|
||||
return {
|
||||
authService: {
|
||||
getSessionToken: mockGetSessionToken,
|
||||
},
|
||||
isAuthenticated: mockCloudIsAuthenticated,
|
||||
}
|
||||
},
|
||||
},
|
||||
}))
|
||||
|
||||
// Mock TelemetryService
|
||||
vi.mock("@roo-code/telemetry", () => ({
|
||||
TelemetryService: {
|
||||
instance: {
|
||||
captureEvent: vi.fn(),
|
||||
},
|
||||
},
|
||||
TelemetryEventName: {},
|
||||
}))
|
||||
|
||||
// Mock i18n
|
||||
vi.mock("../../../../i18n", () => ({
|
||||
t: (key: string, params?: Record<string, any>) => {
|
||||
const translations: Record<string, string> = {
|
||||
"embeddings:validation.apiKeyRequired": "validation.apiKeyRequired",
|
||||
"embeddings:authenticationFailed":
|
||||
"Failed to create embeddings: Authentication failed. Please check your API key.",
|
||||
"embeddings:failedWithStatus": `Failed to create embeddings after ${params?.attempts} attempts: HTTP ${params?.statusCode} - ${params?.errorMessage}`,
|
||||
"embeddings:failedWithError": `Failed to create embeddings after ${params?.attempts} attempts: ${params?.errorMessage}`,
|
||||
"embeddings:failedMaxAttempts": `Failed to create embeddings after ${params?.attempts} attempts`,
|
||||
"embeddings:textExceedsTokenLimit": `Text at index ${params?.index} exceeds maximum token limit (${params?.itemTokens} > ${params?.maxTokens}). Skipping.`,
|
||||
"embeddings:rateLimitRetry": `Rate limit hit, retrying in ${params?.delayMs}ms (attempt ${params?.attempt}/${params?.maxRetries})`,
|
||||
}
|
||||
return translations[key] || key
|
||||
},
|
||||
}))
|
||||
|
||||
const MockedOpenAI = OpenAI as MockedClass<typeof OpenAI>
|
||||
|
||||
describe("RooEmbedder", () => {
|
||||
let mockEmbeddingsCreate: MockedFunction<any>
|
||||
let mockOpenAIInstance: any
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
vi.spyOn(console, "warn").mockImplementation(() => {})
|
||||
vi.spyOn(console, "error").mockImplementation(() => {})
|
||||
|
||||
// Setup mock OpenAI instance
|
||||
mockEmbeddingsCreate = vi.fn()
|
||||
mockOpenAIInstance = {
|
||||
embeddings: {
|
||||
create: mockEmbeddingsCreate,
|
||||
},
|
||||
apiKey: "test-session-token",
|
||||
}
|
||||
|
||||
MockedOpenAI.mockImplementation(() => mockOpenAIInstance)
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks()
|
||||
})
|
||||
|
||||
describe("constructor", () => {
|
||||
it("should create an instance using CloudService session token", () => {
|
||||
const embedder = new RooEmbedder()
|
||||
expect(embedder).toBeInstanceOf(RooEmbedder)
|
||||
})
|
||||
|
||||
it("should use default model when none specified", () => {
|
||||
const embedder = new RooEmbedder()
|
||||
const expectedDefault = getDefaultModelId("roo")
|
||||
expect(expectedDefault).toBe("text-embedding-3-small")
|
||||
expect(embedder.embedderInfo.name).toBe("roo")
|
||||
})
|
||||
|
||||
it("should use custom model when specified", () => {
|
||||
const customModel = "text-embedding-3-large"
|
||||
const embedder = new RooEmbedder(customModel)
|
||||
expect(embedder.embedderInfo.name).toBe("roo")
|
||||
})
|
||||
|
||||
it("should initialize OpenAI client with correct headers and base URL ending in /v1", () => {
|
||||
new RooEmbedder()
|
||||
|
||||
// Verify client was created with correct headers
|
||||
const callArgs = MockedOpenAI.mock.calls[0][0] as any
|
||||
expect(callArgs.defaultHeaders).toEqual({
|
||||
"HTTP-Referer": "https://github.com/RooCodeInc/Roo-Code",
|
||||
"X-Title": "Roo Code",
|
||||
})
|
||||
|
||||
// Verify the baseURL ends with /v1
|
||||
expect(callArgs.baseURL).toMatch(/\/v1$/)
|
||||
})
|
||||
})
|
||||
|
||||
describe("createEmbeddings", () => {
|
||||
it("should create embeddings for a batch of texts", async () => {
|
||||
// Create a proper base64-encoded float32 array
|
||||
const float32Array = new Float32Array([0.1, 0.2, 0.3])
|
||||
const buffer = Buffer.from(float32Array.buffer)
|
||||
const base64Embedding = buffer.toString("base64")
|
||||
|
||||
mockEmbeddingsCreate.mockResolvedValueOnce({
|
||||
data: [{ embedding: base64Embedding }, { embedding: base64Embedding }],
|
||||
usage: { prompt_tokens: 10, total_tokens: 15 },
|
||||
})
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
const result = await embedder.createEmbeddings(["text1", "text2"])
|
||||
|
||||
expect(result.embeddings).toHaveLength(2)
|
||||
expect(result.embeddings[0]).toHaveLength(3) // 3 floats in our test array
|
||||
expect(result.usage?.promptTokens).toBe(10)
|
||||
expect(result.usage?.totalTokens).toBe(15)
|
||||
})
|
||||
|
||||
it("should request base64 encoding format", async () => {
|
||||
const float32Array = new Float32Array([0.1])
|
||||
const buffer = Buffer.from(float32Array.buffer)
|
||||
const base64Embedding = buffer.toString("base64")
|
||||
|
||||
mockEmbeddingsCreate.mockResolvedValueOnce({
|
||||
data: [{ embedding: base64Embedding }],
|
||||
usage: { prompt_tokens: 5, total_tokens: 5 },
|
||||
})
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
await embedder.createEmbeddings(["test"])
|
||||
|
||||
expect(mockEmbeddingsCreate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
encoding_format: "base64",
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should refresh session token before making request", async () => {
|
||||
const float32Array = new Float32Array([0.1])
|
||||
const buffer = Buffer.from(float32Array.buffer)
|
||||
const base64Embedding = buffer.toString("base64")
|
||||
|
||||
mockEmbeddingsCreate.mockResolvedValueOnce({
|
||||
data: [{ embedding: base64Embedding }],
|
||||
usage: { prompt_tokens: 5, total_tokens: 5 },
|
||||
})
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
await embedder.createEmbeddings(["test"])
|
||||
|
||||
// The apiKey should have been refreshed via the setter
|
||||
// (we can't directly verify the setter was called since it's
|
||||
// a mock object, but the call to createEmbeddings succeeds)
|
||||
expect(mockEmbeddingsCreate).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should handle numeric array embeddings (non-base64)", async () => {
|
||||
mockEmbeddingsCreate.mockResolvedValueOnce({
|
||||
data: [{ embedding: [0.1, 0.2, 0.3] }],
|
||||
usage: { prompt_tokens: 5, total_tokens: 5 },
|
||||
})
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
const result = await embedder.createEmbeddings(["test"])
|
||||
|
||||
expect(result.embeddings).toHaveLength(1)
|
||||
expect(result.embeddings[0]).toEqual([0.1, 0.2, 0.3])
|
||||
})
|
||||
|
||||
it("should skip texts exceeding token limit", async () => {
|
||||
const longText = "a".repeat(400000) // Exceeds MAX_ITEM_TOKENS
|
||||
|
||||
const float32Array = new Float32Array([0.1])
|
||||
const buffer = Buffer.from(float32Array.buffer)
|
||||
const base64Embedding = buffer.toString("base64")
|
||||
|
||||
mockEmbeddingsCreate.mockResolvedValueOnce({
|
||||
data: [{ embedding: base64Embedding }],
|
||||
usage: { prompt_tokens: 5, total_tokens: 5 },
|
||||
})
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
const result = await embedder.createEmbeddings([longText, "short text"])
|
||||
|
||||
// Only the short text should have been embedded
|
||||
expect(result.embeddings).toHaveLength(1)
|
||||
})
|
||||
})
|
||||
|
||||
describe("validateConfiguration", () => {
|
||||
it("should return valid when API responds correctly", async () => {
|
||||
const float32Array = new Float32Array([0.1])
|
||||
const buffer = Buffer.from(float32Array.buffer)
|
||||
const base64Embedding = buffer.toString("base64")
|
||||
|
||||
mockEmbeddingsCreate.mockResolvedValueOnce({
|
||||
data: [{ embedding: base64Embedding }],
|
||||
usage: { prompt_tokens: 1, total_tokens: 1 },
|
||||
})
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
const result = await embedder.validateConfiguration()
|
||||
|
||||
expect(result.valid).toBe(true)
|
||||
})
|
||||
|
||||
it("should return invalid when API returns empty response", async () => {
|
||||
mockEmbeddingsCreate.mockResolvedValueOnce({
|
||||
data: [],
|
||||
usage: { prompt_tokens: 0, total_tokens: 0 },
|
||||
})
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
const result = await embedder.validateConfiguration()
|
||||
|
||||
expect(result.valid).toBe(false)
|
||||
})
|
||||
|
||||
it("should handle API errors during validation", async () => {
|
||||
const apiError = new Error("API connection failed")
|
||||
;(apiError as any).status = 500
|
||||
|
||||
mockEmbeddingsCreate.mockRejectedValueOnce(apiError)
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
const result = await embedder.validateConfiguration()
|
||||
|
||||
expect(result.valid).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("embedderInfo", () => {
|
||||
it("should return roo as the embedder name", () => {
|
||||
const embedder = new RooEmbedder()
|
||||
expect(embedder.embedderInfo).toEqual({ name: "roo" })
|
||||
})
|
||||
})
|
||||
|
||||
describe("rate limiting", () => {
|
||||
it("should retry on 429 errors with exponential backoff", async () => {
|
||||
const rateLimitError = new Error("Rate limit exceeded")
|
||||
;(rateLimitError as any).status = 429
|
||||
|
||||
const float32Array = new Float32Array([0.1])
|
||||
const buffer = Buffer.from(float32Array.buffer)
|
||||
const base64Embedding = buffer.toString("base64")
|
||||
|
||||
// First call fails with 429, second succeeds
|
||||
mockEmbeddingsCreate.mockRejectedValueOnce(rateLimitError).mockResolvedValueOnce({
|
||||
data: [{ embedding: base64Embedding }],
|
||||
usage: { prompt_tokens: 5, total_tokens: 5 },
|
||||
})
|
||||
|
||||
const embedder = new RooEmbedder()
|
||||
const result = await embedder.createEmbeddings(["test"])
|
||||
|
||||
expect(result.embeddings).toHaveLength(1)
|
||||
expect(mockEmbeddingsCreate).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
})
|
||||
})
|
||||
393
src/services/code-index/embedders/roo.ts
Normal file
393
src/services/code-index/embedders/roo.ts
Normal file
|
|
@ -0,0 +1,393 @@
|
|||
import { OpenAI } from "openai"
|
||||
import { CloudService } from "@roo-code/cloud"
|
||||
import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder"
|
||||
import {
|
||||
MAX_BATCH_TOKENS,
|
||||
MAX_ITEM_TOKENS,
|
||||
MAX_BATCH_RETRIES as MAX_RETRIES,
|
||||
INITIAL_RETRY_DELAY_MS as INITIAL_DELAY_MS,
|
||||
} from "../constants"
|
||||
import { getDefaultModelId, getModelQueryPrefix } from "../../../shared/embeddingModels"
|
||||
import { t } from "../../../i18n"
|
||||
import { withValidationErrorHandling, HttpError, formatEmbeddingError } from "../shared/validation-helpers"
|
||||
import { TelemetryEventName } from "@roo-code/types"
|
||||
import { TelemetryService } from "@roo-code/telemetry"
|
||||
import { Mutex } from "async-mutex"
|
||||
import { handleOpenAIError } from "../../../api/providers/utils/openai-error-handler"
|
||||
|
||||
interface EmbeddingItem {
|
||||
embedding: string | number[]
|
||||
[key: string]: any
|
||||
}
|
||||
|
||||
interface RooEmbeddingResponse {
|
||||
data: EmbeddingItem[]
|
||||
usage?: {
|
||||
prompt_tokens?: number
|
||||
total_tokens?: number
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the current session token from CloudService, or "unauthenticated" if unavailable.
|
||||
*/
|
||||
function getSessionToken(): string {
|
||||
const token = CloudService.hasInstance() ? CloudService.instance.authService?.getSessionToken() : undefined
|
||||
return token ?? "unauthenticated"
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the base URL for the Roo Code Router proxy, with /v1 suffix.
|
||||
*/
|
||||
function getRooBaseUrl(): string {
|
||||
let baseURL = process.env.ROO_CODE_PROVIDER_URL ?? "https://api.roocode.com/proxy"
|
||||
if (!baseURL.endsWith("/v1")) {
|
||||
baseURL = `${baseURL}/v1`
|
||||
}
|
||||
return baseURL
|
||||
}
|
||||
|
||||
/**
|
||||
* Roo Code Router implementation of the embedder interface with batching and rate limiting.
|
||||
* Uses CloudService session token for authentication against the Roo Code proxy,
|
||||
* which forwards embedding requests to the underlying model provider (e.g. OpenAI).
|
||||
* No third-party API key is required -- users authenticate via Roo Code Cloud.
|
||||
*/
|
||||
export class RooEmbedder implements IEmbedder {
|
||||
private embeddingsClient: OpenAI
|
||||
private readonly defaultModelId: string
|
||||
private readonly maxItemTokens: number
|
||||
|
||||
// Global rate limiting state shared across all instances
|
||||
private static globalRateLimitState = {
|
||||
isRateLimited: false,
|
||||
rateLimitResetTime: 0,
|
||||
consecutiveRateLimitErrors: 0,
|
||||
lastRateLimitError: 0,
|
||||
mutex: new Mutex(),
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a new Roo Code Router embedder.
|
||||
* Authentication is handled via CloudService session token.
|
||||
* @param modelId Optional model identifier (defaults to "text-embedding-3-small")
|
||||
* @param maxItemTokens Optional maximum tokens per item (defaults to MAX_ITEM_TOKENS)
|
||||
*/
|
||||
constructor(modelId?: string, maxItemTokens?: number) {
|
||||
const sessionToken = getSessionToken()
|
||||
const baseURL = getRooBaseUrl()
|
||||
|
||||
try {
|
||||
this.embeddingsClient = new OpenAI({
|
||||
baseURL,
|
||||
apiKey: sessionToken,
|
||||
defaultHeaders: {
|
||||
"HTTP-Referer": "https://github.com/RooCodeInc/Roo-Code",
|
||||
"X-Title": "Roo Code",
|
||||
},
|
||||
})
|
||||
} catch (error) {
|
||||
throw handleOpenAIError(error, "Roo Code Cloud")
|
||||
}
|
||||
|
||||
this.defaultModelId = modelId || getDefaultModelId("roo")
|
||||
this.maxItemTokens = maxItemTokens || MAX_ITEM_TOKENS
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates embeddings for the given texts with batching and rate limiting.
|
||||
* Refreshes the session token before each top-level call to ensure freshness.
|
||||
* @param texts Array of text strings to embed
|
||||
* @param model Optional model identifier
|
||||
* @returns Promise resolving to embedding response
|
||||
*/
|
||||
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
|
||||
// Refresh API key to use the latest session token
|
||||
this.embeddingsClient.apiKey = getSessionToken()
|
||||
|
||||
const modelToUse = model || this.defaultModelId
|
||||
|
||||
// Apply model-specific query prefix if required
|
||||
const queryPrefix = getModelQueryPrefix("roo", modelToUse)
|
||||
const processedTexts = queryPrefix
|
||||
? texts.map((text, index) => {
|
||||
if (text.startsWith(queryPrefix)) {
|
||||
return text
|
||||
}
|
||||
const prefixedText = `${queryPrefix}${text}`
|
||||
const estimatedTokens = Math.ceil(prefixedText.length / 4)
|
||||
if (estimatedTokens > MAX_ITEM_TOKENS) {
|
||||
console.warn(
|
||||
t("embeddings:textWithPrefixExceedsTokenLimit", {
|
||||
index,
|
||||
estimatedTokens,
|
||||
maxTokens: MAX_ITEM_TOKENS,
|
||||
}),
|
||||
)
|
||||
return text
|
||||
}
|
||||
return prefixedText
|
||||
})
|
||||
: texts
|
||||
|
||||
const allEmbeddings: number[][] = []
|
||||
const usage = { promptTokens: 0, totalTokens: 0 }
|
||||
const remainingTexts = [...processedTexts]
|
||||
|
||||
while (remainingTexts.length > 0) {
|
||||
const currentBatch: string[] = []
|
||||
let currentBatchTokens = 0
|
||||
const processedIndices: number[] = []
|
||||
|
||||
for (let i = 0; i < remainingTexts.length; i++) {
|
||||
const text = remainingTexts[i]
|
||||
const itemTokens = Math.ceil(text.length / 4)
|
||||
|
||||
if (itemTokens > this.maxItemTokens) {
|
||||
console.warn(
|
||||
t("embeddings:textExceedsTokenLimit", {
|
||||
index: i,
|
||||
itemTokens,
|
||||
maxTokens: this.maxItemTokens,
|
||||
}),
|
||||
)
|
||||
processedIndices.push(i)
|
||||
continue
|
||||
}
|
||||
|
||||
if (currentBatchTokens + itemTokens <= MAX_BATCH_TOKENS) {
|
||||
currentBatch.push(text)
|
||||
currentBatchTokens += itemTokens
|
||||
processedIndices.push(i)
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Remove processed items from remainingTexts (in reverse order to maintain correct indices)
|
||||
for (let i = processedIndices.length - 1; i >= 0; i--) {
|
||||
remainingTexts.splice(processedIndices[i], 1)
|
||||
}
|
||||
|
||||
if (currentBatch.length > 0) {
|
||||
const batchResult = await this._embedBatchWithRetries(currentBatch, modelToUse)
|
||||
allEmbeddings.push(...batchResult.embeddings)
|
||||
usage.promptTokens += batchResult.usage.promptTokens
|
||||
usage.totalTokens += batchResult.usage.totalTokens
|
||||
}
|
||||
}
|
||||
|
||||
return { embeddings: allEmbeddings, usage }
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper method to handle batch embedding with retries and exponential backoff.
|
||||
* Refreshes the session token on each retry attempt.
|
||||
*/
|
||||
private async _embedBatchWithRetries(
|
||||
batchTexts: string[],
|
||||
model: string,
|
||||
): Promise<{ embeddings: number[][]; usage: { promptTokens: number; totalTokens: number } }> {
|
||||
for (let attempts = 0; attempts < MAX_RETRIES; attempts++) {
|
||||
await this.waitForGlobalRateLimit()
|
||||
|
||||
// Refresh token on each attempt
|
||||
this.embeddingsClient.apiKey = getSessionToken()
|
||||
|
||||
try {
|
||||
const requestParams: any = {
|
||||
input: batchTexts,
|
||||
model: model,
|
||||
// Request base64 encoding to bypass the OpenAI package's parser
|
||||
// which truncates embedding dimensions to 256 for numeric arrays.
|
||||
encoding_format: "base64",
|
||||
}
|
||||
|
||||
const response = (await this.embeddingsClient.embeddings.create(requestParams)) as RooEmbeddingResponse
|
||||
|
||||
// Convert base64 embeddings to float32 arrays
|
||||
const processedEmbeddings = response.data.map((item: EmbeddingItem) => {
|
||||
if (typeof item.embedding === "string") {
|
||||
const buffer = Buffer.from(item.embedding, "base64")
|
||||
const float32Array = new Float32Array(buffer.buffer, buffer.byteOffset, buffer.byteLength / 4)
|
||||
return {
|
||||
...item,
|
||||
embedding: Array.from(float32Array),
|
||||
}
|
||||
}
|
||||
return item
|
||||
})
|
||||
|
||||
response.data = processedEmbeddings
|
||||
|
||||
const embeddings = response.data.map((item) => item.embedding as number[])
|
||||
|
||||
return {
|
||||
embeddings,
|
||||
usage: {
|
||||
promptTokens: response.usage?.prompt_tokens || 0,
|
||||
totalTokens: response.usage?.total_tokens || 0,
|
||||
},
|
||||
}
|
||||
} 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: "RooEmbedder:_embedBatchWithRetries",
|
||||
attempt: attempts + 1,
|
||||
})
|
||||
|
||||
const hasMoreAttempts = attempts < MAX_RETRIES - 1
|
||||
|
||||
const httpError = error as HttpError
|
||||
if (httpError?.status === 429) {
|
||||
await this.updateGlobalRateLimitState(httpError)
|
||||
|
||||
if (hasMoreAttempts) {
|
||||
const baseDelay = INITIAL_DELAY_MS * Math.pow(2, attempts)
|
||||
const globalDelay = await this.getGlobalRateLimitDelay()
|
||||
const delayMs = Math.max(baseDelay, globalDelay)
|
||||
|
||||
console.warn(
|
||||
t("embeddings:rateLimitRetry", {
|
||||
delayMs,
|
||||
attempt: attempts + 1,
|
||||
maxRetries: MAX_RETRIES,
|
||||
}),
|
||||
)
|
||||
await new Promise((resolve) => setTimeout(resolve, delayMs))
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
console.error(`Roo embedder error (attempt ${attempts + 1}/${MAX_RETRIES}):`, error)
|
||||
throw formatEmbeddingError(error, MAX_RETRIES)
|
||||
}
|
||||
}
|
||||
|
||||
throw new Error(t("embeddings:failedMaxAttempts", { attempts: MAX_RETRIES }))
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates the Roo embedder configuration by testing API connectivity.
|
||||
*/
|
||||
async validateConfiguration(): Promise<{ valid: boolean; error?: string }> {
|
||||
return withValidationErrorHandling(async () => {
|
||||
// Refresh token before validation
|
||||
this.embeddingsClient.apiKey = getSessionToken()
|
||||
|
||||
try {
|
||||
const testTexts = ["test"]
|
||||
const modelToUse = this.defaultModelId
|
||||
|
||||
const requestParams: any = {
|
||||
input: testTexts,
|
||||
model: modelToUse,
|
||||
encoding_format: "base64",
|
||||
}
|
||||
|
||||
const response = (await this.embeddingsClient.embeddings.create(requestParams)) as RooEmbeddingResponse
|
||||
|
||||
if (!response?.data || response.data.length === 0) {
|
||||
return {
|
||||
valid: false,
|
||||
error: "embeddings:validation.invalidResponse",
|
||||
}
|
||||
}
|
||||
|
||||
return { valid: true }
|
||||
} 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: "RooEmbedder:validateConfiguration",
|
||||
})
|
||||
throw error
|
||||
}
|
||||
}, "roo")
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns information about this embedder.
|
||||
*/
|
||||
get embedderInfo(): EmbedderInfo {
|
||||
return {
|
||||
name: "roo",
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Waits if there's an active global rate limit.
|
||||
*/
|
||||
private async waitForGlobalRateLimit(): Promise<void> {
|
||||
const release = await RooEmbedder.globalRateLimitState.mutex.acquire()
|
||||
let mutexReleased = false
|
||||
|
||||
try {
|
||||
const state = RooEmbedder.globalRateLimitState
|
||||
|
||||
if (state.isRateLimited && state.rateLimitResetTime > Date.now()) {
|
||||
const waitTime = state.rateLimitResetTime - Date.now()
|
||||
release()
|
||||
mutexReleased = true
|
||||
await new Promise((resolve) => setTimeout(resolve, waitTime))
|
||||
return
|
||||
}
|
||||
|
||||
if (state.isRateLimited && state.rateLimitResetTime <= Date.now()) {
|
||||
state.isRateLimited = false
|
||||
state.consecutiveRateLimitErrors = 0
|
||||
}
|
||||
} finally {
|
||||
if (!mutexReleased) {
|
||||
release()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Updates global rate limit state when a 429 error occurs.
|
||||
*/
|
||||
private async updateGlobalRateLimitState(error: HttpError): Promise<void> {
|
||||
const release = await RooEmbedder.globalRateLimitState.mutex.acquire()
|
||||
try {
|
||||
const state = RooEmbedder.globalRateLimitState
|
||||
const now = Date.now()
|
||||
|
||||
if (now - state.lastRateLimitError < 60000) {
|
||||
state.consecutiveRateLimitErrors++
|
||||
} else {
|
||||
state.consecutiveRateLimitErrors = 1
|
||||
}
|
||||
|
||||
state.lastRateLimitError = now
|
||||
|
||||
const baseDelay = 5000
|
||||
const maxDelay = 300000
|
||||
const exponentialDelay = Math.min(baseDelay * Math.pow(2, state.consecutiveRateLimitErrors - 1), maxDelay)
|
||||
|
||||
state.isRateLimited = true
|
||||
state.rateLimitResetTime = now + exponentialDelay
|
||||
} finally {
|
||||
release()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the current global rate limit delay.
|
||||
*/
|
||||
private async getGlobalRateLimitDelay(): Promise<number> {
|
||||
const release = await RooEmbedder.globalRateLimitState.mutex.acquire()
|
||||
try {
|
||||
const state = RooEmbedder.globalRateLimitState
|
||||
|
||||
if (state.isRateLimited && state.rateLimitResetTime > Date.now()) {
|
||||
return state.rateLimitResetTime - Date.now()
|
||||
}
|
||||
|
||||
return 0
|
||||
} finally {
|
||||
release()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -37,6 +37,7 @@ export type AvailableEmbedders =
|
|||
| "vercel-ai-gateway"
|
||||
| "bedrock"
|
||||
| "openrouter"
|
||||
| "roo"
|
||||
|
||||
export interface EmbedderInfo {
|
||||
name: AvailableEmbedders
|
||||
|
|
|
|||
|
|
@ -84,6 +84,7 @@ export type EmbedderProvider =
|
|||
| "vercel-ai-gateway"
|
||||
| "bedrock"
|
||||
| "openrouter"
|
||||
| "roo"
|
||||
|
||||
export interface IndexProgressUpdate {
|
||||
systemStatus: IndexingState
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import { MistralEmbedder } from "./embedders/mistral"
|
|||
import { VercelAiGatewayEmbedder } from "./embedders/vercel-ai-gateway"
|
||||
import { BedrockEmbedder } from "./embedders/bedrock"
|
||||
import { OpenRouterEmbedder } from "./embedders/openrouter"
|
||||
import { RooEmbedder } from "./embedders/roo"
|
||||
import { QdrantVectorStore } from "./vector-store/qdrant-client"
|
||||
import { codeParser, DirectoryScanner, FileWatcher } from "./processors"
|
||||
import { ICodeParser, IEmbedder, IFileWatcher, IVectorStore } from "./interfaces"
|
||||
|
|
@ -101,6 +102,9 @@ export class CodeIndexServiceFactory {
|
|||
undefined, // maxItemTokens
|
||||
config.openRouterOptions.specificProvider,
|
||||
)
|
||||
} else if (provider === "roo") {
|
||||
// Roo Code Router uses CloudService session token -- no API key needed
|
||||
return new RooEmbedder(config.modelId)
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
|
|
|
|||
|
|
@ -87,6 +87,11 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
|
|||
"qwen/qwen3-embedding-4b": { dimension: 2560, scoreThreshold: 0.4 },
|
||||
"qwen/qwen3-embedding-8b": { dimension: 4096, scoreThreshold: 0.4 },
|
||||
},
|
||||
roo: {
|
||||
// Roo Code Router proxies OpenAI embedding models
|
||||
"text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
|
||||
"text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
|
||||
},
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -185,6 +190,9 @@ export function getDefaultModelId(provider: EmbedderProvider): string {
|
|||
case "openrouter":
|
||||
return "openai/text-embedding-3-large"
|
||||
|
||||
case "roo":
|
||||
return "text-embedding-3-small"
|
||||
|
||||
default:
|
||||
// Fallback for unknown providers
|
||||
console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`)
|
||||
|
|
|
|||
|
|
@ -176,6 +176,14 @@ const createValidationSchema = (provider: EmbedderProvider, t: any) => {
|
|||
.min(1, t("settings:codeIndex.validation.modelSelectionRequired")),
|
||||
})
|
||||
|
||||
case "roo":
|
||||
// Roo Code Cloud uses session token auth -- no API key needed
|
||||
return baseSchema.extend({
|
||||
codebaseIndexEmbedderModelId: z
|
||||
.string()
|
||||
.min(1, t("settings:codeIndex.validation.modelSelectionRequired")),
|
||||
})
|
||||
|
||||
default:
|
||||
return baseSchema
|
||||
}
|
||||
|
|
@ -187,7 +195,8 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
|
|||
}) => {
|
||||
const SECRET_PLACEHOLDER = "••••••••••••••••"
|
||||
const { t } = useAppTranslation()
|
||||
const { codebaseIndexConfig, codebaseIndexModels, cwd, apiConfiguration } = useExtensionState()
|
||||
const { codebaseIndexConfig, codebaseIndexModels, cwd, apiConfiguration, cloudIsAuthenticated } =
|
||||
useExtensionState()
|
||||
const [open, setOpen] = useState(false)
|
||||
const [isAdvancedSettingsOpen, setIsAdvancedSettingsOpen] = useState(false)
|
||||
const [isSetupSettingsOpen, setIsSetupSettingsOpen] = useState(false)
|
||||
|
|
@ -761,6 +770,9 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
|
|||
<SelectItem value="openrouter">
|
||||
{t("settings:codeIndex.openRouterProvider")}
|
||||
</SelectItem>
|
||||
<SelectItem value="roo">
|
||||
{t("settings:codeIndex.rooCodeCloudProvider")}
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
|
|
@ -1430,6 +1442,62 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
|
|||
</>
|
||||
)}
|
||||
|
||||
{currentSettings.codebaseIndexEmbedderProvider === "roo" && (
|
||||
<>
|
||||
{/* Auth status for Roo Code Cloud */}
|
||||
<div className="space-y-2">
|
||||
{cloudIsAuthenticated ? (
|
||||
<p className="text-xs text-vscode-descriptionForeground mt-1 mb-0">
|
||||
{t("settings:codeIndex.rooAuthenticated")}
|
||||
</p>
|
||||
) : (
|
||||
<p className="text-xs text-vscode-errorForeground mt-1 mb-0">
|
||||
{t("settings:codeIndex.rooAuthenticationRequired")}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="space-y-2">
|
||||
<label className="text-sm font-medium">
|
||||
{t("settings:codeIndex.modelLabel")}
|
||||
</label>
|
||||
<VSCodeDropdown
|
||||
value={currentSettings.codebaseIndexEmbedderModelId}
|
||||
onChange={(e: any) =>
|
||||
updateSetting("codebaseIndexEmbedderModelId", e.target.value)
|
||||
}
|
||||
className={cn("w-full", {
|
||||
"border-red-500": formErrors.codebaseIndexEmbedderModelId,
|
||||
})}>
|
||||
<VSCodeOption value="" className="p-2">
|
||||
{t("settings:codeIndex.selectModel")}
|
||||
</VSCodeOption>
|
||||
{getAvailableModels().map((modelId) => {
|
||||
const model =
|
||||
codebaseIndexModels?.[
|
||||
currentSettings.codebaseIndexEmbedderProvider as keyof typeof codebaseIndexModels
|
||||
]?.[modelId]
|
||||
return (
|
||||
<VSCodeOption key={modelId} value={modelId} className="p-2">
|
||||
{modelId}{" "}
|
||||
{model
|
||||
? t("settings:codeIndex.modelDimensions", {
|
||||
dimension: model.dimension,
|
||||
})
|
||||
: ""}
|
||||
</VSCodeOption>
|
||||
)
|
||||
})}
|
||||
</VSCodeDropdown>
|
||||
{formErrors.codebaseIndexEmbedderModelId && (
|
||||
<p className="text-xs text-vscode-errorForeground mt-1 mb-0">
|
||||
{formErrors.codebaseIndexEmbedderModelId}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
||||
{/* Qdrant Settings */}
|
||||
<div className="space-y-2">
|
||||
<label className="text-sm font-medium">
|
||||
|
|
|
|||
|
|
@ -187,6 +187,9 @@
|
|||
"bedrockProfileLabel": "AWS Profile",
|
||||
"bedrockProfilePlaceholder": "default",
|
||||
"bedrockProfileDescription": "AWS profile name from ~/.aws/credentials (required).",
|
||||
"rooCodeCloudProvider": "Roo Code Cloud",
|
||||
"rooAuthenticationRequired": "Sign in to Roo Code Cloud to use this provider for codebase indexing.",
|
||||
"rooAuthenticated": "Authenticated with Roo Code Cloud.",
|
||||
"openRouterProvider": "OpenRouter",
|
||||
"openRouterApiKeyLabel": "OpenRouter API Key",
|
||||
"openRouterApiKeyPlaceholder": "Enter your OpenRouter API key",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue