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:
Roo Code 2025-07-14 06:04:58 +00:00
parent a163053430
commit b2e8141c7e
16 changed files with 4335 additions and 51 deletions

View file

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

File diff suppressed because it is too large Load diff

View file

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

View file

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

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

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

View file

@ -15,6 +15,8 @@ export interface CodeIndexConfig {
geminiOptions?: { apiKey: string }
qdrantUrl?: string
qdrantApiKey?: string
vectorStoreType?: "qdrant" | "local"
localVectorStorePath?: string
searchMinScore?: number
searchMaxResults?: number
}

View file

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

View file

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

View file

@ -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)
}
}
/**

View file

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

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

View file

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

Binary file not shown.

BIN
src/test.db-shm Normal file

Binary file not shown.

0
src/test.db-wal Normal file
View file