Init
This commit is contained in:
Colby McHenry
@@ -0,0 +1,391 @@
|
||||
/**
|
||||
* Text Embedder
|
||||
*
|
||||
* Generates vector embeddings using the nomic-embed-text model via Transformers.js.
|
||||
* Uses ONNX runtime under the hood for fast local inference.
|
||||
*/
|
||||
|
||||
import { pipeline, env } from '@xenova/transformers';
|
||||
import * as path from 'path';
|
||||
import * as fs from 'fs';
|
||||
|
||||
// Type for the feature extraction pipeline
|
||||
type FeatureExtractionPipeline = Awaited<ReturnType<typeof pipeline<'feature-extraction'>>>;
|
||||
|
||||
/**
|
||||
* Default model for embeddings
|
||||
* nomic-embed-text-v1.5 produces 384-dimensional embeddings
|
||||
*/
|
||||
export const DEFAULT_MODEL = 'nomic-ai/nomic-embed-text-v1.5';
|
||||
export const EMBEDDING_DIMENSION = 768; // nomic-embed-text-v1.5 uses 768 dimensions
|
||||
|
||||
/**
|
||||
* Options for the embedder
|
||||
*/
|
||||
export interface EmbedderOptions {
|
||||
/** Model ID to use (default: nomic-ai/nomic-embed-text-v1.5) */
|
||||
modelId?: string;
|
||||
|
||||
/** Directory to cache the model (default: .codegraph/models) */
|
||||
cacheDir?: string;
|
||||
|
||||
/** Whether to show progress during model download */
|
||||
showProgress?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Text embedding result
|
||||
*/
|
||||
export interface EmbeddingResult {
|
||||
/** The embedding vector */
|
||||
embedding: Float32Array;
|
||||
|
||||
/** Dimension of the embedding */
|
||||
dimension: number;
|
||||
|
||||
/** Model used to generate the embedding */
|
||||
model: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Batch embedding result
|
||||
*/
|
||||
export interface BatchEmbeddingResult {
|
||||
/** Array of embeddings in same order as input */
|
||||
embeddings: Float32Array[];
|
||||
|
||||
/** Dimension of each embedding */
|
||||
dimension: number;
|
||||
|
||||
/** Model used to generate embeddings */
|
||||
model: string;
|
||||
|
||||
/** Processing time in milliseconds */
|
||||
durationMs: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Text Embedder using Transformers.js
|
||||
*
|
||||
* Uses the nomic-embed-text-v1.5 model to generate embeddings for code
|
||||
* and natural language queries.
|
||||
*/
|
||||
export class TextEmbedder {
|
||||
private modelId: string;
|
||||
private cacheDir: string;
|
||||
private pipeline: FeatureExtractionPipeline | null = null;
|
||||
private initialized = false;
|
||||
private showProgress: boolean;
|
||||
|
||||
constructor(options: EmbedderOptions = {}) {
|
||||
this.modelId = options.modelId || DEFAULT_MODEL;
|
||||
this.cacheDir = options.cacheDir || '.codegraph/models';
|
||||
this.showProgress = options.showProgress ?? false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize the embedder by loading the model
|
||||
*
|
||||
* This will download the model on first use if not already cached.
|
||||
*/
|
||||
async initialize(): Promise<void> {
|
||||
if (this.initialized) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Configure transformers.js to use local cache
|
||||
env.cacheDir = this.cacheDir;
|
||||
|
||||
// Ensure cache directory exists
|
||||
if (!fs.existsSync(this.cacheDir)) {
|
||||
fs.mkdirSync(this.cacheDir, { recursive: true });
|
||||
}
|
||||
|
||||
// Disable remote model checking if model is already cached
|
||||
// This speeds up initialization significantly
|
||||
const modelCacheExists = fs.existsSync(
|
||||
path.join(this.cacheDir, this.modelId.replace('/', '--'))
|
||||
);
|
||||
if (modelCacheExists) {
|
||||
env.allowRemoteModels = false;
|
||||
}
|
||||
|
||||
// Load the pipeline
|
||||
this.pipeline = await pipeline('feature-extraction', this.modelId, {
|
||||
progress_callback: this.showProgress
|
||||
? (progress: { status: string; file?: string; progress?: number }) => {
|
||||
if (progress.status === 'progress' && progress.file && progress.progress) {
|
||||
const pct = Math.round(progress.progress);
|
||||
process.stdout.write(`\rDownloading ${progress.file}: ${pct}%`);
|
||||
} else if (progress.status === 'done') {
|
||||
process.stdout.write('\n');
|
||||
}
|
||||
}
|
||||
: undefined,
|
||||
});
|
||||
|
||||
this.initialized = true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the embedder is initialized
|
||||
*/
|
||||
isInitialized(): boolean {
|
||||
return this.initialized;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the model ID being used
|
||||
*/
|
||||
getModelId(): string {
|
||||
return this.modelId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the embedding dimension
|
||||
*/
|
||||
getDimension(): number {
|
||||
return EMBEDDING_DIMENSION;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate embedding for a single text
|
||||
*
|
||||
* @param text - Text to embed
|
||||
* @returns Embedding result
|
||||
*/
|
||||
async embed(text: string): Promise<EmbeddingResult> {
|
||||
if (!this.initialized || !this.pipeline) {
|
||||
throw new Error('Embedder not initialized. Call initialize() first.');
|
||||
}
|
||||
|
||||
// Prepare text for nomic-embed-text (it expects specific prefixes)
|
||||
const preparedText = this.prepareText(text, 'document');
|
||||
|
||||
// Generate embedding
|
||||
const output = await this.pipeline(preparedText, {
|
||||
pooling: 'mean',
|
||||
normalize: true,
|
||||
});
|
||||
|
||||
// Extract the embedding array - handle various data formats
|
||||
const data = output.data as unknown;
|
||||
const embedding = this.toFloat32Array(data);
|
||||
|
||||
return {
|
||||
embedding,
|
||||
dimension: embedding.length,
|
||||
model: this.modelId,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate embedding for a query (uses different prefix)
|
||||
*
|
||||
* @param query - Query text to embed
|
||||
* @returns Embedding result
|
||||
*/
|
||||
async embedQuery(query: string): Promise<EmbeddingResult> {
|
||||
if (!this.initialized || !this.pipeline) {
|
||||
throw new Error('Embedder not initialized. Call initialize() first.');
|
||||
}
|
||||
|
||||
// Prepare text for nomic-embed-text query
|
||||
const preparedText = this.prepareText(query, 'search_query');
|
||||
|
||||
// Generate embedding
|
||||
const output = await this.pipeline(preparedText, {
|
||||
pooling: 'mean',
|
||||
normalize: true,
|
||||
});
|
||||
|
||||
// Extract the embedding array - handle various data formats
|
||||
const data = output.data as unknown;
|
||||
const embedding = this.toFloat32Array(data);
|
||||
|
||||
return {
|
||||
embedding,
|
||||
dimension: embedding.length,
|
||||
model: this.modelId,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate embeddings for multiple texts in a batch
|
||||
*
|
||||
* @param texts - Array of texts to embed
|
||||
* @param type - Type of text (document or search_query)
|
||||
* @returns Batch embedding result
|
||||
*/
|
||||
async embedBatch(
|
||||
texts: string[],
|
||||
type: 'document' | 'search_query' = 'document'
|
||||
): Promise<BatchEmbeddingResult> {
|
||||
if (!this.initialized || !this.pipeline) {
|
||||
throw new Error('Embedder not initialized. Call initialize() first.');
|
||||
}
|
||||
|
||||
if (texts.length === 0) {
|
||||
return {
|
||||
embeddings: [],
|
||||
dimension: EMBEDDING_DIMENSION,
|
||||
model: this.modelId,
|
||||
durationMs: 0,
|
||||
};
|
||||
}
|
||||
|
||||
const startTime = Date.now();
|
||||
|
||||
// Prepare all texts
|
||||
const preparedTexts = texts.map((t) => this.prepareText(t, type));
|
||||
|
||||
// Generate embeddings
|
||||
const outputs = await this.pipeline(preparedTexts, {
|
||||
pooling: 'mean',
|
||||
normalize: true,
|
||||
});
|
||||
|
||||
// Extract embeddings
|
||||
const embeddings: Float32Array[] = [];
|
||||
const dims = outputs.dims as number[];
|
||||
const dimension = dims[1] ?? EMBEDDING_DIMENSION;
|
||||
const data = outputs.data as unknown;
|
||||
const flatData = this.toFloat32Array(data);
|
||||
|
||||
for (let i = 0; i < texts.length; i++) {
|
||||
const start = i * dimension;
|
||||
const end = start + dimension;
|
||||
embeddings.push(flatData.slice(start, end));
|
||||
}
|
||||
|
||||
return {
|
||||
embeddings,
|
||||
dimension,
|
||||
model: this.modelId,
|
||||
durationMs: Date.now() - startTime,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert various array formats to Float32Array
|
||||
*/
|
||||
private toFloat32Array(data: unknown): Float32Array {
|
||||
if (data instanceof Float32Array) {
|
||||
return data;
|
||||
}
|
||||
if (Array.isArray(data)) {
|
||||
return new Float32Array(data);
|
||||
}
|
||||
if (data && typeof data === 'object' && 'length' in data) {
|
||||
// Handle TypedArray-like objects
|
||||
const arr = data as ArrayLike<number>;
|
||||
return new Float32Array(arr.length);
|
||||
}
|
||||
throw new Error('Unsupported data format for embedding');
|
||||
}
|
||||
|
||||
/**
|
||||
* Prepare text for the nomic-embed-text model
|
||||
*
|
||||
* The model expects specific prefixes for different tasks:
|
||||
* - "search_document: " for documents to be searched
|
||||
* - "search_query: " for search queries
|
||||
*/
|
||||
private prepareText(text: string, type: 'document' | 'search_query'): string {
|
||||
// Truncate very long texts (model has a max token limit)
|
||||
const maxLength = 8192; // nomic-embed-text-v1.5 supports 8192 tokens
|
||||
const truncatedText = text.length > maxLength ? text.slice(0, maxLength) : text;
|
||||
|
||||
// Add appropriate prefix
|
||||
if (type === 'search_query') {
|
||||
return `search_query: ${truncatedText}`;
|
||||
} else {
|
||||
return `search_document: ${truncatedText}`;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create text representation of a code node for embedding
|
||||
*
|
||||
* Combines name, signature, docstring, and code snippet into
|
||||
* a searchable text representation.
|
||||
*/
|
||||
static createNodeText(node: {
|
||||
name: string;
|
||||
kind: string;
|
||||
qualifiedName?: string;
|
||||
signature?: string;
|
||||
docstring?: string;
|
||||
filePath: string;
|
||||
}): string {
|
||||
const parts: string[] = [];
|
||||
|
||||
// Add kind and name
|
||||
parts.push(`${node.kind}: ${node.name}`);
|
||||
|
||||
// Add qualified name if different from name
|
||||
if (node.qualifiedName && node.qualifiedName !== node.name) {
|
||||
parts.push(`path: ${node.qualifiedName}`);
|
||||
}
|
||||
|
||||
// Add file path
|
||||
parts.push(`file: ${node.filePath}`);
|
||||
|
||||
// Add signature if present
|
||||
if (node.signature) {
|
||||
parts.push(`signature: ${node.signature}`);
|
||||
}
|
||||
|
||||
// Add docstring if present
|
||||
if (node.docstring) {
|
||||
parts.push(`documentation: ${node.docstring}`);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
}
|
||||
|
||||
/**
|
||||
* Compute cosine similarity between two embeddings
|
||||
*/
|
||||
static cosineSimilarity(a: Float32Array, b: Float32Array): number {
|
||||
if (a.length !== b.length) {
|
||||
throw new Error('Embeddings must have the same dimension');
|
||||
}
|
||||
|
||||
let dotProduct = 0;
|
||||
let normA = 0;
|
||||
let normB = 0;
|
||||
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
const aVal = a[i]!;
|
||||
const bVal = b[i]!;
|
||||
dotProduct += aVal * bVal;
|
||||
normA += aVal * aVal;
|
||||
normB += bVal * bVal;
|
||||
}
|
||||
|
||||
normA = Math.sqrt(normA);
|
||||
normB = Math.sqrt(normB);
|
||||
|
||||
if (normA === 0 || normB === 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
return dotProduct / (normA * normB);
|
||||
}
|
||||
|
||||
/**
|
||||
* Release resources
|
||||
*/
|
||||
dispose(): void {
|
||||
this.pipeline = null;
|
||||
this.initialized = false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a text embedder instance
|
||||
*/
|
||||
export function createEmbedder(options?: EmbedderOptions): TextEmbedder {
|
||||
return new TextEmbedder(options);
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
/**
|
||||
* Vectors Module
|
||||
*
|
||||
* Provides text embedding and vector similarity search for semantic code search.
|
||||
*/
|
||||
|
||||
export {
|
||||
TextEmbedder,
|
||||
createEmbedder,
|
||||
DEFAULT_MODEL,
|
||||
EMBEDDING_DIMENSION,
|
||||
EmbedderOptions,
|
||||
EmbeddingResult,
|
||||
BatchEmbeddingResult,
|
||||
} from './embedder';
|
||||
|
||||
export {
|
||||
VectorSearchManager,
|
||||
createVectorSearch,
|
||||
VectorSearchOptions,
|
||||
} from './search';
|
||||
|
||||
export {
|
||||
VectorManager,
|
||||
createVectorManager,
|
||||
VectorManagerOptions,
|
||||
EmbeddingProgress,
|
||||
} from './manager';
|
||||
@@ -0,0 +1,363 @@
|
||||
/**
|
||||
* Vector Manager
|
||||
*
|
||||
* High-level manager that coordinates embedding generation and vector search.
|
||||
*/
|
||||
|
||||
import Database from 'better-sqlite3';
|
||||
import { Node, SearchResult, SearchOptions } from '../types';
|
||||
import { TextEmbedder, createEmbedder, EmbedderOptions, EMBEDDING_DIMENSION } from './embedder';
|
||||
import { VectorSearchManager, createVectorSearch } from './search';
|
||||
import { QueryBuilder } from '../db/queries';
|
||||
|
||||
/**
|
||||
* Progress callback for embedding generation
|
||||
*/
|
||||
export interface EmbeddingProgress {
|
||||
/** Current node index */
|
||||
current: number;
|
||||
|
||||
/** Total nodes to embed */
|
||||
total: number;
|
||||
|
||||
/** Current node being embedded */
|
||||
nodeName?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Options for the vector manager
|
||||
*/
|
||||
export interface VectorManagerOptions {
|
||||
/** Embedder options */
|
||||
embedder?: EmbedderOptions;
|
||||
|
||||
/** Node kinds to embed (default: functions, methods, classes, interfaces) */
|
||||
nodeKinds?: Node['kind'][];
|
||||
|
||||
/** Batch size for embedding generation */
|
||||
batchSize?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Default node kinds to embed
|
||||
*/
|
||||
const DEFAULT_NODE_KINDS: Node['kind'][] = [
|
||||
'function',
|
||||
'method',
|
||||
'class',
|
||||
'interface',
|
||||
'type_alias',
|
||||
'module',
|
||||
'component',
|
||||
];
|
||||
|
||||
/**
|
||||
* Vector Manager
|
||||
*
|
||||
* Provides high-level interface for semantic search:
|
||||
* - Generates embeddings for code nodes
|
||||
* - Stores embeddings in the database
|
||||
* - Performs semantic similarity search
|
||||
*/
|
||||
export class VectorManager {
|
||||
private embedder: TextEmbedder;
|
||||
private searchManager: VectorSearchManager;
|
||||
private queries: QueryBuilder;
|
||||
private nodeKinds: Node['kind'][];
|
||||
private batchSize: number;
|
||||
private initialized = false;
|
||||
|
||||
constructor(
|
||||
db: Database.Database,
|
||||
queries: QueryBuilder,
|
||||
options: VectorManagerOptions = {}
|
||||
) {
|
||||
this.embedder = createEmbedder(options.embedder);
|
||||
this.searchManager = createVectorSearch(db, EMBEDDING_DIMENSION);
|
||||
this.queries = queries;
|
||||
this.nodeKinds = options.nodeKinds || DEFAULT_NODE_KINDS;
|
||||
this.batchSize = options.batchSize || 32;
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize the vector manager
|
||||
*
|
||||
* Loads the embedding model and initializes vector search.
|
||||
*/
|
||||
async initialize(): Promise<void> {
|
||||
if (this.initialized) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Initialize embedder (downloads model if needed)
|
||||
await this.embedder.initialize();
|
||||
|
||||
// Initialize vector search (loads sqlite-vss if available)
|
||||
await this.searchManager.initialize();
|
||||
|
||||
this.initialized = true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the vector manager is initialized
|
||||
*/
|
||||
isInitialized(): boolean {
|
||||
return this.initialized;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate embeddings for all eligible nodes
|
||||
*
|
||||
* @param onProgress - Optional progress callback
|
||||
* @returns Number of nodes embedded
|
||||
*/
|
||||
async embedAllNodes(onProgress?: (progress: EmbeddingProgress) => void): Promise<number> {
|
||||
if (!this.initialized) {
|
||||
throw new Error('VectorManager not initialized. Call initialize() first.');
|
||||
}
|
||||
|
||||
// Get all nodes that should be embedded
|
||||
const nodesToEmbed: Node[] = [];
|
||||
for (const kind of this.nodeKinds) {
|
||||
const nodes = this.queries.getNodesByKind(kind);
|
||||
nodesToEmbed.push(...nodes);
|
||||
}
|
||||
|
||||
// Filter out nodes that already have embeddings
|
||||
const existingIds = new Set(this.searchManager.getIndexedNodeIds());
|
||||
const newNodes = nodesToEmbed.filter((n) => !existingIds.has(n.id));
|
||||
|
||||
if (newNodes.length === 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Process in batches
|
||||
let processed = 0;
|
||||
const model = this.embedder.getModelId();
|
||||
|
||||
for (let i = 0; i < newNodes.length; i += this.batchSize) {
|
||||
const batch = newNodes.slice(i, i + this.batchSize);
|
||||
|
||||
// Create text representations
|
||||
const texts = batch.map((node) => TextEmbedder.createNodeText(node));
|
||||
|
||||
// Generate embeddings
|
||||
const result = await this.embedder.embedBatch(texts, 'document');
|
||||
|
||||
// Store embeddings
|
||||
const entries: Array<{ nodeId: string; embedding: Float32Array }> = [];
|
||||
for (let idx = 0; idx < batch.length; idx++) {
|
||||
const node = batch[idx];
|
||||
const embedding = result.embeddings[idx];
|
||||
if (node && embedding) {
|
||||
entries.push({ nodeId: node.id, embedding });
|
||||
}
|
||||
}
|
||||
this.searchManager.storeVectorBatch(entries, model);
|
||||
|
||||
processed += batch.length;
|
||||
|
||||
// Report progress
|
||||
if (onProgress) {
|
||||
onProgress({
|
||||
current: processed,
|
||||
total: newNodes.length,
|
||||
nodeName: batch[batch.length - 1]?.name,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return processed;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate embedding for a single node
|
||||
*
|
||||
* @param node - Node to embed
|
||||
*/
|
||||
async embedNode(node: Node): Promise<void> {
|
||||
if (!this.initialized) {
|
||||
throw new Error('VectorManager not initialized. Call initialize() first.');
|
||||
}
|
||||
|
||||
const text = TextEmbedder.createNodeText(node);
|
||||
const result = await this.embedder.embed(text);
|
||||
this.searchManager.storeVector(node.id, result.embedding, result.model);
|
||||
}
|
||||
|
||||
/**
|
||||
* Semantic search for nodes matching a query
|
||||
*
|
||||
* @param query - Natural language query
|
||||
* @param options - Search options
|
||||
* @returns Array of search results with similarity scores
|
||||
*/
|
||||
async search(query: string, options: SearchOptions = {}): Promise<SearchResult[]> {
|
||||
if (!this.initialized) {
|
||||
throw new Error('VectorManager not initialized. Call initialize() first.');
|
||||
}
|
||||
|
||||
const { limit = 10, kinds } = options;
|
||||
|
||||
// Generate query embedding
|
||||
const queryResult = await this.embedder.embedQuery(query);
|
||||
|
||||
// Search for similar vectors
|
||||
const vectorResults = this.searchManager.search(queryResult.embedding, {
|
||||
limit: limit * 2, // Get more results to filter
|
||||
minScore: 0.3, // Minimum similarity threshold
|
||||
});
|
||||
|
||||
// Get nodes and filter by kind if specified
|
||||
const results: SearchResult[] = [];
|
||||
for (const vr of vectorResults) {
|
||||
const node = this.queries.getNodeById(vr.nodeId);
|
||||
if (!node) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Filter by node kind if specified
|
||||
if (kinds && kinds.length > 0 && !kinds.includes(node.kind)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
results.push({
|
||||
node,
|
||||
score: vr.score,
|
||||
});
|
||||
|
||||
if (results.length >= limit) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find nodes similar to a given node
|
||||
*
|
||||
* @param nodeId - ID of the node to find similar nodes for
|
||||
* @param options - Search options
|
||||
* @returns Array of similar nodes with similarity scores
|
||||
*/
|
||||
async findSimilar(nodeId: string, options: SearchOptions = {}): Promise<SearchResult[]> {
|
||||
if (!this.initialized) {
|
||||
throw new Error('VectorManager not initialized. Call initialize() first.');
|
||||
}
|
||||
|
||||
const { limit = 10, kinds } = options;
|
||||
|
||||
// Get the node's embedding
|
||||
let embedding = this.searchManager.getVector(nodeId);
|
||||
|
||||
// If no embedding exists, generate one
|
||||
if (!embedding) {
|
||||
const node = this.queries.getNodeById(nodeId);
|
||||
if (!node) {
|
||||
throw new Error(`Node not found: ${nodeId}`);
|
||||
}
|
||||
|
||||
await this.embedNode(node);
|
||||
embedding = this.searchManager.getVector(nodeId);
|
||||
|
||||
if (!embedding) {
|
||||
throw new Error(`Failed to generate embedding for node: ${nodeId}`);
|
||||
}
|
||||
}
|
||||
|
||||
// Search for similar vectors (excluding the source node)
|
||||
const vectorResults = this.searchManager.search(embedding, {
|
||||
limit: limit + 1, // Get one extra to exclude the source
|
||||
minScore: 0.3,
|
||||
});
|
||||
|
||||
// Get nodes and filter
|
||||
const results: SearchResult[] = [];
|
||||
for (const vr of vectorResults) {
|
||||
// Skip the source node
|
||||
if (vr.nodeId === nodeId) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const node = this.queries.getNodeById(vr.nodeId);
|
||||
if (!node) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Filter by node kind if specified
|
||||
if (kinds && kinds.length > 0 && !kinds.includes(node.kind)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
results.push({
|
||||
node,
|
||||
score: vr.score,
|
||||
});
|
||||
|
||||
if (results.length >= limit) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete embedding for a node
|
||||
*
|
||||
* @param nodeId - ID of the node
|
||||
*/
|
||||
deleteNodeEmbedding(nodeId: string): void {
|
||||
this.searchManager.deleteVector(nodeId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get statistics about vector storage
|
||||
*/
|
||||
getStats(): {
|
||||
totalVectors: number;
|
||||
vssEnabled: boolean;
|
||||
modelId: string;
|
||||
dimension: number;
|
||||
} {
|
||||
return {
|
||||
totalVectors: this.searchManager.getVectorCount(),
|
||||
vssEnabled: this.searchManager.isVssEnabled(),
|
||||
modelId: this.embedder.getModelId(),
|
||||
dimension: this.embedder.getDimension(),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear all vectors
|
||||
*/
|
||||
clear(): void {
|
||||
this.searchManager.clear();
|
||||
}
|
||||
|
||||
/**
|
||||
* Rebuild the VSS index
|
||||
*/
|
||||
rebuildIndex(): void {
|
||||
this.searchManager.rebuildVssIndex();
|
||||
}
|
||||
|
||||
/**
|
||||
* Release resources
|
||||
*/
|
||||
dispose(): void {
|
||||
this.embedder.dispose();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a vector manager
|
||||
*/
|
||||
export function createVectorManager(
|
||||
db: Database.Database,
|
||||
queries: QueryBuilder,
|
||||
options?: VectorManagerOptions
|
||||
): VectorManager {
|
||||
return new VectorManager(db, queries, options);
|
||||
}
|
||||
@@ -0,0 +1,470 @@
|
||||
/**
|
||||
* Vector Search
|
||||
*
|
||||
* Provides vector similarity search using sqlite-vss extension.
|
||||
* Falls back to brute-force cosine similarity if sqlite-vss is not available.
|
||||
*/
|
||||
|
||||
import Database from 'better-sqlite3';
|
||||
import { Node } from '../types';
|
||||
import { TextEmbedder, EMBEDDING_DIMENSION } from './embedder';
|
||||
|
||||
/**
|
||||
* Options for vector search
|
||||
*/
|
||||
export interface VectorSearchOptions {
|
||||
/** Maximum number of results to return */
|
||||
limit?: number;
|
||||
|
||||
/** Minimum similarity score (0-1) */
|
||||
minScore?: number;
|
||||
|
||||
/** Node kinds to filter results */
|
||||
nodeKinds?: Node['kind'][];
|
||||
}
|
||||
|
||||
/**
|
||||
* Vector Search Manager
|
||||
*
|
||||
* Handles vector storage and similarity search for semantic code search.
|
||||
*/
|
||||
export class VectorSearchManager {
|
||||
private db: Database.Database;
|
||||
private vssEnabled = false;
|
||||
private embeddingDimension: number;
|
||||
|
||||
constructor(db: Database.Database, dimension: number = EMBEDDING_DIMENSION) {
|
||||
this.db = db;
|
||||
this.embeddingDimension = dimension;
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize vector search
|
||||
*
|
||||
* Attempts to load sqlite-vss extension. Falls back to brute-force
|
||||
* search if the extension is not available.
|
||||
*/
|
||||
async initialize(): Promise<void> {
|
||||
try {
|
||||
// Try to load sqlite-vss extension
|
||||
await this.loadVssExtension();
|
||||
this.vssEnabled = true;
|
||||
console.log('sqlite-vss extension loaded successfully');
|
||||
|
||||
// Create the VSS virtual table
|
||||
this.createVssTable();
|
||||
} catch (error) {
|
||||
// Fall back to brute-force search
|
||||
console.warn(
|
||||
'sqlite-vss extension not available, falling back to brute-force search:',
|
||||
error instanceof Error ? error.message : String(error)
|
||||
);
|
||||
this.vssEnabled = false;
|
||||
}
|
||||
|
||||
// Ensure the vectors table exists (for both VSS and fallback modes)
|
||||
this.ensureVectorsTable();
|
||||
}
|
||||
|
||||
/**
|
||||
* Load the sqlite-vss extension
|
||||
*/
|
||||
private async loadVssExtension(): Promise<void> {
|
||||
try {
|
||||
// The sqlite-vss npm package provides functions to load extensions
|
||||
const vss = await import('sqlite-vss');
|
||||
|
||||
// Use the load function which loads both vector0 and vss0
|
||||
if (typeof vss.load === 'function') {
|
||||
vss.load(this.db);
|
||||
} else if (typeof vss.default?.load === 'function') {
|
||||
vss.default.load(this.db);
|
||||
} else {
|
||||
throw new Error('sqlite-vss load function not found');
|
||||
}
|
||||
} catch (error) {
|
||||
throw new Error(`Failed to load sqlite-vss: ${error instanceof Error ? error.message : String(error)}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create the VSS virtual table for vector search
|
||||
*/
|
||||
private createVssTable(): void {
|
||||
// Check if the table already exists
|
||||
const tableExists = this.db
|
||||
.prepare("SELECT name FROM sqlite_master WHERE type='table' AND name='vss_vectors'")
|
||||
.get();
|
||||
|
||||
if (!tableExists) {
|
||||
// Create VSS virtual table
|
||||
// vss0 is the vector search extension
|
||||
this.db.exec(`
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS vss_vectors USING vss0(
|
||||
embedding(${this.embeddingDimension})
|
||||
);
|
||||
`);
|
||||
|
||||
// Create mapping table to link VSS rowids to node IDs
|
||||
this.db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS vss_map (
|
||||
rowid INTEGER PRIMARY KEY,
|
||||
node_id TEXT NOT NULL UNIQUE
|
||||
);
|
||||
`);
|
||||
|
||||
// Create index on node_id
|
||||
this.db.exec(`
|
||||
CREATE INDEX IF NOT EXISTS idx_vss_map_node ON vss_map(node_id);
|
||||
`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Ensure the basic vectors table exists (for fallback mode)
|
||||
*/
|
||||
private ensureVectorsTable(): void {
|
||||
this.db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS vectors (
|
||||
node_id TEXT PRIMARY KEY,
|
||||
embedding BLOB NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if VSS extension is enabled
|
||||
*/
|
||||
isVssEnabled(): boolean {
|
||||
return this.vssEnabled;
|
||||
}
|
||||
|
||||
/**
|
||||
* Store a vector embedding for a node
|
||||
*
|
||||
* @param nodeId - ID of the node
|
||||
* @param embedding - Vector embedding
|
||||
* @param model - Model used to generate embedding
|
||||
*/
|
||||
storeVector(nodeId: string, embedding: Float32Array, model: string): void {
|
||||
const now = Date.now();
|
||||
|
||||
// Store in the vectors table (always, for persistence)
|
||||
const blob = Buffer.from(embedding.buffer);
|
||||
this.db
|
||||
.prepare(
|
||||
`
|
||||
INSERT OR REPLACE INTO vectors (node_id, embedding, model, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
`
|
||||
)
|
||||
.run(nodeId, blob, model, now);
|
||||
|
||||
// Also store in VSS table if enabled
|
||||
if (this.vssEnabled) {
|
||||
this.storeInVss(nodeId, embedding);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Store vector in VSS virtual table
|
||||
*/
|
||||
private storeInVss(nodeId: string, embedding: Float32Array): void {
|
||||
try {
|
||||
// Check if already exists
|
||||
const existing = this.db
|
||||
.prepare('SELECT rowid FROM vss_map WHERE node_id = ?')
|
||||
.get(nodeId) as { rowid: number } | undefined;
|
||||
|
||||
if (existing) {
|
||||
// Update existing vector
|
||||
const vectorJson = JSON.stringify(Array.from(embedding));
|
||||
this.db
|
||||
.prepare('UPDATE vss_vectors SET embedding = ? WHERE rowid = ?')
|
||||
.run(vectorJson, existing.rowid);
|
||||
} else {
|
||||
// Insert new vector - get max rowid and increment
|
||||
const maxRow = this.db
|
||||
.prepare('SELECT MAX(rowid) as max FROM vss_map')
|
||||
.get() as { max: number | null } | undefined;
|
||||
const newRowid = (maxRow?.max ?? 0) + 1;
|
||||
|
||||
const vectorJson = JSON.stringify(Array.from(embedding));
|
||||
this.db
|
||||
.prepare('INSERT INTO vss_vectors (rowid, embedding) VALUES (?, ?)')
|
||||
.run(newRowid, vectorJson);
|
||||
|
||||
// Map the rowid to node_id
|
||||
this.db
|
||||
.prepare('INSERT INTO vss_map (rowid, node_id) VALUES (?, ?)')
|
||||
.run(newRowid, nodeId);
|
||||
}
|
||||
} catch (error) {
|
||||
// VSS operations can fail for various reasons (dimension mismatch, etc.)
|
||||
// Fall back to brute-force search silently
|
||||
console.warn(
|
||||
'VSS storage failed, using brute-force search:',
|
||||
error instanceof Error ? error.message : String(error)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Store multiple vectors in a batch
|
||||
*
|
||||
* @param entries - Array of node IDs and embeddings
|
||||
* @param model - Model used to generate embeddings
|
||||
*/
|
||||
storeVectorBatch(
|
||||
entries: Array<{ nodeId: string; embedding: Float32Array }>,
|
||||
model: string
|
||||
): void {
|
||||
const now = Date.now();
|
||||
|
||||
// Use a transaction for better performance
|
||||
this.db.transaction(() => {
|
||||
for (const entry of entries) {
|
||||
const blob = Buffer.from(entry.embedding.buffer);
|
||||
this.db
|
||||
.prepare(
|
||||
`
|
||||
INSERT OR REPLACE INTO vectors (node_id, embedding, model, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
`
|
||||
)
|
||||
.run(entry.nodeId, blob, model, now);
|
||||
|
||||
if (this.vssEnabled) {
|
||||
this.storeInVss(entry.nodeId, entry.embedding);
|
||||
}
|
||||
}
|
||||
})();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get vector for a node
|
||||
*
|
||||
* @param nodeId - ID of the node
|
||||
* @returns Embedding or null if not found
|
||||
*/
|
||||
getVector(nodeId: string): Float32Array | null {
|
||||
const row = this.db
|
||||
.prepare('SELECT embedding FROM vectors WHERE node_id = ?')
|
||||
.get(nodeId) as { embedding: Buffer } | undefined;
|
||||
|
||||
if (!row) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return new Float32Array(row.embedding.buffer.slice(
|
||||
row.embedding.byteOffset,
|
||||
row.embedding.byteOffset + row.embedding.byteLength
|
||||
));
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete vector for a node
|
||||
*
|
||||
* @param nodeId - ID of the node
|
||||
*/
|
||||
deleteVector(nodeId: string): void {
|
||||
this.db.prepare('DELETE FROM vectors WHERE node_id = ?').run(nodeId);
|
||||
|
||||
if (this.vssEnabled) {
|
||||
// Get the rowid before deleting
|
||||
const mapping = this.db
|
||||
.prepare('SELECT rowid FROM vss_map WHERE node_id = ?')
|
||||
.get(nodeId) as { rowid: number } | undefined;
|
||||
|
||||
if (mapping) {
|
||||
this.db.prepare('DELETE FROM vss_vectors WHERE rowid = ?').run(mapping.rowid);
|
||||
this.db.prepare('DELETE FROM vss_map WHERE node_id = ?').run(nodeId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Search for similar vectors
|
||||
*
|
||||
* @param queryEmbedding - Query vector to search for
|
||||
* @param options - Search options
|
||||
* @returns Array of node IDs with similarity scores
|
||||
*/
|
||||
search(
|
||||
queryEmbedding: Float32Array,
|
||||
options: VectorSearchOptions = {}
|
||||
): Array<{ nodeId: string; score: number }> {
|
||||
const { limit = 10, minScore = 0 } = options;
|
||||
|
||||
if (this.vssEnabled) {
|
||||
return this.searchWithVss(queryEmbedding, limit, minScore);
|
||||
} else {
|
||||
return this.searchBruteForce(queryEmbedding, limit, minScore);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Search using sqlite-vss KNN search
|
||||
*/
|
||||
private searchWithVss(
|
||||
queryEmbedding: Float32Array,
|
||||
limit: number,
|
||||
minScore: number
|
||||
): Array<{ nodeId: string; score: number }> {
|
||||
try {
|
||||
const vectorJson = JSON.stringify(Array.from(queryEmbedding));
|
||||
// Sanitize limit to prevent SQL injection (ensure it's a positive integer)
|
||||
const safeLimit = Math.max(1, Math.floor(limit));
|
||||
|
||||
// Use VSS KNN search
|
||||
// The distance is L2 (euclidean), we need to convert to similarity score
|
||||
// Note: sqlite-vss requires LIMIT to be a literal, not a parameter
|
||||
const rows = this.db
|
||||
.prepare(
|
||||
`
|
||||
SELECT
|
||||
vss_map.node_id,
|
||||
vss_vectors.distance
|
||||
FROM vss_vectors
|
||||
JOIN vss_map ON vss_map.rowid = vss_vectors.rowid
|
||||
WHERE vss_search(vss_vectors.embedding, ?)
|
||||
LIMIT ${safeLimit}
|
||||
`
|
||||
)
|
||||
.all(vectorJson) as Array<{ node_id: string; distance: number }>;
|
||||
|
||||
// Convert L2 distance to similarity score (1 / (1 + distance))
|
||||
return rows
|
||||
.map((row) => ({
|
||||
nodeId: row.node_id,
|
||||
score: 1 / (1 + row.distance),
|
||||
}))
|
||||
.filter((r) => r.score >= minScore);
|
||||
} catch (error) {
|
||||
// VSS search failed, fall back to brute force
|
||||
console.warn(
|
||||
'VSS search failed, using brute-force:',
|
||||
error instanceof Error ? error.message : String(error)
|
||||
);
|
||||
return this.searchBruteForce(queryEmbedding, limit, minScore);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Brute-force search using cosine similarity
|
||||
*/
|
||||
private searchBruteForce(
|
||||
queryEmbedding: Float32Array,
|
||||
limit: number,
|
||||
minScore: number
|
||||
): Array<{ nodeId: string; score: number }> {
|
||||
// Get all vectors
|
||||
const rows = this.db
|
||||
.prepare('SELECT node_id, embedding FROM vectors')
|
||||
.all() as Array<{ node_id: string; embedding: Buffer }>;
|
||||
|
||||
// Calculate cosine similarity for each
|
||||
const results: Array<{ nodeId: string; score: number }> = [];
|
||||
|
||||
for (const row of rows) {
|
||||
const embedding = new Float32Array(row.embedding.buffer.slice(
|
||||
row.embedding.byteOffset,
|
||||
row.embedding.byteOffset + row.embedding.byteLength
|
||||
));
|
||||
|
||||
const score = TextEmbedder.cosineSimilarity(queryEmbedding, embedding);
|
||||
|
||||
if (score >= minScore) {
|
||||
results.push({ nodeId: row.node_id, score });
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by score descending and limit
|
||||
results.sort((a, b) => b.score - a.score);
|
||||
return results.slice(0, limit);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get count of stored vectors
|
||||
*/
|
||||
getVectorCount(): number {
|
||||
const result = this.db
|
||||
.prepare('SELECT COUNT(*) as count FROM vectors')
|
||||
.get() as { count: number };
|
||||
return result.count;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a node has a vector
|
||||
*/
|
||||
hasVector(nodeId: string): boolean {
|
||||
const result = this.db
|
||||
.prepare('SELECT 1 FROM vectors WHERE node_id = ? LIMIT 1')
|
||||
.get(nodeId);
|
||||
return !!result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all node IDs that have vectors
|
||||
*/
|
||||
getIndexedNodeIds(): string[] {
|
||||
const rows = this.db
|
||||
.prepare('SELECT node_id FROM vectors')
|
||||
.all() as Array<{ node_id: string }>;
|
||||
return rows.map((r) => r.node_id);
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear all vectors
|
||||
*/
|
||||
clear(): void {
|
||||
this.db.prepare('DELETE FROM vectors').run();
|
||||
|
||||
if (this.vssEnabled) {
|
||||
this.db.prepare('DELETE FROM vss_vectors').run();
|
||||
this.db.prepare('DELETE FROM vss_map').run();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Rebuild VSS index from vectors table
|
||||
*
|
||||
* Useful after bulk operations or if VSS index gets out of sync.
|
||||
*/
|
||||
rebuildVssIndex(): void {
|
||||
if (!this.vssEnabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Clear VSS tables
|
||||
this.db.prepare('DELETE FROM vss_vectors').run();
|
||||
this.db.prepare('DELETE FROM vss_map').run();
|
||||
|
||||
// Reload from vectors table
|
||||
const rows = this.db
|
||||
.prepare('SELECT node_id, embedding FROM vectors')
|
||||
.all() as Array<{ node_id: string; embedding: Buffer }>;
|
||||
|
||||
this.db.transaction(() => {
|
||||
for (const row of rows) {
|
||||
const embedding = new Float32Array(row.embedding.buffer.slice(
|
||||
row.embedding.byteOffset,
|
||||
row.embedding.byteOffset + row.embedding.byteLength
|
||||
));
|
||||
this.storeInVss(row.node_id, embedding);
|
||||
}
|
||||
})();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a vector search manager
|
||||
*/
|
||||
export function createVectorSearch(
|
||||
db: Database.Database,
|
||||
dimension?: number
|
||||
): VectorSearchManager {
|
||||
return new VectorSearchManager(db, dimension);
|
||||
}
|
||||
Reference in New Issue
Block a user