@retrivora-ai/rag-engine 1.9.6 → 1.9.8

package/src/core/MultiAgentCoordinator.ts

@@ -0,0 +1,239 @@
+import { ILLMProvider } from '../llm/ILLMProvider';
+import { MCPRegistry } from './mcp';
+import { ChatMessage, ChatResponse } from '../types';
+
+interface CoordinatorOptions {
+  llmProvider: ILLMProvider;
+  mcpRegistry: MCPRegistry;
+  documentSearch: (query: string) => Promise<{ reply: string; sources: any[] }>;
+  maxIterations?: number;
+}
+
+export class MultiAgentCoordinator {
+  private llmProvider: ILLMProvider;
+  private mcpRegistry: MCPRegistry;
+  private documentSearch: (query: string) => Promise<{ reply: string; sources: any[] }>;
+  private maxIterations: number;
+
+  constructor(options: CoordinatorOptions) {
+    this.llmProvider = options.llmProvider;
+    this.mcpRegistry = options.mcpRegistry;
+    this.documentSearch = options.documentSearch;
+    this.maxIterations = options.maxIterations ?? 5;
+  }
+
+  /**
+   * Run the multi-agent coordination loop synchronously and return the final response.
+   */
+  async run(question: string, history: ChatMessage[] = []): Promise<ChatResponse> {
+    const generator = this.runStream(question, history);
+    let finalReply = '';
+    const sources: any[] = [];
+
+    for await (const chunk of generator) {
+      if (typeof chunk === 'string') {
+        finalReply += chunk;
+      } else if (chunk && typeof chunk === 'object') {
+        if ('reply' in chunk && chunk.reply) {
+          finalReply += chunk.reply;
+        }
+        if ('sources' in chunk && chunk.sources) {
+          sources.push(...chunk.sources);
+        }
+      }
+    }
+
+    return {
+      reply: finalReply,
+      sources,
+    };
+  }
+
+  /**
+   * Run the multi-agent coordination loop as an async stream, yielding intermediate thought blocks and tool execution events.
+   */
+  async *runStream(question: string, history: ChatMessage[] = []): AsyncGenerator<string | ChatResponse> {
+    // 1. Gather all tools
+    const mcpTools = await this.mcpRegistry.getAllTools();
+    
+    // Construct system instructions with tools
+    let systemPrompt = `You are a goal-driven supervisor agent coordinating a RAG search engine and external MCP servers to solve the user's task.
+You must think step-by-step before acting. Write your internal thinking process inside a \`<think>...</think>\` block first, then act or answer.
+
+Available Tools:
+- document_search(query: string): Search the vector database and knowledge base for documents.
+
+${mcpTools.length > 0 ? 'MCP Server Tools:' : ''}
+${mcpTools.map(mt => `- ${mt.tool.name}(${JSON.stringify(mt.tool.inputSchema?.properties || {})}): ${mt.tool.description || 'No description provided.'}`).join('\n')}
+
+Tool Calling Protocol:
+If you need to call a tool, write:
+TOOL_CALL: <toolName>({ "argName": "value" })
+And then STOP generating immediately.
+
+Once you have gathered enough information to fully satisfy the user's query, write your final response directly after the closing \`</think>\` block.
+`;
+
+    // 2. Setup message history for the supervisor
+    const supervisorHistory: ChatMessage[] = [
+      { role: 'system' as any, content: systemPrompt },
+      ...history,
+      { role: 'user', content: question }
+    ];
+
+    let iteration = 0;
+    const allSources: any[] = [];
+
+    while (iteration < this.maxIterations) {
+      iteration++;
+      
+      const context = ''; // Context is populated by tool call results directly in message history
+      let fullModelResponse = '';
+      let textBuffer = '';
+      let inThinking = false;
+      let toolCallDetected: { name: string; args: any } | null = null;
+
+      if (!this.llmProvider.chatStream) {
+        // Fallback for non-streaming providers
+        const reply = await this.llmProvider.chat(supervisorHistory, context);
+        fullModelResponse = reply;
+        
+        // Parse simulated thinking / tool call
+        const thinkIndex = reply.indexOf('<think>');
+        const endThinkIndex = reply.indexOf('</think>');
+        
+        let thinkingText = '';
+        let answerText = reply;
+
+        if (thinkIndex !== -1 && endThinkIndex !== -1) {
+          thinkingText = reply.substring(thinkIndex + 7, endThinkIndex);
+          answerText = reply.substring(endThinkIndex + 8);
+          yield { type: 'thinking', text: thinkingText } as any;
+        }
+
+        const toolMatch = answerText.match(/TOOL_CALL:\s*(\w+)\s*\(([\s\S]*)\)/);
+        if (toolMatch) {
+          const toolName = toolMatch[1];
+          const toolArgsStr = toolMatch[2];
+          try {
+            const toolArgs = JSON.parse(toolArgsStr.trim());
+            toolCallDetected = { name: toolName, args: toolArgs };
+          } catch (e) {
+            console.error('[MultiAgentCoordinator] Failed to parse tool args:', toolArgsStr, e);
+          }
+        } else {
+          yield answerText;
+        }
+      } else {
+        // Run with streaming completion
+        const stream = this.llmProvider.chatStream(supervisorHistory, context);
+
+        for await (const chunk of stream) {
+          fullModelResponse += chunk;
+          textBuffer += chunk;
+
+          // Check for <think> tag
+          if (!inThinking && textBuffer.includes('<think>')) {
+            inThinking = true;
+            const idx = textBuffer.indexOf('<think>');
+            textBuffer = textBuffer.slice(idx + 7);
+          }
+
+          // Check for </think> tag
+          if (inThinking && textBuffer.includes('</think>')) {
+            inThinking = false;
+            const idx = textBuffer.indexOf('</think>');
+            const thinkingDelta = textBuffer.slice(0, idx);
+            yield { type: 'thinking', text: thinkingDelta } as any;
+            textBuffer = textBuffer.slice(idx + 8);
+          }
+
+          // If we are in thinking mode, yield thinking delta
+          if (inThinking && textBuffer.length > 0) {
+            yield { type: 'thinking', text: textBuffer } as any;
+            textBuffer = '';
+          }
+
+          // Check for TOOL_CALL during stream
+          if (!inThinking && textBuffer.includes('TOOL_CALL:')) {
+            const idx = textBuffer.indexOf('TOOL_CALL:');
+            const precedingText = textBuffer.slice(0, idx);
+            if (precedingText.trim()) {
+              yield precedingText;
+            }
+            // Await remainder of stream to get the full tool call signature
+            textBuffer = textBuffer.slice(idx);
+          }
+
+          // If not in thinking and no tool call prefix is blocking, yield final answer text
+          if (!inThinking && !textBuffer.includes('TOOL_CALL:') && textBuffer.length > 0) {
+            yield textBuffer;
+            textBuffer = '';
+          }
+        }
+
+        // Flush any remaining buffer text
+        if (textBuffer.trim()) {
+          const toolMatch = textBuffer.match(/TOOL_CALL:\s*(\w+)\s*\(([\s\S]*)\)/);
+          if (toolMatch) {
+            const toolName = toolMatch[1];
+            const toolArgsStr = toolMatch[2];
+            try {
+              const toolArgs = JSON.parse(toolArgsStr.trim());
+              toolCallDetected = { name: toolName, args: toolArgs };
+            } catch (e) {
+              console.error('[MultiAgentCoordinator] Failed to parse tool args:', toolArgsStr, e);
+            }
+          } else {
+            yield textBuffer;
+          }
+        }
+      }
+
+      // 3. Append the model response to the supervisor history
+      supervisorHistory.push({ role: 'assistant', content: fullModelResponse });
+
+      // 4. Handle tool execution if requested, else break the loop
+      if (toolCallDetected) {
+        const { name: toolName, args: toolArgs } = toolCallDetected;
+        
+        yield { 
+          type: 'thinking', 
+          text: `\n[Agent Call] Executing tool "${toolName}" with parameters ${JSON.stringify(toolArgs)}...\n` 
+        } as any;
+
+        let toolResultText = '';
+        try {
+          if (toolName === 'document_search') {
+            const searchRes = await this.documentSearch(toolArgs.query || '');
+            toolResultText = `document_search returned:\n${searchRes.reply}\n\nSources count: ${searchRes.sources.length}`;
+            if (searchRes.sources && searchRes.sources.length > 0) {
+              allSources.push(...searchRes.sources);
+              yield { reply: '', sources: searchRes.sources };
+            }
+          } else {
+            // Call MCP tool
+            const mcpRes = await this.mcpRegistry.callTool(toolName, toolArgs);
+            toolResultText = `MCP Tool "${toolName}" returned:\n${JSON.stringify(mcpRes.content)}`;
+          }
+        } catch (err: any) {
+          toolResultText = `Error calling tool "${toolName}": ${err.message}`;
+        }
+
+        // Append tool output to history
+        supervisorHistory.push({
+          role: 'user',
+          content: `TOOL_RESULT for ${toolName}:\n${toolResultText}`
+        });
+
+        yield {
+          type: 'thinking',
+          text: `[Agent Output] Tool "${toolName}" executed. Continuing reasoning...\n`
+        } as any;
+      } else {
+        // No tool call requested — response is complete
+        break;
+      }
+    }
+  }
+}

