zrald1/dist/core/vector-store.js

Advanced Graph RAG MCP Server with file location identification, graph processing, and result summarization capabilities

export class VectorStore {
    dimension;
    maxElements;
    nodeMap = new Map();
    chunkMap = new Map();
    fileMap = new Map();
    embeddings = new Map();
    currentIndex = 0;
    constructor(dimension = 384, maxElements = 10000) {
        this.dimension = dimension;
        this.maxElements = maxElements;
    }
    async initialize() {
        console.log(`Vector store initialized with dimension ${this.dimension}`);
    }
    async addNode(node) {
        if (!node.embedding || node.embedding.length !== this.dimension) {
            console.warn(`Node ${node.id} has invalid embedding, skipping`);
            return;
        }
        this.nodeMap.set(node.id, node);
        this.embeddings.set(node.id, node.embedding);
    }
    async addChunk(chunk) {
        if (!chunk.embedding || chunk.embedding.length !== this.dimension) {
            console.warn(`Chunk ${chunk.id} has invalid embedding, skipping`);
            return;
        }
        this.chunkMap.set(chunk.id, chunk);
        this.embeddings.set(chunk.id, chunk.embedding);
    }
    async addFile(file, embedding) {
        this.fileMap.set(file.id, file);
        if (embedding && embedding.length === this.dimension) {
            this.embeddings.set(file.id, embedding);
        }
    }
    async searchNodes(queryEmbedding, topK = 10, threshold = 0.7) {
        if (queryEmbedding.length !== this.dimension) {
            throw new Error(`Query embedding dimension ${queryEmbedding.length} does not match store dimension ${this.dimension}`);
        }
        const searchResults = [];
        // Simple brute force search for demonstration
        for (const [id, embedding] of this.embeddings) {
            const similarity = VectorStore.cosineSimilarity(queryEmbedding, embedding);
            if (similarity >= threshold) {
                const node = this.nodeMap.get(id);
                const chunk = this.chunkMap.get(id);
                const file = this.fileMap.get(id);
                searchResults.push({
                    id,
                    score: similarity,
                    node,
                    chunk,
                    file
                });
            }
        }
        return searchResults
            .sort((a, b) => b.score - a.score)
            .slice(0, topK);
    }
    async searchByNodeTypes(queryEmbedding, nodeTypes, topK = 10, threshold = 0.7) {
        const allResults = await this.searchNodes(queryEmbedding, topK * 2, threshold);
        return allResults
            .filter(result => result.node && nodeTypes.includes(result.node.type))
            .slice(0, topK);
    }
    async searchChunks(queryEmbedding, topK = 10, threshold = 0.7) {
        const allResults = await this.searchNodes(queryEmbedding, topK * 2, threshold);
        return allResults
            .filter(result => result.chunk)
            .slice(0, topK);
    }
    async searchFiles(queryEmbedding, topK = 10, threshold = 0.7) {
        const allResults = await this.searchNodes(queryEmbedding, topK * 2, threshold);
        return allResults
            .filter(result => result.file)
            .slice(0, topK);
    }
    getStats() {
        return {
            totalNodes: this.nodeMap.size,
            totalChunks: this.chunkMap.size,
            totalFiles: this.fileMap.size,
            totalEmbeddings: this.embeddings.size,
            dimension: this.dimension,
            maxElements: this.maxElements
        };
    }
    // Utility methods
    static cosineSimilarity(a, b) {
        if (a.length !== b.length) {
            throw new Error('Vectors must have the same length');
        }
        let dotProduct = 0;
        let normA = 0;
        let normB = 0;
        for (let i = 0; i < a.length; i++) {
            dotProduct += a[i] * b[i];
            normA += a[i] * a[i];
            normB += b[i] * b[i];
        }
        normA = Math.sqrt(normA);
        normB = Math.sqrt(normB);
        if (normA === 0 || normB === 0) {
            return 0;
        }
        return dotProduct / (normA * normB);
    }
    static euclideanDistance(a, b) {
        if (a.length !== b.length) {
            throw new Error('Vectors must have the same length');
        }
        let sum = 0;
        for (let i = 0; i < a.length; i++) {
            const diff = a[i] - b[i];
            sum += diff * diff;
        }
        return Math.sqrt(sum);
    }
    static dotProduct(a, b) {
        if (a.length !== b.length) {
            throw new Error('Vectors must have the same length');
        }
        let product = 0;
        for (let i = 0; i < a.length; i++) {
            product += a[i] * b[i];
        }
        return product;
    }
    // Generate random embedding for testing
    static generateRandomEmbedding(dimension) {
        const embedding = new Array(dimension);
        for (let i = 0; i < dimension; i++) {
            embedding[i] = Math.random() * 2 - 1; // Random values between -1 and 1
        }
        // Normalize the embedding
        const norm = Math.sqrt(embedding.reduce((sum, val) => sum + val * val, 0));
        return embedding.map(val => val / norm);
    }
    // Clear all data
    clear() {
        this.nodeMap.clear();
        this.chunkMap.clear();
        this.fileMap.clear();
        this.embeddings.clear();
        this.currentIndex = 0;
    }
    // Get all nodes
    getAllNodes() {
        return Array.from(this.nodeMap.values());
    }
    // Get all chunks
    getAllChunks() {
        return Array.from(this.chunkMap.values());
    }
    // Get all files
    getAllFiles() {
        return Array.from(this.fileMap.values());
    }
    // Remove node
    removeNode(nodeId) {
        const removed = this.nodeMap.delete(nodeId);
        this.embeddings.delete(nodeId);
        return removed;
    }
    // Remove chunk
    removeChunk(chunkId) {
        const removed = this.chunkMap.delete(chunkId);
        this.embeddings.delete(chunkId);
        return removed;
    }
    // Remove file
    removeFile(fileId) {
        const removed = this.fileMap.delete(fileId);
        this.embeddings.delete(fileId);
        return removed;
    }
}
//# sourceMappingURL=vector-store.js.map