Init
This commit is contained in:
Colby McHenry
@@ -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);
|
||||
}
|
||||
Reference in New Issue
Block a user