package/src/core/Pipeline.ts

@@ -2,12 +2,15 @@ import { BaseVectorProvider } from '../providers/vectordb/BaseVectorProvider';
 import { BaseGraphProvider } from '../providers/graphdb/BaseGraphProvider';
 import { ILLMProvider } from '../llm/ILLMProvider';
 import { ChatMessage } from '../types';
+import { EmbeddingFailedException } from '../exceptions';
 import { DocumentChunker, Chunk } from '../rag/DocumentChunker';
 import { EntityExtractor } from '../rag/EntityExtractor';
 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';
@@ -164,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(
@@ -249,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.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.
    */
@@ -327,7 +372,7 @@ You are a helpful assistant. Use the provided context to answer questions accura
     );
 
     if (vectors.length !== chunks.length) {
-      throw new Error(
+      throw new EmbeddingFailedException(
         `[Pipeline] Embedding mismatch: got ${vectors.length} vectors for ${chunks.length} chunks. ` +
         `Check embedding provider logs for individual chunk failures.`
       );
@@ -385,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.mcpServers && this.config.mcpServers.length > 0)) && this.multiAgentCoordinator) {
+      return await this.multiAgentCoordinator.run(question, history);
     }
 
     const stream = this.askStream(question, history, namespace);
@@ -414,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.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.
@@ -424,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;
@@ -546,6 +620,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
@@ -750,15 +846,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,
@@ -774,12 +870,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/ProviderRegistry.ts

