mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
feat: implement local vector store and embedding capabilities
- Add LibSQLVectorStore implementation using @mastra/libsql for local SQLite-based vector storage - Add FastEmbedEmbedder implementation using @mastra/fastembed for local CPU-based embeddings - Support bge-small-en-v1.5 and bge-base-en-v1.5 embedding models - Update configuration system to support "local" vector store type and "fastembed" embedder provider - Add comprehensive test coverage for new implementations - Maintain backward compatibility with existing Qdrant and OpenAI integrations - Enable zero-cost, privacy-focused code indexing without external dependencies Resolves #5682
This commit is contained in:
parent
a163053430
commit
b2e8141c7e
16 changed files with 4335 additions and 51 deletions
|
|
@ -21,7 +21,9 @@ export const CODEBASE_INDEX_DEFAULTS = {
|
|||
export const codebaseIndexConfigSchema = z.object({
|
||||
codebaseIndexEnabled: z.boolean().optional(),
|
||||
codebaseIndexQdrantUrl: z.string().optional(),
|
||||
codebaseIndexEmbedderProvider: z.enum(["openai", "ollama", "openai-compatible", "gemini"]).optional(),
|
||||
codebaseIndexVectorStoreType: z.enum(["qdrant", "local"]).optional(),
|
||||
codebaseIndexLocalVectorStorePath: z.string().optional(),
|
||||
codebaseIndexEmbedderProvider: z.enum(["openai", "ollama", "openai-compatible", "gemini", "fastembed"]).optional(),
|
||||
codebaseIndexEmbedderBaseUrl: z.string().optional(),
|
||||
codebaseIndexEmbedderModelId: z.string().optional(),
|
||||
codebaseIndexEmbedderModelDimension: z.number().optional(),
|
||||
|
|
@ -47,6 +49,7 @@ export const codebaseIndexModelsSchema = z.object({
|
|||
ollama: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
|
||||
"openai-compatible": z.record(z.string(), z.object({ dimension: z.number() })).optional(),
|
||||
gemini: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
|
||||
fastembed: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
|
||||
})
|
||||
|
||||
export type CodebaseIndexModels = z.infer<typeof codebaseIndexModelsSchema>
|
||||
|
|
|
|||
3138
pnpm-lock.yaml
generated
3138
pnpm-lock.yaml
generated
File diff suppressed because it is too large
Load diff
|
|
@ -392,6 +392,8 @@
|
|||
"@aws-sdk/credential-providers": "^3.806.0",
|
||||
"@google/genai": "^1.0.0",
|
||||
"@lmstudio/sdk": "^1.1.1",
|
||||
"@mastra/fastembed": "^0.10.1",
|
||||
"@mastra/libsql": "^0.11.0",
|
||||
"@mistralai/mistralai": "^1.3.6",
|
||||
"@modelcontextprotocol/sdk": "^1.9.0",
|
||||
"@qdrant/js-client-rest": "^1.14.0",
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ export class CodeIndexConfigManager {
|
|||
private geminiOptions?: { apiKey: string }
|
||||
private qdrantUrl?: string = "http://localhost:6333"
|
||||
private qdrantApiKey?: string
|
||||
private vectorStoreType?: "qdrant" | "local"
|
||||
private localVectorStorePath?: string
|
||||
private searchMinScore?: number
|
||||
private searchMaxResults?: number
|
||||
|
||||
|
|
@ -54,6 +56,8 @@ export class CodeIndexConfigManager {
|
|||
const {
|
||||
codebaseIndexEnabled,
|
||||
codebaseIndexQdrantUrl,
|
||||
codebaseIndexVectorStoreType,
|
||||
codebaseIndexLocalVectorStorePath,
|
||||
codebaseIndexEmbedderProvider,
|
||||
codebaseIndexEmbedderBaseUrl,
|
||||
codebaseIndexEmbedderModelId,
|
||||
|
|
@ -72,6 +76,8 @@ export class CodeIndexConfigManager {
|
|||
this.codebaseIndexEnabled = codebaseIndexEnabled ?? true
|
||||
this.qdrantUrl = codebaseIndexQdrantUrl
|
||||
this.qdrantApiKey = qdrantApiKey ?? ""
|
||||
this.vectorStoreType = codebaseIndexVectorStoreType as "qdrant" | "local" | undefined
|
||||
this.localVectorStorePath = codebaseIndexLocalVectorStorePath
|
||||
this.searchMinScore = codebaseIndexSearchMinScore
|
||||
this.searchMaxResults = codebaseIndexSearchMaxResults
|
||||
|
||||
|
|
@ -93,13 +99,15 @@ export class CodeIndexConfigManager {
|
|||
|
||||
this.openAiOptions = { openAiNativeApiKey: openAiKey }
|
||||
|
||||
// Set embedder provider with support for openai-compatible
|
||||
// Set embedder provider with support for openai-compatible and fastembed
|
||||
if (codebaseIndexEmbedderProvider === "ollama") {
|
||||
this.embedderProvider = "ollama"
|
||||
} else if (codebaseIndexEmbedderProvider === "openai-compatible") {
|
||||
this.embedderProvider = "openai-compatible"
|
||||
} else if (codebaseIndexEmbedderProvider === "gemini") {
|
||||
this.embedderProvider = "gemini"
|
||||
} else if (codebaseIndexEmbedderProvider === "fastembed") {
|
||||
this.embedderProvider = "fastembed"
|
||||
} else {
|
||||
this.embedderProvider = "openai"
|
||||
}
|
||||
|
|
@ -188,26 +196,28 @@ export class CodeIndexConfigManager {
|
|||
* Checks if the service is properly configured based on the embedder type.
|
||||
*/
|
||||
public isConfigured(): boolean {
|
||||
// Check if we have a vector store configured (either Qdrant or local)
|
||||
const hasVectorStore = this.qdrantUrl || this.vectorStoreType === "local"
|
||||
|
||||
if (this.embedderProvider === "openai") {
|
||||
const openAiKey = this.openAiOptions?.openAiNativeApiKey
|
||||
const qdrantUrl = this.qdrantUrl
|
||||
return !!(openAiKey && qdrantUrl)
|
||||
return !!(openAiKey && hasVectorStore)
|
||||
} else if (this.embedderProvider === "ollama") {
|
||||
// Ollama model ID has a default, so only base URL is strictly required for config
|
||||
const ollamaBaseUrl = this.ollamaOptions?.ollamaBaseUrl
|
||||
const qdrantUrl = this.qdrantUrl
|
||||
return !!(ollamaBaseUrl && qdrantUrl)
|
||||
return !!(ollamaBaseUrl && hasVectorStore)
|
||||
} else if (this.embedderProvider === "openai-compatible") {
|
||||
const baseUrl = this.openAiCompatibleOptions?.baseUrl
|
||||
const apiKey = this.openAiCompatibleOptions?.apiKey
|
||||
const qdrantUrl = this.qdrantUrl
|
||||
const isConfigured = !!(baseUrl && apiKey && qdrantUrl)
|
||||
const isConfigured = !!(baseUrl && apiKey && hasVectorStore)
|
||||
return isConfigured
|
||||
} else if (this.embedderProvider === "gemini") {
|
||||
const apiKey = this.geminiOptions?.apiKey
|
||||
const qdrantUrl = this.qdrantUrl
|
||||
const isConfigured = !!(apiKey && qdrantUrl)
|
||||
const isConfigured = !!(apiKey && hasVectorStore)
|
||||
return isConfigured
|
||||
} else if (this.embedderProvider === "fastembed") {
|
||||
// FastEmbed is local and doesn't require API keys, just a vector store
|
||||
return !!hasVectorStore
|
||||
}
|
||||
return false // Should not happen if embedderProvider is always set correctly
|
||||
}
|
||||
|
|
@ -353,6 +363,8 @@ export class CodeIndexConfigManager {
|
|||
geminiOptions: this.geminiOptions,
|
||||
qdrantUrl: this.qdrantUrl,
|
||||
qdrantApiKey: this.qdrantApiKey,
|
||||
vectorStoreType: this.vectorStoreType,
|
||||
localVectorStorePath: this.localVectorStorePath,
|
||||
searchMinScore: this.currentSearchMinScore,
|
||||
searchMaxResults: this.currentSearchMaxResults,
|
||||
}
|
||||
|
|
|
|||
243
src/services/code-index/embedders/__tests__/fastembed.spec.ts
Normal file
243
src/services/code-index/embedders/__tests__/fastembed.spec.ts
Normal file
|
|
@ -0,0 +1,243 @@
|
|||
// npx vitest services/code-index/embedders/__tests__/fastembed.spec.ts
|
||||
|
||||
import { describe, it, expect, beforeEach, vi } from "vitest"
|
||||
import { FastEmbedEmbedder } from "../fastembed"
|
||||
|
||||
// Mock TelemetryService
|
||||
vi.mock("@roo-code/telemetry", () => ({
|
||||
TelemetryService: {
|
||||
instance: {
|
||||
captureEvent: vi.fn(),
|
||||
},
|
||||
},
|
||||
}))
|
||||
|
||||
// Mock i18n
|
||||
vi.mock("../../../i18n", () => ({
|
||||
t: vi.fn((key: string, params?: any) => {
|
||||
if (key === "embeddings:fastembed.modelNotSupported") {
|
||||
return `Model "${params?.model}" not supported. Available models: ${params?.availableModels}`
|
||||
}
|
||||
if (key === "embeddings:fastembed.embeddingFailed") {
|
||||
return `Failed to create embeddings with FastEmbed: ${params?.message}`
|
||||
}
|
||||
if (key === "embeddings:fastembed.noValidTexts") {
|
||||
return "No valid texts to embed"
|
||||
}
|
||||
if (key === "embeddings:fastembed.invalidResponseFormat") {
|
||||
return "Invalid response format from FastEmbed"
|
||||
}
|
||||
if (key === "embeddings:fastembed.invalidEmbeddingFormat") {
|
||||
return "Invalid embedding format from FastEmbed"
|
||||
}
|
||||
return key
|
||||
}),
|
||||
}))
|
||||
|
||||
// Mock getModelQueryPrefix
|
||||
vi.mock("../../../shared/embeddingModels", () => ({
|
||||
getModelQueryPrefix: vi.fn(() => null),
|
||||
}))
|
||||
|
||||
// Mock @mastra/fastembed
|
||||
vi.mock("@mastra/fastembed", () => ({
|
||||
fastembed: {
|
||||
small: {
|
||||
doEmbed: vi.fn(),
|
||||
maxEmbeddingsPerCall: 256,
|
||||
},
|
||||
base: {
|
||||
doEmbed: vi.fn(),
|
||||
maxEmbeddingsPerCall: 256,
|
||||
},
|
||||
},
|
||||
}))
|
||||
|
||||
describe("FastEmbedEmbedder", () => {
|
||||
let embedder: FastEmbedEmbedder
|
||||
let mockSmallDoEmbed: any
|
||||
let mockBaseDoEmbed: any
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
|
||||
// Get references to the mocked functions
|
||||
const { fastembed } = require("@mastra/fastembed")
|
||||
mockSmallDoEmbed = fastembed.small.doEmbed
|
||||
mockBaseDoEmbed = fastembed.base.doEmbed
|
||||
})
|
||||
|
||||
describe("constructor", () => {
|
||||
it("should initialize with default model (bge-small-en-v1.5)", () => {
|
||||
embedder = new FastEmbedEmbedder({})
|
||||
expect(embedder.embedderInfo.name).toBe("fastembed")
|
||||
})
|
||||
|
||||
it("should initialize with specified model", () => {
|
||||
embedder = new FastEmbedEmbedder({ fastEmbedModel: "bge-base-en-v1.5" })
|
||||
expect(embedder.embedderInfo.name).toBe("fastembed")
|
||||
})
|
||||
|
||||
it("should use fallback model for unsupported model", () => {
|
||||
const consoleSpy = vi.spyOn(console, "warn").mockImplementation(() => {})
|
||||
|
||||
embedder = new FastEmbedEmbedder({ fastEmbedModel: "unsupported-model" })
|
||||
|
||||
expect(consoleSpy).toHaveBeenCalledWith(expect.stringContaining('Model "unsupported-model" not available'))
|
||||
|
||||
consoleSpy.mockRestore()
|
||||
})
|
||||
})
|
||||
|
||||
describe("createEmbeddings", () => {
|
||||
beforeEach(() => {
|
||||
embedder = new FastEmbedEmbedder({})
|
||||
})
|
||||
|
||||
it("should create embeddings for single text using small model", async () => {
|
||||
const mockEmbeddings = [[0.1, 0.2, 0.3, 0.4]]
|
||||
mockSmallDoEmbed.mockResolvedValue(mockEmbeddings)
|
||||
|
||||
const result = await embedder.createEmbeddings(["test text"])
|
||||
|
||||
expect(mockSmallDoEmbed).toHaveBeenCalledWith({ values: ["test text"] })
|
||||
expect(result).toEqual({
|
||||
embeddings: mockEmbeddings,
|
||||
})
|
||||
})
|
||||
|
||||
it("should create embeddings for multiple texts using small model", async () => {
|
||||
const mockEmbeddings = [
|
||||
[0.1, 0.2, 0.3, 0.4],
|
||||
[0.5, 0.6, 0.7, 0.8],
|
||||
]
|
||||
mockSmallDoEmbed.mockResolvedValue(mockEmbeddings)
|
||||
|
||||
const result = await embedder.createEmbeddings(["text 1", "text 2"])
|
||||
|
||||
expect(mockSmallDoEmbed).toHaveBeenCalledWith({ values: ["text 1", "text 2"] })
|
||||
expect(result).toEqual({
|
||||
embeddings: mockEmbeddings,
|
||||
})
|
||||
})
|
||||
|
||||
it("should create embeddings using base model when specified", async () => {
|
||||
embedder = new FastEmbedEmbedder({ fastEmbedModel: "bge-base-en-v1.5" })
|
||||
const mockEmbeddings = [[0.1, 0.2, 0.3, 0.4]]
|
||||
mockBaseDoEmbed.mockResolvedValue(mockEmbeddings)
|
||||
|
||||
const result = await embedder.createEmbeddings(["test text"])
|
||||
|
||||
expect(mockBaseDoEmbed).toHaveBeenCalledWith({ values: ["test text"] })
|
||||
expect(result).toEqual({
|
||||
embeddings: mockEmbeddings,
|
||||
})
|
||||
})
|
||||
|
||||
it("should handle empty input", async () => {
|
||||
const result = await embedder.createEmbeddings([])
|
||||
|
||||
expect(mockSmallDoEmbed).not.toHaveBeenCalled()
|
||||
expect(result).toEqual({
|
||||
embeddings: [],
|
||||
})
|
||||
})
|
||||
|
||||
it("should handle FastEmbed API errors", async () => {
|
||||
const error = new Error("FastEmbed API error")
|
||||
mockSmallDoEmbed.mockRejectedValue(error)
|
||||
|
||||
await expect(embedder.createEmbeddings(["test text"])).rejects.toThrow(
|
||||
"Failed to create embeddings with FastEmbed: FastEmbed API error",
|
||||
)
|
||||
})
|
||||
|
||||
it("should process large batches correctly", async () => {
|
||||
const texts = Array.from({ length: 150 }, (_, i) => `text ${i}`)
|
||||
const mockEmbeddings = texts.map((_, i) => [i * 0.1, i * 0.2, i * 0.3, i * 0.4])
|
||||
mockSmallDoEmbed.mockResolvedValue(mockEmbeddings)
|
||||
|
||||
const result = await embedder.createEmbeddings(texts)
|
||||
|
||||
expect(mockSmallDoEmbed).toHaveBeenCalledWith({ values: texts })
|
||||
expect(result.embeddings).toHaveLength(150)
|
||||
})
|
||||
})
|
||||
|
||||
describe("validateConfiguration", () => {
|
||||
beforeEach(() => {
|
||||
embedder = new FastEmbedEmbedder({})
|
||||
})
|
||||
|
||||
it("should validate successfully with small model", async () => {
|
||||
const mockEmbeddings = [[0.1, 0.2, 0.3, 0.4]]
|
||||
mockSmallDoEmbed.mockResolvedValue(mockEmbeddings)
|
||||
|
||||
const result = await embedder.validateConfiguration()
|
||||
|
||||
expect(mockSmallDoEmbed).toHaveBeenCalledWith({ values: ["test"] })
|
||||
expect(result).toEqual({ valid: true })
|
||||
})
|
||||
|
||||
it("should validate successfully with base model", async () => {
|
||||
embedder = new FastEmbedEmbedder({ fastEmbedModel: "bge-base-en-v1.5" })
|
||||
const mockEmbeddings = [[0.1, 0.2, 0.3, 0.4]]
|
||||
mockBaseDoEmbed.mockResolvedValue(mockEmbeddings)
|
||||
|
||||
const result = await embedder.validateConfiguration()
|
||||
|
||||
expect(mockBaseDoEmbed).toHaveBeenCalledWith({ values: ["test"] })
|
||||
expect(result).toEqual({ valid: true })
|
||||
})
|
||||
|
||||
it("should return invalid when FastEmbed fails", async () => {
|
||||
const error = new Error("FastEmbed validation error")
|
||||
mockSmallDoEmbed.mockRejectedValue(error)
|
||||
|
||||
const result = await embedder.validateConfiguration()
|
||||
|
||||
expect(result).toEqual({
|
||||
valid: false,
|
||||
error: "FastEmbed validation failed: FastEmbed validation error",
|
||||
})
|
||||
})
|
||||
|
||||
it("should handle unexpected validation errors", async () => {
|
||||
mockSmallDoEmbed.mockRejectedValue("Unexpected error")
|
||||
|
||||
const result = await embedder.validateConfiguration()
|
||||
|
||||
expect(result).toEqual({
|
||||
valid: false,
|
||||
error: "FastEmbed validation failed: Unexpected error",
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe("embedderInfo", () => {
|
||||
it("should return correct embedder info", () => {
|
||||
embedder = new FastEmbedEmbedder({})
|
||||
expect(embedder.embedderInfo).toEqual({
|
||||
name: "fastembed",
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe("model selection", () => {
|
||||
it("should use small model by default", () => {
|
||||
embedder = new FastEmbedEmbedder({})
|
||||
// We can't directly test the private property, but we can test the behavior
|
||||
expect(() => embedder).not.toThrow()
|
||||
})
|
||||
|
||||
it("should use base model when specified", () => {
|
||||
embedder = new FastEmbedEmbedder({ fastEmbedModel: "bge-base-en-v1.5" })
|
||||
expect(() => embedder).not.toThrow()
|
||||
})
|
||||
|
||||
it("should use small model when explicitly specified", () => {
|
||||
embedder = new FastEmbedEmbedder({ fastEmbedModel: "bge-small-en-v1.5" })
|
||||
expect(() => embedder).not.toThrow()
|
||||
})
|
||||
})
|
||||
})
|
||||
230
src/services/code-index/embedders/fastembed.ts
Normal file
230
src/services/code-index/embedders/fastembed.ts
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
import { fastembed } from "@mastra/fastembed"
|
||||
import { ApiHandlerOptions } from "../../../shared/api"
|
||||
import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces"
|
||||
import { getModelQueryPrefix } from "../../../shared/embeddingModels"
|
||||
import { MAX_ITEM_TOKENS } from "../constants"
|
||||
import { t } from "../../../i18n"
|
||||
import { withValidationErrorHandling, sanitizeErrorMessage } from "../shared/validation-helpers"
|
||||
import { TelemetryService } from "@roo-code/telemetry"
|
||||
import { TelemetryEventName } from "@roo-code/types"
|
||||
|
||||
/**
|
||||
* FastEmbed implementation of the embedder interface for local CPU-based embeddings
|
||||
*/
|
||||
export class FastEmbedEmbedder implements IEmbedder {
|
||||
private readonly defaultModel: string
|
||||
private readonly availableModels: Record<string, any>
|
||||
|
||||
constructor(options: ApiHandlerOptions & { fastEmbedModel?: string }) {
|
||||
// Available FastEmbed models
|
||||
this.availableModels = {
|
||||
"bge-small-en-v1.5": fastembed.small,
|
||||
"bge-base-en-v1.5": fastembed.base,
|
||||
}
|
||||
|
||||
// Set default model
|
||||
this.defaultModel = options.fastEmbedModel || "bge-small-en-v1.5"
|
||||
|
||||
// Validate that the selected model is available
|
||||
if (!this.availableModels[this.defaultModel]) {
|
||||
console.warn(
|
||||
`[FastEmbedEmbedder] Model "${this.defaultModel}" not available. Using "bge-small-en-v1.5" as fallback.`,
|
||||
)
|
||||
this.defaultModel = "bge-small-en-v1.5"
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates embeddings for the given texts using FastEmbed
|
||||
* @param texts Array of text strings to embed
|
||||
* @param model Optional model identifier (currently supports bge-small-en-v1.5 and bge-base-en-v1.5)
|
||||
* @returns Promise resolving to embedding response
|
||||
*/
|
||||
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
|
||||
const modelToUse = model || this.defaultModel
|
||||
|
||||
// Get the FastEmbed model instance
|
||||
const fastEmbedModel = this.availableModels[modelToUse]
|
||||
if (!fastEmbedModel) {
|
||||
throw new Error(
|
||||
t("embeddings:fastembed.modelNotSupported", {
|
||||
model: modelToUse,
|
||||
availableModels: Object.keys(this.availableModels).join(", "),
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
// Apply model-specific query prefix if required
|
||||
const queryPrefix = getModelQueryPrefix("fastembed", modelToUse)
|
||||
const processedTexts = queryPrefix
|
||||
? texts.map((text, index) => {
|
||||
// Prevent double-prefixing
|
||||
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 original text if adding prefix would exceed limit
|
||||
return text
|
||||
}
|
||||
return prefixedText
|
||||
})
|
||||
: texts
|
||||
|
||||
try {
|
||||
// Filter out texts that are too long
|
||||
const validTexts = processedTexts.filter((text, index) => {
|
||||
const estimatedTokens = Math.ceil(text.length / 4)
|
||||
if (estimatedTokens > MAX_ITEM_TOKENS) {
|
||||
console.warn(
|
||||
t("embeddings:textExceedsTokenLimit", {
|
||||
index,
|
||||
itemTokens: estimatedTokens,
|
||||
maxTokens: MAX_ITEM_TOKENS,
|
||||
}),
|
||||
)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
if (validTexts.length === 0) {
|
||||
throw new Error(t("embeddings:fastembed.noValidTexts"))
|
||||
}
|
||||
|
||||
// Process texts in batches according to model's maxEmbeddingsPerCall
|
||||
const maxBatchSize = fastEmbedModel.maxEmbeddingsPerCall || 256
|
||||
const allEmbeddings: number[][] = []
|
||||
|
||||
for (let i = 0; i < validTexts.length; i += maxBatchSize) {
|
||||
const batch = validTexts.slice(i, i + maxBatchSize)
|
||||
|
||||
// Call FastEmbed's doEmbed method
|
||||
const batchResult = await fastEmbedModel.doEmbed({
|
||||
values: batch,
|
||||
})
|
||||
|
||||
// FastEmbed returns embeddings in the format we expect
|
||||
if (Array.isArray(batchResult) && batchResult.length > 0) {
|
||||
allEmbeddings.push(...batchResult)
|
||||
} else {
|
||||
throw new Error(t("embeddings:fastembed.invalidResponseFormat"))
|
||||
}
|
||||
}
|
||||
|
||||
// FastEmbed doesn't provide usage statistics, so we estimate
|
||||
const estimatedTokens = validTexts.reduce((total, text) => total + Math.ceil(text.length / 4), 0)
|
||||
|
||||
return {
|
||||
embeddings: allEmbeddings,
|
||||
usage: {
|
||||
promptTokens: estimatedTokens,
|
||||
totalTokens: estimatedTokens,
|
||||
},
|
||||
}
|
||||
} catch (error: any) {
|
||||
// Capture telemetry before reformatting the error
|
||||
TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
|
||||
error: sanitizeErrorMessage(error instanceof Error ? error.message : String(error)),
|
||||
stack: error instanceof Error ? sanitizeErrorMessage(error.stack || "") : undefined,
|
||||
location: "FastEmbedEmbedder:createEmbeddings",
|
||||
})
|
||||
|
||||
// Log the original error for debugging purposes
|
||||
console.error("FastEmbed embedding failed:", error)
|
||||
|
||||
// Re-throw a more specific error for the caller
|
||||
throw new Error(t("embeddings:fastembed.embeddingFailed", { message: error.message }))
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates the FastEmbed embedder configuration by testing a simple embedding
|
||||
* @returns Promise resolving to validation result with success status and optional error message
|
||||
*/
|
||||
async validateConfiguration(): Promise<{ valid: boolean; error?: string }> {
|
||||
return withValidationErrorHandling(
|
||||
async () => {
|
||||
try {
|
||||
// Get the default model
|
||||
const fastEmbedModel = this.availableModels[this.defaultModel]
|
||||
if (!fastEmbedModel) {
|
||||
return {
|
||||
valid: false,
|
||||
error: t("embeddings:fastembed.modelNotSupported", {
|
||||
model: this.defaultModel,
|
||||
availableModels: Object.keys(this.availableModels).join(", "),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// Test with a simple embedding request
|
||||
const testResult = await fastEmbedModel.doEmbed({
|
||||
values: ["test"],
|
||||
})
|
||||
|
||||
// Check if we got a valid response
|
||||
if (!Array.isArray(testResult) || testResult.length === 0) {
|
||||
return {
|
||||
valid: false,
|
||||
error: t("embeddings:fastembed.invalidResponseFormat"),
|
||||
}
|
||||
}
|
||||
|
||||
// Check if the embedding has the expected structure
|
||||
const firstEmbedding = testResult[0]
|
||||
if (!Array.isArray(firstEmbedding) || firstEmbedding.length === 0) {
|
||||
return {
|
||||
valid: false,
|
||||
error: t("embeddings:fastembed.invalidEmbeddingFormat"),
|
||||
}
|
||||
}
|
||||
|
||||
return { valid: true }
|
||||
} catch (error) {
|
||||
// Capture telemetry for validation errors
|
||||
TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
|
||||
error: sanitizeErrorMessage(error instanceof Error ? error.message : String(error)),
|
||||
stack: error instanceof Error ? sanitizeErrorMessage(error.stack || "") : undefined,
|
||||
location: "FastEmbedEmbedder:validateConfiguration",
|
||||
})
|
||||
throw error
|
||||
}
|
||||
},
|
||||
"fastembed",
|
||||
{
|
||||
beforeStandardHandling: (error: any) => {
|
||||
// Handle FastEmbed-specific errors
|
||||
if (
|
||||
error?.message?.includes("model not found") ||
|
||||
error?.message?.includes("Model not supported")
|
||||
) {
|
||||
return {
|
||||
valid: false,
|
||||
error: t("embeddings:fastembed.modelNotSupported", {
|
||||
model: this.defaultModel,
|
||||
availableModels: Object.keys(this.availableModels).join(", "),
|
||||
}),
|
||||
}
|
||||
}
|
||||
// Let standard handling take over
|
||||
return undefined
|
||||
},
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
get embedderInfo(): EmbedderInfo {
|
||||
return {
|
||||
name: "fastembed",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -15,6 +15,8 @@ export interface CodeIndexConfig {
|
|||
geminiOptions?: { apiKey: string }
|
||||
qdrantUrl?: string
|
||||
qdrantApiKey?: string
|
||||
vectorStoreType?: "qdrant" | "local"
|
||||
localVectorStorePath?: string
|
||||
searchMinScore?: number
|
||||
searchMaxResults?: number
|
||||
}
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ export interface EmbeddingResponse {
|
|||
}
|
||||
}
|
||||
|
||||
export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini"
|
||||
export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "fastembed"
|
||||
|
||||
export interface EmbedderInfo {
|
||||
name: AvailableEmbedders
|
||||
|
|
|
|||
|
|
@ -70,7 +70,7 @@ export interface ICodeIndexManager {
|
|||
}
|
||||
|
||||
export type IndexingState = "Standby" | "Indexing" | "Indexed" | "Error"
|
||||
export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini"
|
||||
export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "fastembed"
|
||||
|
||||
export interface IndexProgressUpdate {
|
||||
systemStatus: IndexingState
|
||||
|
|
|
|||
|
|
@ -3,8 +3,10 @@ import { OpenAiEmbedder } from "./embedders/openai"
|
|||
import { CodeIndexOllamaEmbedder } from "./embedders/ollama"
|
||||
import { OpenAICompatibleEmbedder } from "./embedders/openai-compatible"
|
||||
import { GeminiEmbedder } from "./embedders/gemini"
|
||||
import { FastEmbedEmbedder } from "./embedders/fastembed"
|
||||
import { EmbedderProvider, getDefaultModelId, getModelDimension } from "../../shared/embeddingModels"
|
||||
import { QdrantVectorStore } from "./vector-store/qdrant-client"
|
||||
import { LibSQLVectorStore } from "./vector-store/libsql-vector-store"
|
||||
import { codeParser, DirectoryScanner, FileWatcher } from "./processors"
|
||||
import { ICodeParser, IEmbedder, IFileWatcher, IVectorStore } from "./interfaces"
|
||||
import { CodeIndexConfigManager } from "./config-manager"
|
||||
|
|
@ -64,6 +66,10 @@ export class CodeIndexServiceFactory {
|
|||
throw new Error(t("embeddings:serviceFactory.geminiConfigMissing"))
|
||||
}
|
||||
return new GeminiEmbedder(config.geminiOptions.apiKey)
|
||||
} else if (provider === "fastembed") {
|
||||
return new FastEmbedEmbedder({
|
||||
fastEmbedModel: config.modelId,
|
||||
})
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
|
|
@ -129,12 +135,21 @@ export class CodeIndexServiceFactory {
|
|||
}
|
||||
}
|
||||
|
||||
if (!config.qdrantUrl) {
|
||||
throw new Error(t("embeddings:serviceFactory.qdrantUrlMissing"))
|
||||
}
|
||||
// Check if using local vector store
|
||||
const vectorStoreType = config.vectorStoreType || "qdrant" // Default to qdrant for backward compatibility
|
||||
|
||||
// Assuming constructor is updated: new QdrantVectorStore(workspacePath, url, vectorSize, apiKey?)
|
||||
return new QdrantVectorStore(this.workspacePath, config.qdrantUrl, vectorSize, config.qdrantApiKey)
|
||||
if (vectorStoreType === "local") {
|
||||
const localStorePath = config.localVectorStorePath || `${this.workspacePath}/.roo/vector-store`
|
||||
return new LibSQLVectorStore(localStorePath, "codebase_index", vectorSize)
|
||||
} else {
|
||||
// Default to Qdrant
|
||||
if (!config.qdrantUrl) {
|
||||
throw new Error(t("embeddings:serviceFactory.qdrantUrlMissing"))
|
||||
}
|
||||
|
||||
// Assuming constructor is updated: new QdrantVectorStore(workspacePath, url, vectorSize, apiKey?)
|
||||
return new QdrantVectorStore(this.workspacePath, config.qdrantUrl, vectorSize, config.qdrantApiKey)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -0,0 +1,436 @@
|
|||
// npx vitest services/code-index/vector-store/__tests__/libsql-vector-store.spec.ts
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"
|
||||
import { LibSQLVectorStore } from "../libsql-vector-store"
|
||||
import { PointStruct } from "../../interfaces/vector-store"
|
||||
import * as fs from "fs"
|
||||
import * as path from "path"
|
||||
|
||||
// Mock @mastra/libsql
|
||||
vi.mock("@mastra/libsql", () => ({
|
||||
LibSQLVector: vi.fn().mockImplementation(() => ({
|
||||
createIndex: vi.fn(),
|
||||
upsert: vi.fn(),
|
||||
query: vi.fn(),
|
||||
deleteVector: vi.fn(),
|
||||
truncateIndex: vi.fn(),
|
||||
})),
|
||||
}))
|
||||
|
||||
// Mock fs for cleanup
|
||||
vi.mock("fs", () => ({
|
||||
existsSync: vi.fn(),
|
||||
rmSync: vi.fn(),
|
||||
}))
|
||||
|
||||
describe("LibSQLVectorStore", () => {
|
||||
let vectorStore: LibSQLVectorStore
|
||||
let mockLibSQLVector: any
|
||||
const testDbPath = "/tmp/test-vector-store.db"
|
||||
const testIndexName = "test_index"
|
||||
const testDimension = 384
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
|
||||
// Get reference to the mocked LibSQLVector constructor
|
||||
const { LibSQLVector } = require("@mastra/libsql")
|
||||
mockLibSQLVector = {
|
||||
createIndex: vi.fn(),
|
||||
upsert: vi.fn(),
|
||||
query: vi.fn(),
|
||||
deleteVector: vi.fn(),
|
||||
truncateIndex: vi.fn(),
|
||||
}
|
||||
LibSQLVector.mockReturnValue(mockLibSQLVector)
|
||||
|
||||
vectorStore = new LibSQLVectorStore(testDbPath, testIndexName, testDimension)
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
describe("constructor", () => {
|
||||
it("should initialize with correct parameters", () => {
|
||||
const { LibSQLVector } = require("@mastra/libsql")
|
||||
expect(LibSQLVector).toHaveBeenCalledWith(testDbPath)
|
||||
})
|
||||
|
||||
it("should store configuration parameters", () => {
|
||||
expect(vectorStore).toBeDefined()
|
||||
})
|
||||
})
|
||||
|
||||
describe("initialize", () => {
|
||||
it("should create index with correct parameters", async () => {
|
||||
mockLibSQLVector.createIndex.mockResolvedValue(undefined)
|
||||
|
||||
await vectorStore.initialize()
|
||||
|
||||
expect(mockLibSQLVector.createIndex).toHaveBeenCalledWith({
|
||||
indexName: testIndexName,
|
||||
dimension: testDimension,
|
||||
})
|
||||
})
|
||||
|
||||
it("should handle initialization errors", async () => {
|
||||
const error = new Error("Failed to create index")
|
||||
mockLibSQLVector.createIndex.mockRejectedValue(error)
|
||||
|
||||
await expect(vectorStore.initialize()).rejects.toThrow(
|
||||
"Failed to initialize LibSQL vector store: Failed to create index",
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("upsertPoints", () => {
|
||||
const testPoints: PointStruct[] = [
|
||||
{
|
||||
id: "test-1",
|
||||
vector: [0.1, 0.2, 0.3, 0.4],
|
||||
payload: {
|
||||
filePath: "/test/file1.ts",
|
||||
content: "test content 1",
|
||||
startLine: 1,
|
||||
endLine: 10,
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "test-2",
|
||||
vector: [0.5, 0.6, 0.7, 0.8],
|
||||
payload: {
|
||||
filePath: "/test/file2.ts",
|
||||
content: "test content 2",
|
||||
startLine: 11,
|
||||
endLine: 20,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
beforeEach(async () => {
|
||||
mockLibSQLVector.createIndex.mockResolvedValue(undefined)
|
||||
await vectorStore.initialize()
|
||||
})
|
||||
|
||||
it("should upsert points correctly", async () => {
|
||||
mockLibSQLVector.upsert.mockResolvedValue(undefined)
|
||||
|
||||
await vectorStore.upsertPoints(testPoints)
|
||||
|
||||
expect(mockLibSQLVector.upsert).toHaveBeenCalledWith({
|
||||
indexName: testIndexName,
|
||||
vectors: [
|
||||
[0.1, 0.2, 0.3, 0.4],
|
||||
[0.5, 0.6, 0.7, 0.8],
|
||||
],
|
||||
ids: ["test-1", "test-2"],
|
||||
metadata: [
|
||||
{
|
||||
filePath: "/test/file1.ts",
|
||||
content: "test content 1",
|
||||
startLine: 1,
|
||||
endLine: 10,
|
||||
},
|
||||
{
|
||||
filePath: "/test/file2.ts",
|
||||
content: "test content 2",
|
||||
startLine: 11,
|
||||
endLine: 20,
|
||||
},
|
||||
],
|
||||
})
|
||||
})
|
||||
|
||||
it("should handle empty points array", async () => {
|
||||
await vectorStore.upsertPoints([])
|
||||
|
||||
expect(mockLibSQLVector.upsert).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should handle upsert errors", async () => {
|
||||
const error = new Error("Upsert failed")
|
||||
mockLibSQLVector.upsert.mockRejectedValue(error)
|
||||
|
||||
await expect(vectorStore.upsertPoints(testPoints)).rejects.toThrow(
|
||||
"Failed to upsert points to LibSQL vector store: Upsert failed",
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle points with missing payload fields", async () => {
|
||||
const pointsWithMissingFields: PointStruct[] = [
|
||||
{
|
||||
id: "test-1",
|
||||
vector: [0.1, 0.2, 0.3, 0.4],
|
||||
payload: {
|
||||
filePath: "/test/file1.ts",
|
||||
content: "test content 1",
|
||||
// Missing startLine and endLine
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
mockLibSQLVector.upsert.mockResolvedValue(undefined)
|
||||
|
||||
await vectorStore.upsertPoints(pointsWithMissingFields)
|
||||
|
||||
expect(mockLibSQLVector.upsert).toHaveBeenCalledWith({
|
||||
indexName: testIndexName,
|
||||
vectors: [[0.1, 0.2, 0.3, 0.4]],
|
||||
ids: ["test-1"],
|
||||
metadata: [
|
||||
{
|
||||
filePath: "/test/file1.ts",
|
||||
content: "test content 1",
|
||||
},
|
||||
],
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe("search", () => {
|
||||
const testQueryVector = [0.1, 0.2, 0.3, 0.4]
|
||||
|
||||
beforeEach(async () => {
|
||||
mockLibSQLVector.createIndex.mockResolvedValue(undefined)
|
||||
await vectorStore.initialize()
|
||||
})
|
||||
|
||||
it("should search with correct parameters", async () => {
|
||||
const mockResults = [
|
||||
{
|
||||
id: "test-1",
|
||||
score: 0.95,
|
||||
metadata: {
|
||||
filePath: "/test/file1.ts",
|
||||
content: "test content 1",
|
||||
startLine: 1,
|
||||
endLine: 10,
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "test-2",
|
||||
score: 0.85,
|
||||
metadata: {
|
||||
filePath: "/test/file2.ts",
|
||||
content: "test content 2",
|
||||
startLine: 11,
|
||||
endLine: 20,
|
||||
},
|
||||
},
|
||||
]
|
||||
mockLibSQLVector.query.mockResolvedValue(mockResults)
|
||||
|
||||
const results = await vectorStore.search(testQueryVector, undefined, 0.5, 10)
|
||||
|
||||
expect(mockLibSQLVector.query).toHaveBeenCalledWith({
|
||||
indexName: testIndexName,
|
||||
queryVector: testQueryVector,
|
||||
topK: 10,
|
||||
})
|
||||
|
||||
expect(results).toEqual([
|
||||
{
|
||||
id: "test-1",
|
||||
score: 0.95,
|
||||
payload: {
|
||||
filePath: "/test/file1.ts",
|
||||
content: "test content 1",
|
||||
startLine: 1,
|
||||
endLine: 10,
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "test-2",
|
||||
score: 0.85,
|
||||
payload: {
|
||||
filePath: "/test/file2.ts",
|
||||
content: "test content 2",
|
||||
startLine: 11,
|
||||
endLine: 20,
|
||||
},
|
||||
},
|
||||
])
|
||||
})
|
||||
|
||||
it("should filter results by minimum score", async () => {
|
||||
const mockResults = [
|
||||
{
|
||||
id: "test-1",
|
||||
score: 0.95,
|
||||
metadata: {
|
||||
filePath: "/test/file1.ts",
|
||||
codeChunk: "test content 1",
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "test-2",
|
||||
score: 0.3, // Below threshold
|
||||
metadata: {
|
||||
filePath: "/test/file2.ts",
|
||||
content: "test content 2",
|
||||
},
|
||||
},
|
||||
]
|
||||
mockLibSQLVector.query.mockResolvedValue(mockResults)
|
||||
|
||||
const results = await vectorStore.search(testQueryVector, 10, 0.5)
|
||||
|
||||
expect(results).toHaveLength(1)
|
||||
expect(results[0].id).toBe("test-1")
|
||||
})
|
||||
|
||||
it("should handle search errors", async () => {
|
||||
const error = new Error("Search failed")
|
||||
mockLibSQLVector.query.mockRejectedValue(error)
|
||||
|
||||
await expect(vectorStore.search(testQueryVector, 10, 0.5)).rejects.toThrow(
|
||||
"Failed to search LibSQL vector store: Search failed",
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle empty search results", async () => {
|
||||
mockLibSQLVector.query.mockResolvedValue([])
|
||||
|
||||
const results = await vectorStore.search(testQueryVector, 10, 0.5)
|
||||
|
||||
expect(results).toEqual([])
|
||||
})
|
||||
|
||||
it("should use default minimum score when not provided", async () => {
|
||||
const mockResults = [
|
||||
{
|
||||
id: "test-1",
|
||||
score: 0.95,
|
||||
metadata: {
|
||||
filePath: "/test/file1.ts",
|
||||
content: "test content 1",
|
||||
},
|
||||
},
|
||||
]
|
||||
mockLibSQLVector.query.mockResolvedValue(mockResults)
|
||||
|
||||
const results = await vectorStore.search(testQueryVector, 10)
|
||||
|
||||
expect(results).toHaveLength(1)
|
||||
})
|
||||
})
|
||||
|
||||
describe("deletePointsByFilePath", () => {
|
||||
beforeEach(async () => {
|
||||
mockLibSQLVector.createIndex.mockResolvedValue(undefined)
|
||||
await vectorStore.initialize()
|
||||
})
|
||||
|
||||
it("should delete points by file path", async () => {
|
||||
mockLibSQLVector.deleteVector.mockResolvedValue(undefined)
|
||||
|
||||
await vectorStore.deletePointsByFilePath("/test/file1.ts")
|
||||
|
||||
expect(mockLibSQLVector.deleteVector).toHaveBeenCalledWith({
|
||||
indexName: testIndexName,
|
||||
where: "metadata->>'filePath' = '/test/file1.ts'",
|
||||
})
|
||||
})
|
||||
|
||||
it("should handle deletion errors", async () => {
|
||||
const error = new Error("Delete failed")
|
||||
mockLibSQLVector.deleteVector.mockRejectedValue(error)
|
||||
|
||||
await expect(vectorStore.deletePointsByFilePath("/test/file1.ts")).rejects.toThrow(
|
||||
"Failed to delete points by file path from LibSQL vector store: Delete failed",
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle file paths with special characters", async () => {
|
||||
mockLibSQLVector.deleteVector.mockResolvedValue(undefined)
|
||||
|
||||
await vectorStore.deletePointsByFilePath("/test/file with spaces & symbols.ts")
|
||||
|
||||
expect(mockLibSQLVector.deleteVector).toHaveBeenCalledWith({
|
||||
indexName: testIndexName,
|
||||
where: "metadata->>'filePath' = '/test/file with spaces & symbols.ts'",
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe("clearCollection", () => {
|
||||
beforeEach(async () => {
|
||||
mockLibSQLVector.createIndex.mockResolvedValue(undefined)
|
||||
await vectorStore.initialize()
|
||||
})
|
||||
|
||||
it("should clear collection", async () => {
|
||||
mockLibSQLVector.truncateIndex.mockResolvedValue(undefined)
|
||||
|
||||
await vectorStore.clearCollection()
|
||||
|
||||
expect(mockLibSQLVector.truncateIndex).toHaveBeenCalledWith({
|
||||
indexName: testIndexName,
|
||||
})
|
||||
})
|
||||
|
||||
it("should handle clear errors", async () => {
|
||||
const error = new Error("Clear failed")
|
||||
mockLibSQLVector.truncateIndex.mockRejectedValue(error)
|
||||
|
||||
await expect(vectorStore.clearCollection()).rejects.toThrow(
|
||||
"Failed to clear LibSQL vector store collection: Clear failed",
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("deleteCollection", () => {
|
||||
beforeEach(async () => {
|
||||
mockLibSQLVector.createIndex.mockResolvedValue(undefined)
|
||||
await vectorStore.initialize()
|
||||
})
|
||||
|
||||
it("should delete collection and database file", async () => {
|
||||
const mockFs = require("fs")
|
||||
mockFs.existsSync.mockReturnValue(true)
|
||||
mockFs.rmSync.mockReturnValue(undefined)
|
||||
|
||||
await vectorStore.deleteCollection()
|
||||
|
||||
expect(mockFs.existsSync).toHaveBeenCalledWith(testDbPath)
|
||||
expect(mockFs.rmSync).toHaveBeenCalledWith(testDbPath, { force: true })
|
||||
})
|
||||
|
||||
it("should handle case when database file does not exist", async () => {
|
||||
const mockFs = require("fs")
|
||||
mockFs.existsSync.mockReturnValue(false)
|
||||
|
||||
await vectorStore.deleteCollection()
|
||||
|
||||
expect(mockFs.existsSync).toHaveBeenCalledWith(testDbPath)
|
||||
expect(mockFs.rmSync).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should handle deletion errors", async () => {
|
||||
const mockFs = require("fs")
|
||||
mockFs.existsSync.mockReturnValue(true)
|
||||
const error = new Error("File deletion failed")
|
||||
mockFs.rmSync.mockImplementation(() => {
|
||||
throw error
|
||||
})
|
||||
|
||||
await expect(vectorStore.deleteCollection()).rejects.toThrow(
|
||||
"Failed to delete LibSQL vector store collection: File deletion failed",
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("error handling", () => {
|
||||
it("should handle LibSQLVector constructor errors", () => {
|
||||
const { LibSQLVector } = require("@mastra/libsql")
|
||||
LibSQLVector.mockImplementation(() => {
|
||||
throw new Error("Constructor failed")
|
||||
})
|
||||
|
||||
expect(() => {
|
||||
new LibSQLVectorStore(testDbPath, testIndexName, testDimension)
|
||||
}).toThrow("Constructor failed")
|
||||
})
|
||||
})
|
||||
})
|
||||
262
src/services/code-index/vector-store/libsql-vector-store.ts
Normal file
262
src/services/code-index/vector-store/libsql-vector-store.ts
Normal file
|
|
@ -0,0 +1,262 @@
|
|||
import { LibSQLVector } from "@mastra/libsql"
|
||||
import { createHash } from "crypto"
|
||||
import * as path from "path"
|
||||
import * as fs from "fs"
|
||||
import { IVectorStore, PointStruct, VectorStoreSearchResult, Payload } from "../interfaces/vector-store"
|
||||
import { DEFAULT_MAX_SEARCH_RESULTS, DEFAULT_SEARCH_MIN_SCORE } from "../constants"
|
||||
import { t } from "../../../i18n"
|
||||
|
||||
/**
|
||||
* LibSQL implementation of the vector store interface for local file-based vector storage
|
||||
*/
|
||||
export class LibSQLVectorStore implements IVectorStore {
|
||||
private vectorStore: LibSQLVector
|
||||
private readonly collectionName: string
|
||||
private readonly vectorSize: number
|
||||
private readonly databasePath: string
|
||||
|
||||
/**
|
||||
* Creates a new LibSQL vector store
|
||||
* @param workspacePath Path to the workspace
|
||||
* @param databasePath Path to the SQLite database file
|
||||
* @param vectorSize Size of the vectors to store
|
||||
*/
|
||||
constructor(workspacePath: string, databasePath: string, vectorSize: number) {
|
||||
this.vectorSize = vectorSize
|
||||
this.databasePath = databasePath
|
||||
|
||||
// Ensure the directory exists
|
||||
const dbDir = path.dirname(databasePath)
|
||||
if (!fs.existsSync(dbDir)) {
|
||||
fs.mkdirSync(dbDir, { recursive: true })
|
||||
}
|
||||
|
||||
// Initialize LibSQL vector store
|
||||
this.vectorStore = new LibSQLVector({
|
||||
connectionUrl: `file:${databasePath}`,
|
||||
})
|
||||
|
||||
// Generate collection name from workspace path
|
||||
const hash = createHash("sha256").update(workspacePath).digest("hex")
|
||||
this.collectionName = `ws_${hash.substring(0, 16)}`
|
||||
}
|
||||
|
||||
/**
|
||||
* Initializes the vector store by creating necessary indexes
|
||||
* @returns Promise resolving to boolean indicating if a new collection was created
|
||||
*/
|
||||
async initialize(): Promise<boolean> {
|
||||
try {
|
||||
// Check if the index already exists
|
||||
const existingIndexes = await this.vectorStore.listIndexes()
|
||||
const indexExists = existingIndexes.some((index: any) => index.name === this.collectionName)
|
||||
|
||||
if (!indexExists) {
|
||||
// Create the vector index
|
||||
await this.vectorStore.createIndex({
|
||||
indexName: this.collectionName,
|
||||
dimension: this.vectorSize,
|
||||
})
|
||||
return true // New collection created
|
||||
}
|
||||
|
||||
return false // Collection already existed
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : String(error)
|
||||
console.error(`[LibSQLVectorStore] Failed to initialize collection "${this.collectionName}":`, errorMessage)
|
||||
throw new Error(
|
||||
t("embeddings:vectorStore.libsqlInitializationFailed", {
|
||||
databasePath: this.databasePath,
|
||||
errorMessage,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Upserts points into the vector store
|
||||
* @param points Array of points to upsert
|
||||
*/
|
||||
async upsertPoints(points: PointStruct[]): Promise<void> {
|
||||
if (points.length === 0) return
|
||||
|
||||
try {
|
||||
// Extract vectors and metadata separately
|
||||
const vectors = points.map((point) => point.vector)
|
||||
const ids = points.map((point) => point.id)
|
||||
const metadata = points.map((point) => ({
|
||||
filePath: point.payload.filePath || "",
|
||||
codeChunk: point.payload.codeChunk || "",
|
||||
startLine: point.payload.startLine || 0,
|
||||
endLine: point.payload.endLine || 0,
|
||||
pathSegments: point.payload.filePath ? point.payload.filePath.split(path.sep).filter(Boolean) : [],
|
||||
}))
|
||||
|
||||
// Upsert all points to the index
|
||||
await this.vectorStore.upsert({
|
||||
indexName: this.collectionName,
|
||||
vectors: vectors,
|
||||
ids: ids,
|
||||
metadata: metadata,
|
||||
})
|
||||
} catch (error) {
|
||||
console.error("[LibSQLVectorStore] Failed to upsert points:", error)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Searches for similar vectors using the LibSQL vector search
|
||||
* @param queryVector Vector to search for
|
||||
* @param directoryPrefix Optional directory prefix to filter results
|
||||
* @param minScore Optional minimum score threshold
|
||||
* @param maxResults Optional maximum number of results to return
|
||||
* @returns Promise resolving to search results
|
||||
*/
|
||||
async search(
|
||||
queryVector: number[],
|
||||
directoryPrefix?: string,
|
||||
minScore?: number,
|
||||
maxResults?: number,
|
||||
): Promise<VectorStoreSearchResult[]> {
|
||||
try {
|
||||
const actualMinScore = minScore ?? DEFAULT_SEARCH_MIN_SCORE
|
||||
const actualMaxResults = maxResults ?? DEFAULT_MAX_SEARCH_RESULTS
|
||||
|
||||
// Build filter for directory prefix if provided
|
||||
let filter: any = undefined
|
||||
if (directoryPrefix) {
|
||||
const segments = directoryPrefix.split(path.sep).filter(Boolean)
|
||||
// Create a filter that checks if pathSegments starts with the directory segments
|
||||
filter = {
|
||||
pathSegments: {
|
||||
$in: segments,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Perform vector search
|
||||
const searchResults = await this.vectorStore.query({
|
||||
indexName: this.collectionName,
|
||||
queryVector: queryVector,
|
||||
topK: actualMaxResults,
|
||||
filter,
|
||||
includeVector: false,
|
||||
minScore: actualMinScore,
|
||||
})
|
||||
|
||||
// Transform results to our format
|
||||
const results: VectorStoreSearchResult[] = []
|
||||
for (const result of searchResults) {
|
||||
if (result.metadata) {
|
||||
results.push({
|
||||
id: result.id,
|
||||
score: result.score || 0,
|
||||
payload: {
|
||||
filePath: result.metadata.filePath as string,
|
||||
codeChunk: result.metadata.codeChunk as string,
|
||||
startLine: result.metadata.startLine as number,
|
||||
endLine: result.metadata.endLine as number,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by similarity score (descending)
|
||||
results.sort((a, b) => b.score - a.score)
|
||||
return results
|
||||
} catch (error) {
|
||||
console.error("[LibSQLVectorStore] Failed to search points:", error)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes points by file path
|
||||
* @param filePath Path of the file to delete points for
|
||||
*/
|
||||
async deletePointsByFilePath(filePath: string): Promise<void> {
|
||||
return this.deletePointsByMultipleFilePaths([filePath])
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes points by multiple file paths
|
||||
* @param filePaths Array of file paths to delete points for
|
||||
*/
|
||||
async deletePointsByMultipleFilePaths(filePaths: string[]): Promise<void> {
|
||||
if (filePaths.length === 0) return
|
||||
|
||||
try {
|
||||
// LibSQL vector store doesn't have bulk delete by metadata filter
|
||||
// We need to query first to get the IDs, then delete them
|
||||
for (const filePath of filePaths) {
|
||||
// Query to find vectors with this file path
|
||||
const searchResults = await this.vectorStore.query({
|
||||
indexName: this.collectionName,
|
||||
queryVector: new Array(this.vectorSize).fill(0), // Dummy vector for metadata search
|
||||
topK: 10000, // Large number to get all matches
|
||||
filter: {
|
||||
filePath: { $eq: filePath },
|
||||
},
|
||||
includeVector: false,
|
||||
})
|
||||
|
||||
// Delete each found vector by ID
|
||||
for (const result of searchResults) {
|
||||
await this.vectorStore.deleteVector({
|
||||
indexName: this.collectionName,
|
||||
id: result.id,
|
||||
})
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("[LibSQLVectorStore] Failed to delete points by file paths:", error)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Clears all points from the collection
|
||||
*/
|
||||
async clearCollection(): Promise<void> {
|
||||
try {
|
||||
// LibSQL doesn't have a direct clear method, so we truncate the index
|
||||
await this.vectorStore.truncateIndex({
|
||||
indexName: this.collectionName,
|
||||
})
|
||||
} catch (error) {
|
||||
console.error("[LibSQLVectorStore] Failed to clear collection:", error)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes the entire collection (drops the index)
|
||||
*/
|
||||
async deleteCollection(): Promise<void> {
|
||||
try {
|
||||
if (await this.collectionExists()) {
|
||||
await this.vectorStore.deleteIndex({
|
||||
indexName: this.collectionName,
|
||||
})
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(`[LibSQLVectorStore] Failed to delete collection ${this.collectionName}:`, error)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if the collection exists
|
||||
* @returns Promise resolving to boolean indicating if the collection exists
|
||||
*/
|
||||
async collectionExists(): Promise<boolean> {
|
||||
try {
|
||||
const indexes = await this.vectorStore.listIndexes()
|
||||
return indexes.some((index: any) => index.name === this.collectionName)
|
||||
} catch (error) {
|
||||
console.error("[LibSQLVectorStore] Failed to check collection existence:", error)
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -2,7 +2,7 @@
|
|||
* Defines profiles for different embedding models, including their dimensions.
|
||||
*/
|
||||
|
||||
export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" // Add other providers as needed
|
||||
export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "fastembed" // Add other providers as needed
|
||||
|
||||
export interface EmbeddingModelProfile {
|
||||
dimension: number
|
||||
|
|
@ -49,6 +49,10 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
|
|||
gemini: {
|
||||
"text-embedding-004": { dimension: 768 },
|
||||
},
|
||||
fastembed: {
|
||||
"bge-small-en-v1.5": { dimension: 384, scoreThreshold: 0.4 },
|
||||
"bge-base-en-v1.5": { dimension: 768, scoreThreshold: 0.4 },
|
||||
},
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -136,6 +140,9 @@ export function getDefaultModelId(provider: EmbedderProvider): string {
|
|||
case "gemini":
|
||||
return "text-embedding-004"
|
||||
|
||||
case "fastembed":
|
||||
return "bge-small-en-v1.5"
|
||||
|
||||
default:
|
||||
// Fallback for unknown providers
|
||||
console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`)
|
||||
|
|
|
|||
BIN
src/test.db
Normal file
BIN
src/test.db
Normal file
Binary file not shown.
BIN
src/test.db-shm
Normal file
BIN
src/test.db-shm
Normal file
Binary file not shown.
0
src/test.db-wal
Normal file
0
src/test.db-wal
Normal file
Loading…
Add table
Reference in a new issue