@retrivora-ai/rag-engine 1.0.4 → 1.0.5

package/src/core/LangChainAgent.ts

@@ -3,76 +3,109 @@ import { RagConfig } from "../config/RagConfig";
 
 /**
  * LangChainAgent — Orchestration layer that turns the RAG pipeline into an Agent.
- * 
- * Uses dynamic imports to remain an optional dependency.
+ *
+ * Uses dynamic imports to keep LangChain an optional peer dependency.
+ * Requires `langchain` (v1.x) and `@langchain/core` to be installed.
+ *
+ * LangChain v1.x migration notes:
+ * - `AgentExecutor` and `createOpenAIFunctionsAgent` are removed.
+ * - The replacement is `createAgent` from the `langchain` root package,
+ *   which returns a `ReactAgent` (LangGraph-backed ReAct agent).
+ * - Invocation uses `{ messages: [...] }` instead of `{ input, chat_history }`.
  */
 export class LangChainAgent {
-  private executor?: unknown;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  private agent?: any;
 
   constructor(private pipeline: Pipeline, private config: RagConfig) {}
 
   /**
-   * Initializes the agent with the RAG pipeline as a primary tool.
-   * Dynamically imports LangChain dependencies to avoid build errors if missing.
+   * Initializes the LangChain ReAct agent with the RAG pipeline as a tool.
+   * Dynamically imports langchain so the package doesn't crash if it isn't installed.
    */
   async initialize(chatModel: unknown) {
     try {
-      // Dynamic imports for optional dependency support using template strings to hide from TS static analysis
-      const { DynamicTool } = await import(`${"@langchain/core/tools"}`);
-      const { ChatPromptTemplate, MessagesPlaceholder } = await import(`${"@langchain/core/prompts"}`);
-      const { AgentExecutor, createOpenAIFunctionsAgent } = await import(`${"langchain/agents"}`);
+      const [{ DynamicTool }, { HumanMessage }, { createAgent }] = await Promise.all([
+        import('@langchain/core/tools'),
+        import('@langchain/core/messages'),
+        import('langchain'),
+      ]);
 
       // 1. Define the RAG Search Tool
       const searchTool = new DynamicTool({
-        name: "document_search",
-        description: "Use this tool to search through the knowledge base and document repository. Input should be a specific search query.",
+        name: 'document_search',
+        description:
+          'Use this tool to search through the knowledge base and document repository. Input should be a specific search query.',
         func: async (query: string) => {
           const response = await this.pipeline.ask(query);
-          return `Search Results:\n${response.reply}\n\nSources Used: ${JSON.stringify(response.sources.map(s => s.id))}`;
+          return `Search Results:\n${response.reply}\n\nSources Used: ${JSON.stringify(
+            response.sources.map((s) => s.id)
+          )}`;
         },
       });
 
-      const tools = [searchTool];
-
-      // 2. Define the Agent Prompt
-      const prompt = ChatPromptTemplate.fromMessages([
-        ["system", this.config.llm.systemPrompt || "You are a helpful AI assistant with access to a document search tool."],
-        new MessagesPlaceholder("chat_history"),
-        ["human", "{input}"],
-        new MessagesPlaceholder("agent_scratchpad"),
-      ]);
-
-      // 3. Create the Agent
-      const agent = await createOpenAIFunctionsAgent({
-        llm: chatModel,
-        tools,
-        prompt,
+      // 2. Create the ReAct agent (langchain v1.x API)
+      //    - `model` replaces the old `llm` parameter
+      //    - `createAgent` replaces `createOpenAIFunctionsAgent` + `AgentExecutor`
+      this.agent = createAgent({
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        model: chatModel as any,
+        tools: [searchTool],
+        systemPrompt: this.config.llm.systemPrompt ||
+          'You are a helpful AI assistant with access to a document search tool.',
       });
 
-      this.executor = new AgentExecutor({
-        agent,
-        tools,
-      });
+      // Suppress unused import warning — HumanMessage is used in run()
+      void HumanMessage;
     } catch (error) {
-      console.error("[LangChainAgent] Failed to initialize. Ensure 'langchain' and '@langchain/core' are installed.");
-      throw 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"
+        : '';
+      throw new Error(
+        `[LangChainAgent] Failed to initialize.${hint}\n${
+          error instanceof Error ? error.message : String(error)
+        }`
+      );
     }
   }
 
   /**
    * Run the agentic flow.
+   *
+   * LangChain v1.x: invoke takes `{ messages: [...] }` instead of `{ input, chat_history }`.
+   * The agent returns `{ messages: [...] }` — the last message is the final answer.
    */
   async run(input: string, chatHistory: unknown[] = []) {
-    if (!this.executor) {
-      throw new Error("[LangChainAgent] Agent not initialized. Call initialize() first.");
+    if (!this.agent) {
+      throw new Error('[LangChainAgent] Agent not initialized. Call initialize() first.');
     }
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const response = await (this.executor as any).invoke({
-      input,
-      chat_history: chatHistory,
+    // Dynamically import HumanMessage to construct the message list
+    const { HumanMessage, AIMessage } = await import('@langchain/core/messages');
+
+    const historyMessages = (chatHistory as Array<{ role: string; content: string }>).map((m) =>
+      m.role === 'user' ? new HumanMessage(m.content) : new AIMessage(m.content)
+    );
+
+    const response = await this.agent.invoke({
+      messages: [...historyMessages, new HumanMessage(input)],
     });
 
-    return response.output;
+    // Extract the text from the last message in the response
+    const lastMessage = response?.messages?.at(-1);
+    if (lastMessage && typeof lastMessage.content === 'string') {
+      return lastMessage.content;
+    }
+    // Fallback for structured content arrays
+    if (Array.isArray(lastMessage?.content)) {
+      return lastMessage.content
+        .filter((c: { type: string }) => c.type === 'text')
+        .map((c: { text: string }) => c.text)
+        .join('');
+    }
+
+    return String(response?.output ?? response ?? '');
   }
 }

package/src/core/Pipeline.ts

@@ -13,6 +13,52 @@ import { EmbeddingStrategyResolver } from '../config/EmbeddingStrategy';
 import { QueryProcessor } from './QueryProcessor';
 import { IngestDocument, ChatResponse, VectorMatch, GraphSearchResult, RetrievalResult, UpsertDocument } from '../types';
 
+// ─── LRU Embedding Cache ───────────────────────────────────────────────────────
+
+/**
+ * A simple size-capped LRU cache for embedding vectors.
+ * Prevents the unbounded `Map` from consuming infinite memory under production load.
+ */
+class LRUEmbeddingCache {
+  private readonly cache = new Map<string, number[]>();
+  private readonly maxSize: number;
+
+  constructor(maxSize = 500) {
+    this.maxSize = maxSize;
+  }
+
+  get(key: string): number[] | undefined {
+    const value = this.cache.get(key);
+    if (value !== undefined) {
+      // Promote to most-recently-used position
+      this.cache.delete(key);
+      this.cache.set(key, value);
+    }
+    return value;
+  }
+
+  set(key: string, value: number[]): void {
+    if (this.cache.has(key)) {
+      this.cache.delete(key);
+    } else if (this.cache.size >= this.maxSize) {
+      // Evict the oldest (first) entry
+      const oldestKey = this.cache.keys().next().value;
+      if (oldestKey !== undefined) this.cache.delete(oldestKey);
+    }
+    this.cache.set(key, value);
+  }
+
+  clear(): void {
+    this.cache.clear();
+  }
+
+  get size(): number {
+    return this.cache.size;
+  }
+}
+
+// ─── Pipeline ─────────────────────────────────────────────────────────────────
+
 /**
  * Pipeline — orchestrates the RAG flow: Embed → Search → Augment → Generate.
  *
@@ -22,6 +68,7 @@ import { IngestDocument, ChatResponse, VectorMatch, GraphSearchResult, Retrieval
  * - Smart embedding strategy (integrated / separate / external)
  * - Error recovery for transient failures
  * - Multi-tenancy support via namespacing
+ * - LRU-bounded embedding cache (max 500 entries, prevents memory leaks)
  */
 export class Pipeline {
   private vectorDB!: BaseVectorProvider;
@@ -33,7 +80,8 @@ export class Pipeline {
   private entityExtractor?: EntityExtractor;
   private reranker: Reranker;
   private agent?: LangChainAgent;
-  private embeddingCache = new Map<string, number[]>();
+  /** LRU-bounded cache: avoids re-embedding identical queries within the same process. */
+  private embeddingCache = new LRUEmbeddingCache(500);
   private initialised = false;
 
   constructor(private config: RagConfig) {
@@ -144,6 +192,7 @@ export class Pipeline {
 
   /**
    * Step 2: Generate embeddings for chunks with retry logic.
+   * Uses batchEmbed when available for efficiency; falls back to sequential embedding.
    */
   private async processEmbeddings(chunks: Chunk[]): Promise<number[][]> {
     const embedBatchOptions: BatchOptions = {
@@ -158,8 +207,12 @@ export class Pipeline {
       embedBatchOptions
     );
 
+    // mapWithRetry filters out undefined on failure; guard against mismatch explicitly
     if (vectors.length !== chunks.length) {
-      throw new Error(`Embedding mismatch: got ${vectors.length} vectors for ${chunks.length} chunks`);
+      throw new Error(
+        `[Pipeline] Embedding mismatch: got ${vectors.length} vectors for ${chunks.length} chunks. ` +
+        `Check embedding provider logs for individual chunk failures.`
+      );
     }
 
     return vectors;
@@ -313,12 +366,12 @@ export class Pipeline {
 
   /**
    * Universal retrieval method combining all enabled providers.
+   * Uses an LRU-bounded embedding cache to avoid re-embedding the same query.
    */
   async retrieve(query: string, options: { namespace?: string; topK?: number; filter?: Record<string, unknown> }): Promise<RetrievalResult> {
     const ns = options.namespace ?? this.config.projectId;
     const topK = options.topK ?? 5;
 
-    // Performance Optimization: Cache query embeddings to avoid redundant LLM calls
     const cacheKey = `${ns}::${query}`;
     let queryVector = this.embeddingCache.get(cacheKey);
     
@@ -354,10 +407,13 @@ New Question: ${question}
 
 Optimized Search Query:`;
 
-    const rewrite = await this.llmProvider.chat([
-      { role: 'system', content: 'You are an assistant that optimizes search queries for RAG systems.' },
-      { role: 'user', content: prompt }
-    ], '');
+    const rewrite = await this.llmProvider.chat(
+      [
+        { role: 'system', content: 'You are an assistant that optimizes search queries for RAG systems.' },
+        { role: 'user', content: prompt }
+      ],
+      ''
+    );
 
     return rewrite.trim() || question;
   }

package/src/core/ProviderRegistry.ts

@@ -1,5 +1,5 @@
 import { BaseVectorProvider } from '../providers/vectordb/BaseVectorProvider';
-import { BaseGraphProvider } from '@/providers/graphdb/BaseGraphProvider';
+import { BaseGraphProvider } from '../providers/graphdb/BaseGraphProvider';
 import { VectorDBConfig, GraphDBConfig, LLMConfig, EmbeddingConfig } from '../config/RagConfig';
 import { LLMFactory } from '../llm/LLMFactory';
 import { ILLMProvider } from '../llm/ILLMProvider';
@@ -127,7 +127,7 @@ export class ProviderRegistry {
 
     switch (provider) {
       case 'simple': {
-        const { SimpleGraphProvider } = await import('@/providers/graphdb/SimpleGraphProvider');
+        const { SimpleGraphProvider } = await import('../providers/graphdb/SimpleGraphProvider');
         return new SimpleGraphProvider(config);
       }
       default:

package/src/core/QueryProcessor.ts

@@ -137,7 +137,7 @@ export class QueryProcessor {
       }
     }
 
-    // 6. Capitalized proper nouns
+    // 6. Capitalized proper nouns — added here so buildQueryFilter can reuse them
     for (const match of question.matchAll(/\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+){0,3}\b/g)) {
       addHint(match[0]);
     }
@@ -147,6 +147,10 @@ export class QueryProcessor {
 
   /**
    * Constructs a QueryFilter object from extracted hints.
+   *
+   * Note: proper nouns are extracted inside `extractQueryFieldHints` and surfaced
+   * as field-less hints; `buildQueryFilter` adds them to `keywords` without
+   * re-running the proper noun regex to avoid duplication.
    */
   static buildQueryFilter(question: string, hints: QueryFieldHint[]): QueryFilter {
     const filter: QueryFilter = { metadata: {}, keywords: [], queryText: question };
@@ -155,16 +159,13 @@ export class QueryProcessor {
       if (hint.field) {
         filter.metadata![hint.field] = hint.value;
       } else {
-        filter.keywords!.push(hint.value);
+        // hint.value already includes proper nouns from extractQueryFieldHints
+        if (!filter.keywords!.includes(hint.value)) {
+          filter.keywords!.push(hint.value);
+        }
       }
     }
 
-    // Capture proper-noun sequences as keywords
-    for (const match of question.matchAll(/\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+){0,3}\b/g)) {
-      const term = this.normalizeHintValue(match[0]);
-      if (term && !filter.keywords!.includes(term)) filter.keywords!.push(term);
-    }
-
     if (Object.keys(filter.metadata || {}).length === 0) delete filter.metadata;
     if (filter.keywords && filter.keywords.length === 0) delete filter.keywords;
 

package/src/handlers/index.ts

@@ -4,11 +4,60 @@ import { RagConfig } from '../config/RagConfig';
 import { ChatMessage } from '../llm/ILLMProvider';
 import { DocumentParser } from '../utils/DocumentParser';
 
+// ─── SSE helpers ──────────────────────────────────────────────────────────────
+
+/**
+ * Encode a payload as a Server-Sent Events frame.
+ *   data: <json>\n\n
+ */
+function sseFrame(payload: unknown): string {
+  return `data: ${JSON.stringify(payload)}\n\n`;
+}
+
+/** Encode a plain text chunk as an SSE frame. */
+function sseTextFrame(text: string): string {
+  return `data: ${JSON.stringify({ type: 'text', text })}\n\n`;
+}
+
+/** Encode the retrieval metadata as an SSE frame. */
+function sseMetaFrame(meta: unknown): string {
+  return `data: ${JSON.stringify({ type: 'metadata', ...((meta as object) ?? {}) })}\n\n`;
+}
+
+/** Encode a stream error as an SSE frame. */
+function sseErrorFrame(message: string): string {
+  return `data: ${JSON.stringify({ type: 'error', error: message })}\n\n`;
+}
+
+/** SSE response headers. */
+const SSE_HEADERS: Record<string, string> = {
+  'Content-Type': 'text/event-stream; charset=utf-8',
+  'Cache-Control': 'no-cache, no-transform',
+  Connection: 'keep-alive',
+  'X-Accel-Buffering': 'no', // Disable Nginx buffering for streaming
+};
+
+// ─── Handler Factories ────────────────────────────────────────────────────────
+
 /**
- * createChatHandler — factory that returns a Next.js App Router POST handler
+ * createChatHandler — factory that returns a Next.js App Router POST handler.
+ *
+ * Accepts either a full/partial `RagConfig` **or** a pre-built `VectorPlugin`
+ * instance (preferred for singleton/multi-tenant setups).
+ *
+ * @example
+ * // Option A — pass config (plugin created once at module load time)
+ * export const POST = createChatHandler(getRagConfig());
+ *
+ * // Option B — pass a pre-built plugin (most flexible)
+ * const plugin = new VectorPlugin(getRagConfig());
+ * export const POST = createChatHandler(plugin);
  */
-export function createChatHandler(config?: Partial<RagConfig>) {
-  const plugin = new VectorPlugin(config);
+export function createChatHandler(configOrPlugin?: Partial<RagConfig> | VectorPlugin) {
+  const plugin =
+    configOrPlugin instanceof VectorPlugin
+      ? configOrPlugin
+      : new VectorPlugin(configOrPlugin);
 
   return async function POST(req: NextRequest) {
     try {
@@ -33,63 +82,94 @@ export function createChatHandler(config?: Partial<RagConfig>) {
 }
 
 /**
- * createStreamHandler — factory that returns a streaming Next.js App Router POST handler
+ * createStreamHandler — factory for a streaming SSE Next.js App Router POST handler.
+ *
+ * Streams the assistant response as Server-Sent Events (SSE). Each frame is a
+ * JSON object with a `type` discriminant:
+ *   - `{ type: "text",     text: "..." }`       — incremental text chunk
+ *   - `{ type: "metadata", sources: [...] }`    — retrieval metadata (last frame)
+ *   - `{ type: "error",    error: "..." }`       — stream-level error
+ *
+ * @example
+ * // src/app/api/chat/route.ts
+ * export const POST = createStreamHandler(getRagConfig());
+ *
+ * // Client-side parsing in useRagChat.ts:
+ * // const lines = chunk.split('\n');
+ * // for (const line of lines) {
+ * //   if (line.startsWith('data: ')) {
+ * //     const frame = JSON.parse(line.slice(6));
+ * //     if (frame.type === 'text') { ... }
+ * //   }
+ * // }
  */
-export function createStreamHandler(config?: Partial<RagConfig>) {
-  const plugin = new VectorPlugin(config);
+export function createStreamHandler(configOrPlugin?: Partial<RagConfig> | VectorPlugin) {
+  const plugin =
+    configOrPlugin instanceof VectorPlugin
+      ? configOrPlugin
+      : new VectorPlugin(configOrPlugin);
 
   return async function POST(req: NextRequest) {
+    // Parse body first so we can return a normal JSON 400 before opening the stream
+    let body: { message?: string; history?: ChatMessage[]; namespace?: string };
     try {
-      const body = await req.json();
-      const { message, history = [], namespace } = body;
-      
-      const encoder = new TextEncoder();
-      const stream = new ReadableStream({
-        async start(controller) {
-          try {
-            const pipelineStream = plugin.chatStream(message, history, namespace);
-            
-            for await (const chunk of pipelineStream) {
-              if (typeof chunk === 'string') {
-                controller.enqueue(encoder.encode(chunk));
-              } else {
-                // Yield metadata/sources at the end
-                controller.enqueue(encoder.encode(`\n\n__METADATA__${JSON.stringify(chunk)}`));
-              }
-            }
-            controller.close();
-          } catch (streamError) {
-            console.error('[createStreamHandler] Stream processing error:', streamError);
-            const errorMessage = streamError instanceof Error ? streamError.message : String(streamError);
-            controller.enqueue(encoder.encode(`\n\n__ERROR__${JSON.stringify({ error: errorMessage })}`));
-            controller.close();
-          }
-        },
+      body = await req.json();
+    } catch {
+      return new Response(JSON.stringify({ error: 'Invalid JSON body' }), {
+        status: 400,
+        headers: { 'Content-Type': 'application/json' },
       });
+    }
 
-      return new Response(stream, {
-        headers: {
-          'Content-Type': 'text/event-stream',
-          'Cache-Control': 'no-cache',
-          'Connection': 'keep-alive',
-        },
-      });
-    } catch (err) {
-      const message = err instanceof Error ? err.message : 'Internal server error';
-      return new Response(JSON.stringify({ error: message }), {
-        status: 500,
+    const { message, history = [], namespace } = body;
+
+    if (!message?.trim()) {
+      return new Response(JSON.stringify({ error: 'message is required' }), {
+        status: 400,
         headers: { 'Content-Type': 'application/json' },
       });
     }
+
+    const encoder = new TextEncoder();
+
+    const stream = new ReadableStream({
+      async start(controller) {
+        const enqueue = (text: string) => controller.enqueue(encoder.encode(text));
+
+        try {
+          const pipelineStream = plugin.chatStream(message, history, namespace);
+
+          for await (const chunk of pipelineStream) {
+            if (typeof chunk === 'string') {
+              enqueue(sseTextFrame(chunk));
+            } else {
+              // Retrieval metadata object — always the final frame
+              enqueue(sseMetaFrame(chunk));
+            }
+          }
+        } catch (streamError) {
+          const errorMessage =
+            streamError instanceof Error ? streamError.message : String(streamError);
+          console.error('[createStreamHandler] Stream error:', streamError);
+          enqueue(sseErrorFrame(errorMessage));
+        } finally {
+          controller.close();
+        }
+      },
+    });
+
+    return new Response(stream, { headers: SSE_HEADERS });
   };
 }
 
-
 /**
  * createIngestHandler — factory for the document ingestion endpoint.
  */
-export function createIngestHandler(config?: Partial<RagConfig>) {
-  const plugin = new VectorPlugin(config);
+export function createIngestHandler(configOrPlugin?: Partial<RagConfig> | VectorPlugin) {
+  const plugin =
+    configOrPlugin instanceof VectorPlugin
+      ? configOrPlugin
+      : new VectorPlugin(configOrPlugin);
 
   return async function POST(req: NextRequest) {
     try {
@@ -115,8 +195,11 @@ export function createIngestHandler(config?: Partial<RagConfig>) {
 /**
  * createHealthHandler — factory for the health-check endpoint.
  */
-export function createHealthHandler(config?: Partial<RagConfig>) {
-  const plugin = new VectorPlugin(config);
+export function createHealthHandler(configOrPlugin?: Partial<RagConfig> | VectorPlugin) {
+  const plugin =
+    configOrPlugin instanceof VectorPlugin
+      ? configOrPlugin
+      : new VectorPlugin(configOrPlugin);
 
   return async function GET() {
     try {
@@ -137,8 +220,11 @@ export function createHealthHandler(config?: Partial<RagConfig>) {
 /**
  * createUploadHandler — factory for the file upload ingestion endpoint.
  */
-export function createUploadHandler(config?: Partial<RagConfig>) {
-  const plugin = new VectorPlugin(config);
+export function createUploadHandler(configOrPlugin?: Partial<RagConfig> | VectorPlugin) {
+  const plugin =
+    configOrPlugin instanceof VectorPlugin
+      ? configOrPlugin
+      : new VectorPlugin(configOrPlugin);
 
   return async function POST(req: NextRequest) {
     try {
@@ -174,3 +260,6 @@ export function createUploadHandler(config?: Partial<RagConfig>) {
     }
   };
 }
+
+// Re-export SSE helper so host apps can use it in custom handlers
+export { sseFrame, sseTextFrame, sseMetaFrame, sseErrorFrame };