@@ -10,7 +10,7 @@ import { ProviderNotFoundException } from '../exceptions';
  * ProviderRegistry — dynamic provider loader for Vector DBs and LLMs.
  */
 type VectorProviderClass = {
-  new (config: VectorDBConfig): BaseVectorProvider;
+  new(config: VectorDBConfig): BaseVectorProvider;
   getValidator?: () => IProviderValidator;
   getHealthChecker?: () => IProviderHealthChecker;
 };
@@ -20,10 +20,10 @@ type GraphProviderClass = new (config: GraphDBConfig) => BaseGraphProvider;
 export class ProviderRegistry {
   private static vectorProviders: Record<string, VectorProviderClass> = {};
   private static graphProviders: Record<string, GraphProviderClass> = {};
-  
+
   private static vectorValidators: Record<string, IProviderValidator> = {};
   private static vectorHealthCheckers: Record<string, IProviderHealthChecker> = {};
-  
+
   private static llmValidators: Record<string, IProviderValidator> = {};
   private static llmHealthCheckers: Record<string, IProviderHealthChecker> = {};
 
@@ -39,7 +39,7 @@ export class ProviderRegistry {
 
   static async getVectorValidator(provider: string): Promise<IProviderValidator | null> {
     if (this.vectorValidators[provider]) return this.vectorValidators[provider];
-    
+
     // Try to load built-in validator
     try {
       const providerClass = await this.loadVectorProviderClass(provider);
@@ -55,7 +55,7 @@ export class ProviderRegistry {
 
   static async getVectorHealthChecker(provider: string): Promise<IProviderHealthChecker | null> {
     if (this.vectorHealthCheckers[provider]) return this.vectorHealthCheckers[provider];
-    
+
     try {
       const providerClass = await this.loadVectorProviderClass(provider);
       if (providerClass.getHealthChecker) {

package/src/core/Retrivora.ts

@@ -2,7 +2,9 @@ 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';
 
 /**
  * Public SDK facade matching the prompt-level Retrivora API.
@@ -20,27 +22,70 @@ export class Retrivora {
   }
 
   async initialize(): Promise<void> {
-    await ConfigValidator.validateAndThrow(this.config);
-    await this.pipeline.initialize();
+    try {
+      await ConfigValidator.validateAndThrow(this.config);
+      LicenseVerifier.verify(
+        this.config.licenseKey,
+        this.config.projectId,
+        this.config.vectorDb?.provider
+      );
+    } catch (err) {
+      throw wrapError(err, 'CONFIGURATION_ERROR');
+    }
+    try {
+      await this.pipeline.initialize();
+    } catch (err) {
+      throw wrapError(err, 'AUTHENTICATION_ERROR');
+    }
   }
 
   async ingest(
     documents: IngestDocument[],
     namespace?: string,
   ): Promise<Array<{ docId: string | number; chunksIngested: number }>> {
-    return this.pipeline.ingest(documents, namespace);
+    try {
+      return await this.pipeline.ingest(documents, namespace);
+    } catch (err) {
+      const msg = String(err);
+      let defaultCode: RetrivoraErrorCode = 'RETRIEVAL_FAILED';
+      if (msg.includes('Embed') || msg.includes('embed')) {
+        defaultCode = 'EMBEDDING_FAILED';
+      }
+      throw wrapError(err, defaultCode);
+    }
   }
 
   async ask(question: string, history: ChatMessage[] = [], namespace?: string): Promise<ChatResponse> {
-    return this.pipeline.ask(question, history, namespace);
+    try {
+      return await this.pipeline.ask(question, history, namespace);
+    } catch (err) {
+      const msg = String(err);
+      let defaultCode: RetrivoraErrorCode = 'RETRIEVAL_FAILED';
+      if (msg.includes('Embed') || msg.includes('embed')) {
+        defaultCode = 'EMBEDDING_FAILED';
+      }
+      throw wrapError(err, defaultCode);
+    }
   }
 
-  askStream(
+  async *askStream(
     question: string,
     history: ChatMessage[] = [],
     namespace?: string,
   ): AsyncIterable<string | ChatResponse> {
-    return this.pipeline.askStream(question, history, namespace);
+    try {
+      const stream = this.pipeline.askStream(question, history, namespace);
+      for await (const chunk of stream) {
+        yield chunk;
+      }
+    } catch (err) {
+      const msg = String(err);
+      let defaultCode: RetrivoraErrorCode = 'RETRIEVAL_FAILED';
+      if (msg.includes('Embed') || msg.includes('embed')) {
+        defaultCode = 'EMBEDDING_FAILED';
+      }
+      throw wrapError(err, defaultCode);
+    }
   }
 
   getPipeline(): Pipeline {
@@ -49,3 +94,4 @@ export class Retrivora {
 }
 
 export type { RetrievalConfig, UniversalRagConfig };
+