@retrivora-ai/rag-engine 0.1.5 → 0.1.7

package/src/core/Pipeline.ts

@@ -2,14 +2,21 @@ import { BaseVectorProvider } from '../providers/vectordb/BaseVectorProvider';
 import { ILLMProvider, ChatMessage } from '../llm/ILLMProvider';
 import { DocumentChunker } from '../rag/DocumentChunker';
 import { RagConfig } from '../config/RagConfig';
-import { VectorMatch } from '../types';
 import { ProviderRegistry } from './ProviderRegistry';
+import { BatchProcessor, BatchOptions } from './BatchProcessor';
 
 // Moved to src/types/index.ts
 import { IngestDocument, ChatResponse } from '../types';
 
 /**
  * Pipeline — orchestrates the RAG flow: Embed → Search → Augment → Generate.
+ *
+ * Features:
+ * - Lazy initialization of providers
+ * - Batch processing with retry logic
+ * - Configurable embedding strategies
+ * - Error recovery for transient failures
+ * - Support for multi-tenancy via namespacing
  */
 export class Pipeline {
   private vectorDB!: BaseVectorProvider;
@@ -34,7 +41,7 @@ export class Pipeline {
     this.llmProvider = ProviderRegistry.createLLMProvider(this.config.llm, this.config.embedding);
 
     if (this.config.llm.provider === 'anthropic') {
-      // Logic for Anthropic as seen in RAGPipeline
+      // Anthropic doesn't support embeddings, so use separate provider
       const { LLMFactory } = await import('../llm/LLMFactory');
       this.embeddingProvider = LLMFactory.createEmbeddingProvider(this.config.embedding);
     } else {
@@ -45,28 +52,76 @@ export class Pipeline {
     this.initialised = true;
   }
 
-  async ingest(documents: IngestDocument[], namespace?: string): Promise<Array<{ docId: string; chunksIngested: number }>> {
+  /**
+   * Ingest documents with automatic chunking, embedding, and batch processing.
+   * Handles retries for transient failures.
+   */
+  async ingest(documents: IngestDocument[], namespace?: string): Promise<Array<{ docId: string | number; chunksIngested: number }>> {
     await this.initialize();
     const ns = namespace ?? this.config.projectId;
     const results = [];
 
     for (const doc of documents) {
-      const chunks = this.chunker.chunk(doc.content, {
-        docId: doc.docId,
-        metadata: doc.metadata,
-      });
-
-      const vectors = await this.embeddingProvider.batchEmbed(chunks.map(c => c.content));
-      const upsertDocs = chunks.map((chunk, i) => ({
-        id: chunk.id,
-        vector: vectors[i],
-        content: chunk.content,
-        metadata: chunk.metadata,
-      }));
-
-      await this.vectorDB.batchUpsert(upsertDocs, ns);
-      results.push({ docId: doc.docId, chunksIngested: chunks.length });
+      try {
+        const chunks = this.chunker.chunk(doc.content, {
+          docId: doc.docId,
+          metadata: doc.metadata,
+        });
+
+        // Batch embed chunks with retry logic
+        const batchOptions: BatchOptions = {
+          batchSize: 50, // Embedding batch size
+          maxRetries: 3,
+          initialDelayMs: 100,
+        };
+
+        const vectors = await BatchProcessor.mapWithRetry(
+          chunks.map(c => c.content),
+          (text) => this.embeddingProvider.embed(text),
+          batchOptions
+        );
+
+        if (vectors.length !== chunks.length) {
+          throw new Error(`Embedding failed: got ${vectors.length} vectors for ${chunks.length} chunks`);
+        }
+
+        const upsertDocs = chunks.map((chunk, i) => ({
+          id: chunk.id,
+          vector: vectors[i],
+          content: chunk.content,
+          metadata: chunk.metadata,
+        }));
+
+        // Batch upsert with retry logic
+        const upsertBatchOptions: BatchOptions = {
+          batchSize: 100,
+          maxRetries: 3,
+          initialDelayMs: 100,
+        };
+
+        const upsertResult = await BatchProcessor.processBatch(
+          upsertDocs,
+          (batch) => this.vectorDB.batchUpsert(batch, ns),
+          upsertBatchOptions
+        );
+
+        if (upsertResult.errors.length > 0) {
+          console.warn(`[Pipeline] Ingestion: ${upsertResult.errors.length} batches failed for doc ${doc.docId}`);
+        }
+
+        results.push({
+          docId: doc.docId,
+          chunksIngested: upsertResult.totalProcessed,
+        });
+      } catch (error) {
+        console.error(`[Pipeline] Failed to ingest document ${doc.docId}:`, error);
+        results.push({
+          docId: doc.docId,
+          chunksIngested: 0,
+        });
+      }
     }
+
     return results;
   }
 
@@ -76,26 +131,21 @@ export class Pipeline {
     const topK = this.config.rag?.topK ?? 5;
     const scoreThreshold = this.config.rag?.scoreThreshold ?? 0.0;
 
-    const queryVector = await this.embeddingProvider.embed(question);
-    const rawMatches = await this.vectorDB.query(queryVector, topK, ns);
-    
-    const sources = rawMatches.filter((m) => m.score >= scoreThreshold);
-    const context = sources.length
-      ? sources.map((m, i) => `[Source ${i + 1}]\n${m.content}`).join('\n\n---\n\n')
-      : 'No relevant context found.';
-
-    const messages: ChatMessage[] = [...history, { role: 'user', content: question }];
-    const reply = await this.llmProvider.chat(messages, context);
-
-    return { reply, sources };
-  }
-
-  async health(): Promise<Record<string, boolean>> {
-    await this.initialize();
-    const [vectorDB, llm] = await Promise.all([
-      this.vectorDB.ping().catch(() => false),
-      this.llmProvider.ping().catch(() => false),
-    ]);
-    return { vectorDB, llm };
+    try {
+      const queryVector = await this.embeddingProvider.embed(question);
+      const rawMatches = await this.vectorDB.query(queryVector, topK, ns);
+      
+      const sources = rawMatches.filter((m) => m.score >= scoreThreshold);
+      const context = sources.length
+        ? sources.map((m, i) => `[Source ${i + 1}]\n${m.content}`).join('\n\n---\n\n')
+        : 'No relevant context found.';
+
+      const messages: ChatMessage[] = [...history, { role: 'user', content: question }];
+      const reply = await this.llmProvider.chat(messages, context);
+
+      return { reply, sources };
+    } catch (error) {
+      throw new Error(`[Pipeline] Chat failed: ${error instanceof Error ? error.message : String(error)}`);
+    }
   }
 }