@retrivora-ai/rag-engine 1.9.7 → 1.9.9

package/src/core/Pipeline.ts

@@ -9,6 +9,8 @@ import { Reranker } from '../rag/Reranker';
 import { LlamaIndexIngestor } from '../rag/LlamaIndexIngestor';
 import { LangChainAgent } from './LangChainAgent';
 import { RagConfig } from '../config/RagConfig';
+import { MCPRegistry } from './mcp';
+import { MultiAgentCoordinator } from './MultiAgentCoordinator';
 import { ProviderRegistry } from './ProviderRegistry';
 import { BatchProcessor, BatchOptions } from './BatchProcessor';
 import { EmbeddingStrategyResolver } from '../config/EmbeddingStrategy';
@@ -165,9 +167,13 @@ export class Pipeline {
   private entityExtractor?: EntityExtractor;
   private reranker: Reranker;
   private agent?: LangChainAgent;
+  private mcpRegistry?: MCPRegistry;
+  private multiAgentCoordinator?: MultiAgentCoordinator;
   /** LRU-bounded cache: avoids re-embedding identical queries within the same process. */
   private embeddingCache = new LRUEmbeddingCache(500);
   private initialised = false;
+  /** Namespace-specific static cold context cache for CAG */
+  private coldContexts = new Map<string, string>();
 
   constructor(private config: RagConfig) {
     this.chunker = new DocumentChunker(
@@ -250,14 +256,52 @@ You are a helpful assistant. Use the provided context to answer questions accura
 
     await this.vectorDB.initialize();
 
-    if (this.config.rag?.architecture === 'agentic') {
+    if (this.config.rag?.architecture === 'agentic' || this.config.rag?.architecture === 'multi-agent' || (this.config.mcpServers && this.config.mcpServers.length > 0)) {
+      this.mcpRegistry = new MCPRegistry(this.config.mcpServers || []);
+      this.multiAgentCoordinator = new MultiAgentCoordinator({
+        llmProvider: this.llmProvider,
+        mcpRegistry: this.mcpRegistry,
+        documentSearch: async (query: string) => {
+          return this.runNormalQuery(query);
+        }
+      });
       this.agent = new LangChainAgent(this, this.config);
       await this.agent.initialize(this.llmProvider as unknown);
     }
 
+    if (this.config.rag?.cag?.enabled) {
+      await this.loadColdContext(this.config.projectId);
+    }
+
     this.initialised = true;
   }
 
+  private async loadColdContext(ns: string): Promise<void> {
+    const cagConfig = this.config.rag?.cag;
+    if (!cagConfig || !cagConfig.enabled) return;
+
+    try {
+      const coldNs = cagConfig.coldNamespace || ns;
+      const limit = cagConfig.coldLimit ?? 20;
+
+      const queryText = cagConfig.coldQuery || 'documentation';
+      const queryVector = await this.embeddingProvider.embed(queryText, { taskType: 'query' });
+      const coldMatches = await this.vectorDB.query(queryVector, limit, coldNs);
+
+      if (coldMatches.length > 0) {
+        const formatted = coldMatches.map((m, i) => `[Cold Source ${i + 1}]\n${m.content}`).join('\n\n---\n\n');
+        this.coldContexts.set(ns, formatted);
+        console.log(`[Pipeline] Cold RAG (CAG) loaded ${coldMatches.length} documents into cached context for namespace: ${ns}`);
+      } else {
+        this.coldContexts.set(ns, 'No cold context found.');
+        console.warn(`[Pipeline] Cold RAG (CAG) initialized but no documents were found in namespace: ${coldNs}`);
+      }
+    } catch (err) {
+      console.error(`[Pipeline] Failed to load Cold RAG (CAG) context for namespace ${ns}:`, err);
+      this.coldContexts.set(ns, 'Failed to load cold context.');
+    }
+  }
+
   /**
    * Ingest documents with automatic chunking, embedding, and batch upsert.
    */
@@ -386,12 +430,26 @@ You are a helpful assistant. Use the provided context to answer questions accura
     );
   }
 
+  async runNormalQuery(question: string, history: ChatMessage[] = [], namespace?: string): Promise<{ reply: string; sources: any[] }> {
+    const stream = this.askStreamInternal(question, history, namespace);
+    let reply = '';
+    let sources: any[] = [];
+    for await (const chunk of stream) {
+      if (typeof chunk === 'string') {
+        reply += chunk;
+      } else if (typeof chunk === 'object' && chunk !== null) {
+        if ('reply' in chunk && chunk.reply) reply += chunk.reply;
+        if ('sources' in chunk && chunk.sources) sources = chunk.sources;
+      }
+    }
+    return { reply, sources };
+  }
+
   async ask(question: string, history: ChatMessage[] = [], namespace?: string): Promise<ChatResponse> {
     await this.initialize();
 
-    if (this.config.rag?.architecture === 'agentic' && this.agent) {
-      const agentReply = await this.agent.run(question, history);
-      return { reply: agentReply, sources: [] };
+    if ((this.config.rag?.architecture === 'agentic' || this.config.rag?.architecture === 'multi-agent' || (this.config.mcpServers && this.config.mcpServers.length > 0)) && this.multiAgentCoordinator) {
+      return await this.multiAgentCoordinator.run(question, history);
     }
 
     const stream = this.askStream(question, history, namespace);
@@ -415,6 +473,21 @@ You are a helpful assistant. Use the provided context to answer questions accura
     return { reply, sources, graphData, ui_transformation: uiTransformation, trace };
   }
 
+  async *askStream(question: string, history: ChatMessage[] = [], namespace?: string): AsyncIterable<string | ChatResponse> {
+    await this.initialize();
+    if ((this.config.rag?.architecture === 'agentic' || this.config.rag?.architecture === 'multi-agent' || (this.config.mcpServers && this.config.mcpServers.length > 0)) && this.multiAgentCoordinator) {
+      const stream = this.multiAgentCoordinator.runStream(question, history);
+      for await (const chunk of stream) {
+        yield chunk;
+      }
+      return;
+    }
+    const stream = this.askStreamInternal(question, history, namespace);
+    for await (const chunk of stream) {
+      yield chunk;
+    }
+  }
+
   /**
    * High-performance streaming RAG flow.
    * Yields text chunks first, then the retrieval metadata + observability trace at the end.
@@ -425,7 +498,7 @@ You are a helpful assistant. Use the provided context to answer questions accura
    * - UITransformation is computed after text streaming and emitted with metadata
    * - SchemaMapper.train runs while answer generation streams
    */
-  async *askStream(question: string, history: ChatMessage[] = [], namespace?: string): AsyncIterable<string | ChatResponse> {
+  async *askStreamInternal(question: string, history: ChatMessage[] = [], namespace?: string): AsyncIterable<string | ChatResponse> {
     await this.initialize();
     const ns = namespace ?? this.config.projectId;
     const topK = this.config.rag?.topK ?? 5;
@@ -462,7 +535,8 @@ You are a helpful assistant. Use the provided context to answer questions accura
       const cacheKey = `${ns}::${searchQuery}`;
       const cachedVector = this.embeddingCache.get(cacheKey);
 
-      const useGraph = this.graphDB && this.config.rag?.useGraphRetrieval;
+      const arch = this.config.rag?.architecture;
+      const useGraph = arch !== 'naive' && arch !== 'retrieve-and-rerank' && this.graphDB && this.config.rag?.useGraphRetrieval;
       const [strategyResult, embeddedVector] = await Promise.all([
         QueryProcessor.determineRetrievalStrategy(
           searchQuery,
@@ -506,7 +580,8 @@ You are a helpful assistant. Use the provided context to answer questions accura
         ? structuredSources
         : structuredSources.filter((m) => m.score >= scoreThreshold);
       const rerankLimit = wantsExhaustiveList ? retrievalLimit : topK;
-      if (!hasNumericPredicates && this.config.rag?.useReranking) {
+      const useReranking = arch !== 'naive' && (arch === 'retrieve-and-rerank' || this.config.rag?.useReranking);
+      if (!hasNumericPredicates && useReranking) {
         fullSources = await this.reranker.rerank(fullSources, question, rerankLimit);
       } else if (!wantsExhaustiveList) {
         fullSources = fullSources.slice(0, topK);
@@ -547,6 +622,28 @@ You are a helpful assistant. Use the provided context to answer questions accura
         context = `GRAPH KNOWLEDGE:\n${graphContext}\n\nVECTOR CONTEXT:\n${context}`;
       }
 
+      if (this.config.rag?.cag?.enabled) {
+        if (!this.coldContexts.has(ns)) {
+          await this.loadColdContext(ns);
+        }
+        const coldContext = this.coldContexts.get(ns);
+        if (coldContext && coldContext !== 'No cold context found.' && coldContext !== 'Failed to load cold context.') {
+          context = `COLD CONTEXT (STATIC KNOWLEDGE):\n${coldContext}\n\nRETRIEVED HOT CONTEXT (REAL-TIME KNOWLEDGE):\n${context}`;
+        }
+      }
+
+      // Yield sources immediately to support early UI citation rendering (highly responsive)
+      yield {
+        reply: '',
+        sources: sources.map((s) => ({
+          id: s.id,
+          score: s.score,
+          content: s.content,
+          metadata: s.metadata ?? {},
+          namespace: ns,
+        } satisfies RetrievedChunk)),
+      } as ChatResponse;
+
       // 4.5 Schema Training — true fire-and-forget, never gates anything.
       const allMetadataKeys = Array.from(new Set(sources.flatMap(s => Object.keys(s.metadata || {}))));
       const cachedSchema = allMetadataKeys.length > 0
@@ -735,7 +832,7 @@ You are a helpful assistant. Use the provided context to answer questions accura
       const latency: LatencyBreakdown = {
         embedMs: Math.round(embedMs),
         retrieveMs: Math.round(retrieveMs),
-        rerankMs: this.config.rag?.useReranking ? Math.round(rerankMs) : undefined,
+        rerankMs: (arch !== 'naive' && (arch === 'retrieve-and-rerank' || this.config.rag?.useReranking)) ? Math.round(rerankMs) : undefined,
         generateMs: Math.round(generateMs),
         totalMs: Math.round(totalMs),
       };
@@ -751,15 +848,15 @@ You are a helpful assistant. Use the provided context to answer questions accura
         estimatedCostUsd: estimateCostUsd(promptTokens, completionTokens, this.config.llm.model),
       };
 
-      // 7. Await both UI transformation and hallucination score concurrently
-      // before yielding the metadata frame. This guarantees the trace is always
-      // complete when the consumer receives it — no post-yield mutation needed.
+      // 7. Await UI transformation and optionally hallucination score before yielding the metadata frame.
+      const awaitHallucination = this.config.llm.options?.awaitHallucination === true;
+      
       const [uiTransformation, hallucinationResult] = await Promise.all([
         uiTransformationPromise,
-        hallucinationScoringPromise,
+        awaitHallucination ? hallucinationScoringPromise : Promise.resolve(undefined),
       ]);
 
-      const trace: ObservabilityTrace = {
+      const buildTrace = (hScore?: { score: number; reason: string }): ObservabilityTrace => ({
         requestId,
         query: question,
         rewrittenQuery,
@@ -775,12 +872,48 @@ You are a helpful assistant. Use the provided context to answer questions accura
         latency,
         tokens,
         timestamp: new Date().toISOString(),
-        // Hallucination score is now always present when available
-        ...(hallucinationResult ? {
-          hallucinationScore: hallucinationResult.score,
-          hallucinationReason: hallucinationResult.reason,
+        ...(hScore ? {
+          hallucinationScore: hScore.score,
+          hallucinationReason: hScore.reason,
         } : {}),
-      };
+      });
+
+      const trace = buildTrace(hallucinationResult);
+
+      // Asynchronously trigger telemetry logging without blocking response yield
+      if (this.config.telemetry?.enabled) {
+        const telemetryUrl = this.config.telemetry.url || '/api/telemetry';
+        const absoluteUrl = telemetryUrl.startsWith('http')
+          ? telemetryUrl
+          : (process.env.NEXT_PUBLIC_APP_URL || `http://localhost:${process.env.PORT || 3000}`) + telemetryUrl;
+
+        // Start background process for telemetry, awaiting hallucination scoring if it was backgrounded
+        (async () => {
+          try {
+            let finalTrace = trace;
+            if (!awaitHallucination && runHallucination) {
+              const backgroundScoreResult = await hallucinationScoringPromise;
+              if (backgroundScoreResult) {
+                finalTrace = buildTrace(backgroundScoreResult);
+              }
+            }
+            
+            await fetch(absoluteUrl, {
+              method: 'POST',
+              headers: {
+                'Content-Type': 'application/json',
+              },
+              body: JSON.stringify({
+                trace: finalTrace,
+                licenseKey: this.config.licenseKey,
+                projectId: ns
+              })
+            });
+          } catch (err: any) {
+            console.warn(`[Pipeline Telemetry] Failed to send trace async to ${absoluteUrl}:`, err.message);
+          }
+        })();
+      }
 
       yield {
         reply: '',

package/src/core/Retrivora.ts

@@ -2,6 +2,7 @@ import { RagConfig, RetrievalConfig, UniversalRagConfig } from '../config/RagCon
 import { ChatMessage, ChatResponse, IngestDocument } from '../types';
 import { ConfigResolver } from './ConfigResolver';
 import { ConfigValidator } from './ConfigValidator';
+import { LicenseVerifier } from './LicenseVerifier';
 import { Pipeline } from './Pipeline';
 import { wrapError, RetrivoraErrorCode } from '../exceptions';
 
@@ -23,6 +24,12 @@ export class Retrivora {
   async initialize(): Promise<void> {
     try {
       await ConfigValidator.validateAndThrow(this.config);
+      LicenseVerifier.verify(
+        this.config.licenseKey,
+        this.config.projectId,
+        this.config.vectorDb?.provider
+      );
+      await LicenseVerifier.verifyIP(this.config.licenseKey);
     } catch (err) {
       throw wrapError(err, 'CONFIGURATION_ERROR');
     }

package/src/core/VectorPlugin.ts

@@ -5,6 +5,7 @@ import { Pipeline } from './Pipeline';
 import { ProviderHealthCheck, HealthCheckResult } from './ProviderHealthCheck';
 import { IngestDocument, ChatResponse } from '../types';
 import { ChatMessage } from '../types';
+import { LicenseVerifier } from './LicenseVerifier';
 import { wrapError, RetrivoraErrorCode } from '../exceptions';
 
 /**
@@ -24,14 +25,22 @@ export class VectorPlugin {
   private validationPromise: Promise<void>;
 
   constructor(hostConfig?: Partial<RagConfig>) {
-    this.config = ConfigResolver.resolve(hostConfig);
+    const resolvedConfig = ConfigResolver.resolve(hostConfig);
+    this.config = resolvedConfig;
     
     // Start validation early
-    this.validationPromise = ConfigValidator.validateAndThrow(this.config);
+    this.validationPromise = (async () => {
+      await ConfigValidator.validateAndThrow(resolvedConfig);
+      LicenseVerifier.verify(
+        resolvedConfig.licenseKey,
+        resolvedConfig.projectId,
+        resolvedConfig.vectorDb?.provider
+      );
+    })();
     // Prevent unhandled promise rejection warning
     this.validationPromise.catch(() => {});
     
-    this.pipeline = new Pipeline(this.config);
+    this.pipeline = new Pipeline(resolvedConfig);
   }
 
   /**

package/src/core/mcp.ts

@@ -0,0 +1,261 @@
+import { spawn, ChildProcess } from 'child_process';
+import { MCPServerConfig } from '../config/RagConfig';
+
+interface MCPTool {
+  name: string;
+  description?: string;
+  inputSchema: any;
+}
+
+interface MCPCallResult {
+  content: Array<{ type: string; text?: string; [key: string]: any }>;
+  isError?: boolean;
+}
+
+export class MCPClient {
+  private childProcess?: ChildProcess;
+  private sseUrl?: string;
+  private messageIdCounter = 0;
+  private pendingRequests = new Map<number, { resolve: (val: any) => void; reject: (err: any) => void }>();
+  private stdioBuffer = '';
+  private initialized = false;
+
+  constructor(private config: MCPServerConfig) {}
+
+  private getNextMessageId(): number {
+    return ++this.messageIdCounter;
+  }
+
+  /**
+   * Connect and initialize the MCP server connection.
+   */
+  async connect(): Promise<void> {
+    if (this.initialized) return;
+
+    if (this.config.transport === 'stdio') {
+      await this.connectStdio();
+    } else {
+      await this.connectSSE();
+    }
+
+    // Call initialize
+    try {
+      const response = await this.sendJsonRpc('initialize', {
+        protocolVersion: '2024-11-05',
+        capabilities: {},
+        clientInfo: {
+          name: 'retrivora-sdk',
+          version: '1.0.0',
+        },
+      });
+
+      // Send initialized notification
+      await this.sendNotification('notifications/initialized');
+      this.initialized = true;
+      console.log(`[MCPClient] Connected and initialized server "${this.config.name}" successfully.`);
+    } catch (err: any) {
+      this.disconnect();
+      throw new Error(`[MCPClient] Failed to initialize server "${this.config.name}": ${err.message}`);
+    }
+  }
+
+  private async connectStdio(): Promise<void> {
+    const cmd = this.config.command;
+    if (!cmd) {
+      throw new Error(`[MCPClient] Command option is required for stdio transport on "${this.config.name}"`);
+    }
+
+    const args = this.config.args || [];
+    const env = { ...process.env, ...(this.config.env || {}) };
+
+    this.childProcess = spawn(cmd, args, { env, shell: true });
+
+    if (!this.childProcess || !this.childProcess.stdout || !this.childProcess.stdin) {
+      throw new Error(`[MCPClient] Failed to spawn stdio child process for "${this.config.name}"`);
+    }
+
+    this.childProcess.stdout.on('data', (data: Buffer) => {
+      this.stdioBuffer += data.toString('utf-8');
+      this.processStdioLines();
+    });
+
+    this.childProcess.stderr?.on('data', (data: Buffer) => {
+      console.warn(`[MCPClient] [stderr] [${this.config.name}]:`, data.toString('utf-8'));
+    });
+
+    this.childProcess.on('close', (code) => {
+      console.log(`[MCPClient] Connection closed for "${this.config.name}" with exit code: ${code}`);
+      this.disconnect();
+    });
+
+    // Wait a brief tick to ensure stdout is set up
+    await new Promise((resolve) => setTimeout(resolve, 100));
+  }
+
+  private processStdioLines() {
+    let newlineIndex: number;
+    while ((newlineIndex = this.stdioBuffer.indexOf('\n')) !== -1) {
+      const line = this.stdioBuffer.substring(0, newlineIndex).trim();
+      this.stdioBuffer = this.stdioBuffer.substring(newlineIndex + 1);
+
+      if (!line) continue;
+
+      try {
+        const payload = JSON.parse(line);
+        if (payload.id !== undefined) {
+          const req = this.pendingRequests.get(Number(payload.id));
+          if (req) {
+            this.pendingRequests.delete(Number(payload.id));
+            if (payload.error) {
+              req.reject(new Error(payload.error.message || JSON.stringify(payload.error)));
+            } else {
+              req.resolve(payload.result);
+            }
+          }
+        }
+      } catch (e) {
+        console.warn(`[MCPClient] Error parsing stdio JSON-RPC line from "${this.config.name}":`, e, 'Line content:', line);
+      }
+    }
+  }
+
+  private async connectSSE(): Promise<void> {
+    if (!this.config.url) {
+      throw new Error(`[MCPClient] URL is required for sse transport on "${this.config.name}"`);
+    }
+
+    // Set up standard SSE handshake / SSE connection
+    // In serverless contexts, we pull initial SSE transport configuration,
+    // and route HTTP POSTs to the client URL.
+    this.sseUrl = this.config.url;
+  }
+
+  private async sendJsonRpc(method: string, params: any): Promise<any> {
+    const id = this.getNextMessageId();
+    const request = {
+      jsonrpc: '2.0',
+      id,
+      method,
+      params,
+    };
+
+    if (this.config.transport === 'stdio') {
+      if (!this.childProcess || !this.childProcess.stdin) {
+        throw new Error(`[MCPClient] Server "${this.config.name}" is not connected (stdio process missing)`);
+      }
+      return new Promise((resolve, reject) => {
+        this.pendingRequests.set(id, { resolve, reject });
+        this.childProcess!.stdin!.write(JSON.stringify(request) + '\n', 'utf-8');
+      });
+    } else {
+      // SSE / HTTP POST fallback for cloud/serverless setups
+      if (!this.sseUrl) {
+        throw new Error(`[MCPClient] Server "${this.config.name}" SSE URL is not initialized`);
+      }
+      const response = await fetch(this.sseUrl, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify(request),
+      });
+      if (!response.ok) {
+        throw new Error(`SSE JSON-RPC POST failed with status ${response.status}`);
+      }
+      const payload = await response.json();
+      if (payload.error) {
+        throw new Error(payload.error.message || JSON.stringify(payload.error));
+      }
+      return payload.result;
+    }
+  }
+
+  private async sendNotification(method: string, params?: any): Promise<void> {
+    const notification = {
+      jsonrpc: '2.0',
+      method,
+      params,
+    };
+
+    if (this.config.transport === 'stdio') {
+      if (this.childProcess?.stdin) {
+        this.childProcess.stdin.write(JSON.stringify(notification) + '\n', 'utf-8');
+      }
+    } else if (this.sseUrl) {
+      await fetch(this.sseUrl, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify(notification),
+      }).catch(err => console.warn('[MCPClient] Failed to send notification over SSE:', err));
+    }
+  }
+
+  /**
+   * List all tools offered by the MCP server.
+   */
+  async listTools(): Promise<MCPTool[]> {
+    await this.connect();
+    const result = await this.sendJsonRpc('tools/list', {});
+    return result?.tools || [];
+  }
+
+  /**
+   * Execute a tool on the MCP server.
+   */
+  async callTool(name: string, args: Record<string, any> = {}): Promise<MCPCallResult> {
+    await this.connect();
+    return await this.sendJsonRpc('tools/call', { name, arguments: args });
+  }
+
+  disconnect() {
+    this.initialized = false;
+    this.pendingRequests.forEach(req => req.reject(new Error('MCPClient disconnected')));
+    this.pendingRequests.clear();
+    
+    if (this.childProcess) {
+      try {
+        this.childProcess.kill();
+      } catch { /* silent */ }
+      this.childProcess = undefined;
+    }
+  }
+}
+
+export class MCPRegistry {
+  private clients: MCPClient[] = [];
+
+  constructor(servers: MCPServerConfig[] = []) {
+    this.clients = servers.map(cfg => new MCPClient(cfg));
+  }
+
+  async getAllTools(): Promise<Array<{ serverName: string; tool: MCPTool }>> {
+    const list: Array<{ serverName: string; tool: MCPTool }> = [];
+    await Promise.all(
+      this.clients.map(async (client) => {
+        try {
+          const tools = await client.listTools();
+          tools.forEach(t => list.push({ serverName: client['config'].name, tool: t }));
+        } catch (e) {
+          console.warn(`[MCPRegistry] Failed to fetch tools from server "${client['config'].name}":`, e);
+        }
+      })
+    );
+    return list;
+  }
+
+  async callTool(toolName: string, args: Record<string, any> = {}): Promise<MCPCallResult> {
+    // Find client hosting this tool. (In standard MCP, tool names are unique per server block)
+    // We will search across all connected clients.
+    for (const client of this.clients) {
+      try {
+        const tools = await client.listTools();
+        if (tools.some(t => t.name === toolName)) {
+          return await client.callTool(toolName, args);
+        }
+      } catch { /* silent fallback */ }
+    }
+    throw new Error(`[MCPRegistry] Tool "${toolName}" not found on any registered MCP server.`);
+  }
+
+  disconnectAll() {
+    this.clients.forEach(c => c.disconnect());
+  }
+}