@juspay/neurolink 9.1.1 → 9.3.0

package/dist/lib/rag/retrieval/hybridSearch.js

@@ -0,0 +1,314 @@
+/**
+ * Hybrid Search Implementation
+ *
+ * Combines vector (dense) search with BM25 (sparse) search for improved retrieval.
+ * Supports multiple fusion methods: Reciprocal Rank Fusion (RRF) and Linear Combination.
+ */
+import { ProviderFactory } from "../../factories/providerFactory.js";
+import { logger } from "../../utils/logger.js";
+import { rerank } from "../reranker/reranker.js";
+/**
+ * In-memory BM25 implementation for testing and development
+ */
+export class InMemoryBM25Index {
+    documents = new Map();
+    avgDocLength = 0;
+    k1 = 1.5; // BM25 parameter
+    b = 0.75; // BM25 parameter
+    async search(query, topK = 10) {
+        const queryTokens = this.tokenize(query);
+        if (queryTokens.length === 0 || this.documents.size === 0) {
+            return [];
+        }
+        // Calculate IDF for each query term
+        const idfValues = new Map();
+        for (const token of queryTokens) {
+            const docCount = this.countDocumentsWithTerm(token);
+            const idf = Math.log((this.documents.size - docCount + 0.5) / (docCount + 0.5) + 1);
+            idfValues.set(token, idf);
+        }
+        // Calculate BM25 score for each document
+        const scores = [];
+        for (const [id, doc] of this.documents) {
+            let score = 0;
+            const docLength = doc.tokens.length;
+            for (const token of queryTokens) {
+                const tf = this.countTermFrequency(doc.tokens, token);
+                const idf = idfValues.get(token) || 0;
+                // BM25 scoring formula
+                const numerator = tf * (this.k1 + 1);
+                const denominator = tf +
+                    this.k1 * (1 - this.b + this.b * (docLength / this.avgDocLength));
+                score += idf * (numerator / denominator);
+            }
+            if (score > 0) {
+                scores.push({
+                    id,
+                    score,
+                    text: doc.text,
+                    metadata: doc.metadata,
+                });
+            }
+        }
+        // Sort by score descending
+        scores.sort((a, b) => b.score - a.score);
+        return scores.slice(0, topK);
+    }
+    async addDocuments(documents) {
+        for (const doc of documents) {
+            const tokens = this.tokenize(doc.text);
+            this.documents.set(doc.id, {
+                text: doc.text,
+                tokens,
+                metadata: doc.metadata || {},
+            });
+        }
+        // Recalculate average document length
+        let totalLength = 0;
+        for (const doc of this.documents.values()) {
+            totalLength += doc.tokens.length;
+        }
+        this.avgDocLength =
+            this.documents.size > 0 ? totalLength / this.documents.size : 0;
+    }
+    tokenize(text) {
+        return text
+            .toLowerCase()
+            .replace(/[^\w\s]/g, " ")
+            .split(/\s+/)
+            .filter((t) => t.length > 0);
+    }
+    countTermFrequency(tokens, term) {
+        return tokens.filter((t) => t === term).length;
+    }
+    countDocumentsWithTerm(term) {
+        let count = 0;
+        for (const doc of this.documents.values()) {
+            if (doc.tokens.includes(term)) {
+                count++;
+            }
+        }
+        return count;
+    }
+}
+/**
+ * Reciprocal Rank Fusion
+ * Combines rankings from multiple retrieval methods
+ *
+ * @param rankings - Array of ranking lists, each with id and rank
+ * @param k - RRF constant (default: 60)
+ * @returns Map of document IDs to fused scores
+ */
+export function reciprocalRankFusion(rankings, k = 60) {
+    const scores = new Map();
+    for (const ranking of rankings) {
+        for (const { id, rank } of ranking) {
+            const currentScore = scores.get(id) || 0;
+            scores.set(id, currentScore + 1 / (k + rank));
+        }
+    }
+    return scores;
+}
+/**
+ * Linear Combination of normalized scores
+ *
+ * @param vectorScores - Vector search scores
+ * @param bm25Scores - BM25 search scores
+ * @param alpha - Weight for vector scores (0-1), bm25 gets 1-alpha
+ * @returns Map of document IDs to combined scores
+ */
+export function linearCombination(vectorScores, bm25Scores, alpha = 0.5) {
+    const combined = new Map();
+    // Get all document IDs
+    const allIds = new Set([...vectorScores.keys(), ...bm25Scores.keys()]);
+    // Normalize scores
+    const normalizedVector = normalizeScores(vectorScores);
+    const normalizedBM25 = normalizeScores(bm25Scores);
+    for (const id of allIds) {
+        const vectorScore = normalizedVector.get(id) || 0;
+        const bm25Score = normalizedBM25.get(id) || 0;
+        combined.set(id, alpha * vectorScore + (1 - alpha) * bm25Score);
+    }
+    return combined;
+}
+/**
+ * Normalize scores to 0-1 range
+ */
+function normalizeScores(scores) {
+    const values = Array.from(scores.values());
+    if (values.length === 0) {
+        return new Map();
+    }
+    const min = Math.min(...values);
+    const max = Math.max(...values);
+    const range = max - min || 1;
+    const normalized = new Map();
+    for (const [id, score] of scores) {
+        normalized.set(id, (score - min) / range);
+    }
+    return normalized;
+}
+/**
+ * Create a hybrid search function
+ *
+ * @param options - Search options
+ * @returns Hybrid search function
+ */
+export function createHybridSearch(options) {
+    const { vectorStore, bm25Index, indexName, embeddingModel, defaultConfig = {}, } = options;
+    /**
+     * Execute hybrid search combining vector and BM25 retrieval
+     *
+     * @param query - Search query
+     * @param config - Search configuration
+     * @returns Hybrid search results
+     */
+    return async function hybridSearch(query, config) {
+        const startTime = Date.now();
+        const { vectorWeight = defaultConfig.vectorWeight ?? 0.5, bm25Weight = defaultConfig.bm25Weight ?? 0.5, fusionMethod = defaultConfig.fusionMethod ?? "rrf", rrfK = defaultConfig.rrfK ?? 60, topK = defaultConfig.topK ?? 10, enableReranking = defaultConfig.enableReranking ?? false, reranker: rerankerConfig = defaultConfig.reranker, } = config || {};
+        try {
+            // Generate query embedding
+            const embeddingProvider = await ProviderFactory.createProvider(embeddingModel?.provider, embeddingModel?.modelName);
+            if (typeof embeddingProvider
+                .embed !== "function") {
+                throw new Error(`Embedding provider does not support the embed() method. ` +
+                    `Please use a provider that supports embeddings (e.g., OpenAI text-embedding-3-small, Vertex text-embedding-004).`);
+            }
+            const queryEmbedding = await embeddingProvider.embed(query);
+            // Parallel retrieval
+            const [vectorResults, bm25Results] = await Promise.all([
+                vectorStore.query({
+                    indexName,
+                    queryVector: queryEmbedding,
+                    topK: topK * 2, // Get more for fusion
+                }),
+                bm25Index.search(query, topK * 2),
+            ]);
+            // Fuse results
+            let fusedResults;
+            if (fusionMethod === "rrf") {
+                // Reciprocal Rank Fusion
+                const vectorRanking = vectorResults.map((r, i) => ({
+                    id: r.id,
+                    rank: i + 1,
+                }));
+                const bm25Ranking = bm25Results.map((r, i) => ({
+                    id: r.id,
+                    rank: i + 1,
+                }));
+                const rrfScores = reciprocalRankFusion([vectorRanking, bm25Ranking], rrfK);
+                // Combine with original data
+                const resultMap = new Map();
+                for (const r of vectorResults) {
+                    resultMap.set(r.id, { text: r.text || "", metadata: r.metadata });
+                }
+                for (const r of bm25Results) {
+                    if (!resultMap.has(r.id)) {
+                        resultMap.set(r.id, { text: r.text, metadata: r.metadata });
+                    }
+                }
+                fusedResults = Array.from(rrfScores.entries())
+                    .sort((a, b) => b[1] - a[1])
+                    .slice(0, topK)
+                    .map(([id, score]) => ({
+                    id,
+                    score,
+                    text: resultMap.get(id)?.text || "",
+                    metadata: resultMap.get(id)?.metadata,
+                    scores: {
+                        combined: score,
+                    },
+                }));
+            }
+            else {
+                // Linear combination
+                const vectorScoreMap = new Map(vectorResults.map((r) => [r.id, r.score || 0]));
+                const bm25ScoreMap = new Map(bm25Results.map((r) => [r.id, r.score]));
+                // Adjust weights based on config
+                const totalWeight = vectorWeight + bm25Weight;
+                const normalizedVectorWeight = vectorWeight / totalWeight;
+                const combinedScores = linearCombination(vectorScoreMap, bm25ScoreMap, normalizedVectorWeight);
+                // Combine with original data
+                const resultMap = new Map();
+                for (const r of vectorResults) {
+                    resultMap.set(r.id, {
+                        text: r.text || "",
+                        metadata: r.metadata,
+                        vectorScore: r.score,
+                    });
+                }
+                for (const r of bm25Results) {
+                    const existing = resultMap.get(r.id);
+                    if (existing) {
+                        existing.bm25Score = r.score;
+                    }
+                    else {
+                        resultMap.set(r.id, {
+                            text: r.text,
+                            metadata: r.metadata,
+                            bm25Score: r.score,
+                        });
+                    }
+                }
+                fusedResults = Array.from(combinedScores.entries())
+                    .sort((a, b) => b[1] - a[1])
+                    .slice(0, topK)
+                    .map(([id, score]) => {
+                    const data = resultMap.get(id);
+                    return {
+                        id,
+                        score,
+                        text: data?.text || "",
+                        metadata: data?.metadata,
+                        scores: {
+                            vector: data?.vectorScore,
+                            bm25: data?.bm25Score,
+                            combined: score,
+                        },
+                    };
+                });
+            }
+            // Apply reranking if configured
+            if (enableReranking && rerankerConfig && fusedResults.length > 0) {
+                const rerankerModel = await ProviderFactory.createProvider(rerankerConfig.model.provider, rerankerConfig.model.modelName);
+                const rerankedResults = await rerank(fusedResults.map((r) => ({
+                    id: r.id,
+                    text: r.text,
+                    score: r.score,
+                    metadata: r.metadata,
+                })), query, rerankerModel, {
+                    weights: rerankerConfig.weights,
+                    topK: rerankerConfig.topK || topK,
+                });
+                fusedResults = rerankedResults.map((r) => ({
+                    id: r.result.id,
+                    score: r.score,
+                    text: r.result.text || "",
+                    metadata: r.result.metadata,
+                    scores: {
+                        ...(fusedResults.find((f) => f.id === r.result.id)?.scores || {}),
+                        reranked: r.score,
+                    },
+                }));
+            }
+            const queryTime = Date.now() - startTime;
+            logger.info("[HybridSearch] Search completed", {
+                query: query.slice(0, 50),
+                vectorResults: vectorResults.length,
+                bm25Results: bm25Results.length,
+                fusedResults: fusedResults.length,
+                fusionMethod,
+                queryTime,
+            });
+            return fusedResults;
+        }
+        catch (error) {
+            logger.error("[HybridSearch] Search failed", {
+                query: query.slice(0, 50),
+                error: error instanceof Error ? error.message : String(error),
+            });
+            throw error;
+        }
+    };
+}
+//# sourceMappingURL=hybridSearch.js.map

