@retrivora-ai/rag-engine 1.8.2 → 1.8.3

package/dist/handlers/index.js

@@ -3380,11 +3380,70 @@ Text to extract from:
 // src/rag/Reranker.ts
 var Reranker = class {
   /**
-   * Re-ranks matches based on a secondary relevance score.
-   * In a production environment, this would call a Cross-Encoder model.
-   * Here we implement a placeholder that filters by score and limits count.
+   * Re-ranks matches based on a secondary relevance score using an LLM.
    */
-  async rerank(matches, query, limit = 5) {
+  async rerank(matches, query, limit = 5, llm) {
+    if (!llm || matches.length <= 1) {
+      const keywords = query.toLowerCase().replace(/[^\w\s]/g, "").split(/\s+/).filter((w) => w.length > 2);
+      if (keywords.length === 0) {
+        return matches.sort((a, b) => b.score - a.score).slice(0, limit);
+      }
+      const scoredMatches = matches.map((match) => {
+        const contentLower = match.content.toLowerCase();
+        let keywordScore = 0;
+        keywords.forEach((keyword) => {
+          if (contentLower.includes(keyword)) {
+            keywordScore += 1;
+          }
+        });
+        const normalizedKeywordScore = keywordScore / keywords.length;
+        const combinedScore = match.score + normalizedKeywordScore * 0.5;
+        return __spreadProps(__spreadValues({}, match), {
+          score: combinedScore
+        });
+      });
+      return scoredMatches.sort((a, b) => b.score - a.score).slice(0, limit);
+    }
+    try {
+      const topN = matches.slice(0, 10);
+      const prompt = `You are a relevance ranking expert.
+Given the following user query and a list of retrieved document chunks, rank the chunks by relevance to the query.
+Return the result as a comma-separated list of document indices in order of relevance, from most relevant to least relevant.
+Use the indices provided in brackets like [0], [1], etc.
+Only return the indices (e.g., "2,0,1"), nothing else. Do not include brackets in the output.
+
+Query: "${query}"
+
+Documents:
+${topN.map((m, i) => `[${i}] ${m.content.replace(/\n/g, " ")}`).join("\n")}`;
+      const response = await llm.chat(
+        [{ role: "user", content: prompt }],
+        "",
+        { temperature: 0, maxTokens: 50 }
+      );
+      const cleanedResponse = response.trim().replace(/[\[\]\s]/g, "");
+      const rankedIndices = cleanedResponse.split(",").map(Number).filter((n) => !isNaN(n));
+      if (rankedIndices.length > 0) {
+        const rankedMatches = [];
+        const usedIndices = /* @__PURE__ */ new Set();
+        for (const index of rankedIndices) {
+          if (index >= 0 && index < topN.length && !usedIndices.has(index)) {
+            rankedMatches.push(topN[index]);
+            usedIndices.add(index);
+          }
+        }
+        for (let i = 0; i < topN.length; i++) {
+          if (!usedIndices.has(i)) {
+            rankedMatches.push(topN[i]);
+          }
+        }
+        const rest = matches.slice(10);
+        rankedMatches.push(...rest);
+        return rankedMatches.slice(0, limit);
+      }
+    } catch (error) {
+      console.warn("[Reranker] LLM re-ranking failed, falling back to score-based sorting:", error);
+    }
     return matches.sort((a, b) => b.score - a.score).slice(0, limit);
   }
 };
@@ -3444,7 +3503,7 @@ var LangChainAgent = class {
    */
   async initialize(chatModel) {
     try {
-      const [{ DynamicTool }, { HumanMessage }, { createAgent }] = await Promise.all([
+      const [{ DynamicTool }, , { createAgent }] = await Promise.all([
         import("@langchain/core/tools"),
         import("@langchain/core/messages"),
         import("langchain")
@@ -3494,7 +3553,6 @@ When presenting structured data, statistics, or comparisons, decide if it is bes
         tools: [searchTool],
         systemPrompt: finalSystemPrompt
       });
-      void HumanMessage;
     } catch (error) {
       const isMissing = error instanceof Error && error.message.includes("Cannot find module");
       const hint = isMissing ? " Make sure 'langchain' and '@langchain/core' are installed:\n  npm install langchain @langchain/core" : "";
@@ -3972,6 +4030,69 @@ var QueryProcessor = class {
     if (filter.keywords && filter.keywords.length === 0) delete filter.keywords;
     return filter;
   }
+  /**
+   * Determines the retrieval strategy based on the query content.
+   */
+  static async determineRetrievalStrategy(query, llm, customGraphKeywords, customVectorKeywords) {
+    if (llm) {
+      try {
+        const prompt = `You are a routing expert in a RAG system.
+Given the user query, determine which retrieval mechanism is needed:
+- 'vector': For semantic search, finding specific documents or content based on similarity.
+- 'graph': For finding relationships between entities, multi-hop questions, or network structures.
+- 'both': If the query requires both specific content and relationship/structural analysis.
+
+Query: "${query}"
+
+Return ONLY 'vector', 'graph', or 'both'. No explanation.`;
+        const response = await llm.chat(
+          [{ role: "user", content: prompt }],
+          "",
+          { temperature: 0, maxTokens: 10 }
+        );
+        const cleanResponse = response.trim().toLowerCase();
+        if (cleanResponse.includes("vector")) return "vector";
+        if (cleanResponse.includes("graph")) return "graph";
+        if (cleanResponse.includes("both")) return "both";
+      } catch (error) {
+        console.warn("[QueryProcessor] LLM strategy classification failed, falling back to keywords:", error);
+      }
+    }
+    const normalized = query.toLowerCase();
+    const graphKeywords = customGraphKeywords || [
+      "relationship",
+      "connect",
+      "link",
+      "network",
+      "friend",
+      "colleague",
+      "manager",
+      "hierarchy",
+      "multi-hop",
+      "shortest path",
+      "between",
+      "related to",
+      "associated with",
+      "belongs to",
+      "part of"
+    ];
+    const hasGraphKeyword = graphKeywords.some((kw) => normalized.includes(kw));
+    const vectorKeywords = customVectorKeywords || [
+      "find",
+      "search",
+      "tell me about",
+      "what is",
+      "how to",
+      "documents about"
+    ];
+    const hasVectorKeyword = vectorKeywords.some((kw) => normalized.includes(kw));
+    if (hasGraphKeyword && !hasVectorKeyword) {
+      return "graph";
+    } else if (hasGraphKeyword && hasVectorKeyword) {
+      return "both";
+    }
+    return "vector";
+  }
 };
 
 // src/utils/synonyms.ts
@@ -5081,6 +5202,7 @@ var Pipeline = class {
         let context = sources.length ? sources.map((m, i) => `[Source ${i + 1}]
 ${m.content}`).join("\n\n---\n\n") : "No relevant context found.";
         if (graphData && graphData.nodes.length > 0) {
+          console.log(`[Graph Retrieval] Found ${graphData.nodes.length} relevant entities.`);
           const graphContext = graphData.nodes.map(
             (n) => `Entity: ${n.label} (${n.id})${n.properties ? " - " + JSON.stringify(n.properties) : ""}`
           ).join("\n");
@@ -5146,7 +5268,7 @@ ${context}`;
         };
         const hallucinationResult = yield new __await(scoreHallucination(this.llmProvider, fullReply, context).catch(() => void 0));
         const trainedSchema = yield new __await(trainingPromise);
-        const uiTransformation = yield new __await(this.generateUiTransformation(question, sources, context, trainedSchema));
+        const uiTransformation = yield new __await(this.generateUiTransformation(question, sources, trainedSchema));
         const trace = {
           requestId,
           query: question,
@@ -5185,7 +5307,7 @@ ${context}`;
    * Universal retrieval method combining all enabled providers.
    * Uses an LRU-bounded embedding cache to avoid re-embedding the same query.
    */
-  async generateUiTransformation(question, sources, _context, trainedSchema) {
+  async generateUiTransformation(question, sources, trainedSchema) {
     if (!sources || sources.length === 0) {
       return UITransformer.transform(question, sources, this.config, trainedSchema);
     }
@@ -5202,15 +5324,19 @@ ${context}`;
     const topK = (_b = options.topK) != null ? _b : 5;
     const cacheKey = `${ns}::${query}`;
     let queryVector = this.embeddingCache.get(cacheKey);
+    const strategy = await QueryProcessor.determineRetrievalStrategy(query, this.llmProvider);
+    console.debug(`[Pipeline] Determined retrieval strategy: ${strategy}`);
     const [retrievedVector, graphData] = await Promise.all([
-      queryVector ? Promise.resolve(queryVector) : this.embeddingProvider.embed(query, { taskType: "query" }),
-      this.graphDB && ((_c = this.config.rag) == null ? void 0 : _c.useGraphRetrieval) ? this.graphDB.query(query) : Promise.resolve(void 0)
+      // Only embed if we need vector search (strategy is 'vector' or 'both')
+      strategy === "vector" || strategy === "both" ? queryVector ? Promise.resolve(queryVector) : this.embeddingProvider.embed(query, { taskType: "query" }) : Promise.resolve([]),
+      // Only query graph if we need graph search (strategy is 'graph' or 'both')
+      (strategy === "graph" || strategy === "both") && this.graphDB && ((_c = this.config.rag) == null ? void 0 : _c.useGraphRetrieval) ? this.graphDB.query(query) : Promise.resolve(void 0)
     ]);
-    if (!queryVector) {
+    if ((strategy === "vector" || strategy === "both") && !queryVector && retrievedVector && retrievedVector.length > 0) {
       this.embeddingCache.set(cacheKey, retrievedVector);
       queryVector = retrievedVector;
     }
-    const sources = await this.vectorDB.query(queryVector, topK, ns, options.filter);
+    const sources = (strategy === "vector" || strategy === "both") && queryVector && queryVector.length > 0 ? await this.vectorDB.query(queryVector, topK, ns, options.filter) : [];
     return { sources, graphData };
   }
   /** Rewrite the user query for better retrieval performance. */
@@ -5239,13 +5365,13 @@ Optimized Search Query:`;
     await this.initialize();
     const ns = namespace != null ? namespace : this.config.projectId;
     try {
-      const { sources } = await this.retrieve(query, { namespace: ns, topK: 3 });
+      const { sources } = await this.retrieve(query, { namespace: ns, topK: 5 });
       if (sources.length === 0) return [];
       const context = sources.map((s) => s.content).join("\n\n---\n\n");
-      const prompt = `Based on the following snippets from a document, what are 3 short, relevant questions a user might ask? 
+      const prompt = `Based on the following snippets from a document, what are at least 5 short, relevant questions a user might ask? 
 Focus on questions that can be answered by the context.
 Keep each question under 10 words and make them very specific to the content.
-Return ONLY a JSON array of strings like ["Question 1", "Question 2", "Question 3"].
+Return ONLY a JSON array of strings like ["Question 1", "Question 2", "Question 3", "Question 4", "Question 5"].
 
 Context:
 ${context}
@@ -5262,7 +5388,7 @@ Suggestions:`;
       if (match) {
         const suggestions = JSON.parse(match[0]);
         if (Array.isArray(suggestions)) {
-          return suggestions.map((s) => String(s)).slice(0, 3);
+          return suggestions.map((s) => String(s)).slice(0, 5);
         }
       }
     } catch (error) {