package/dist/lib/rag/retrieval/index.d.ts

@@ -0,0 +1,5 @@
+/**
+ * Retrieval Module Exports
+ */
+export { createVectorQueryTool, InMemoryVectorStore, type VectorStore, } from "./vectorQueryTool.js";
+export { createHybridSearch, InMemoryBM25Index, reciprocalRankFusion, linearCombination, type BM25Index, type HybridSearchOptions, } from "./hybridSearch.js";

package/dist/lib/rag/retrieval/index.js

@@ -0,0 +1,6 @@
+/**
+ * Retrieval Module Exports
+ */
+export { createVectorQueryTool, InMemoryVectorStore, } from "./vectorQueryTool.js";
+export { createHybridSearch, InMemoryBM25Index, reciprocalRankFusion, linearCombination, } from "./hybridSearch.js";
+//# sourceMappingURL=index.js.map

package/dist/lib/rag/retrieval/vectorQueryTool.d.ts

@@ -0,0 +1,93 @@
+/**
+ * Vector Query Tool
+ *
+ * Provides semantic search capabilities for RAG pipelines.
+ * Integrates with vector stores and supports metadata filtering and reranking.
+ */
+import { z } from "zod";
+import type { MetadataFilter, RequestContext, VectorQueryResponse, VectorQueryResult, VectorQueryToolConfig } from "../types.js";
+/**
+ * Abstract vector store interface
+ * Vector stores should implement this interface to work with the query tool
+ */
+export interface VectorStore {
+    query(params: {
+        indexName: string;
+        queryVector: number[];
+        topK?: number;
+        filter?: MetadataFilter;
+        includeVectors?: boolean;
+    }): Promise<VectorQueryResult[]>;
+}
+/**
+ * Creates a vector query tool for semantic search
+ * Follows NeuroLink's factory pattern
+ *
+ * @param config - Tool configuration
+ * @param vectorStore - Vector store instance or resolver function
+ * @returns Tool object with execute method
+ */
+export declare function createVectorQueryTool(config: VectorQueryToolConfig, vectorStore: VectorStore | ((context: RequestContext) => VectorStore)): {
+    name: string;
+    description: string;
+    parameters: z.ZodObject<{
+        topK: z.ZodOptional<z.ZodNumber>;
+        filter?: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodUnknown>> | undefined;
+        query: z.ZodString;
+    }, "strip", z.ZodTypeAny, {
+        query: string;
+        filter?: unknown;
+        topK?: number | undefined;
+    }, {
+        query: string;
+        filter?: unknown;
+        topK?: number | undefined;
+    }>;
+    /**
+     * Execute the vector query
+     * @param params - Query parameters
+     * @param context - Optional request context
+     * @returns Query results with relevant context
+     */
+    execute: (params: {
+        query: string;
+        filter?: MetadataFilter;
+        topK?: number;
+    }, context?: RequestContext) => Promise<VectorQueryResponse>;
+};
+/**
+ * In-memory vector store implementation for testing and development
+ */
+export declare class InMemoryVectorStore implements VectorStore {
+    private vectors;
+    /**
+     * Add vectors to an index
+     */
+    upsert(indexName: string, items: Array<{
+        id: string;
+        vector: number[];
+        metadata?: Record<string, unknown>;
+    }>): Promise<void>;
+    /**
+     * Query vectors by similarity
+     */
+    query(params: {
+        indexName: string;
+        queryVector: number[];
+        topK?: number;
+        filter?: MetadataFilter;
+        includeVectors?: boolean;
+    }): Promise<VectorQueryResult[]>;
+    /**
+     * Delete vectors from an index
+     */
+    delete(indexName: string, ids: string[]): Promise<void>;
+    /**
+     * Check if metadata matches filter
+     */
+    private matchesFilter;
+    /**
+     * Calculate cosine similarity between two vectors
+     */
+    private cosineSimilarity;
+}