@retrivora-ai/rag-engine 1.8.0 → 1.8.2

package/src/core/Pipeline.ts

@@ -12,8 +12,13 @@ import { ProviderRegistry } from './ProviderRegistry';
 import { BatchProcessor, BatchOptions } from './BatchProcessor';
 import { EmbeddingStrategyResolver } from '../config/EmbeddingStrategy';
 import { QueryProcessor } from './QueryProcessor';
-import { IngestDocument, ChatResponse, VectorMatch, GraphSearchResult, RetrievalResult, UpsertDocument } from '../types';
+import {
+  IngestDocument, ChatResponse, VectorMatch, GraphSearchResult,
+  RetrievalResult, UpsertDocument, ObservabilityTrace, LatencyBreakdown,
+  TokenUsage, RetrievedChunk,
+} from '../types';
 import { UITransformer } from '../utils/UITransformer';
+import { SchemaMapper, SchemaMap } from '../utils/SchemaMapper';
 
 // ─── LRU Embedding Cache ───────────────────────────────────────────────────────
 
@@ -59,6 +64,89 @@ class LRUEmbeddingCache {
   }
 }
 
+// ─── Token estimation helpers ──────────────────────────────────────────────────
+
+/** Rough token count estimate when the provider doesn't expose usage data. */
+function estimateTokens(text: string): number {
+  return Math.ceil(text.length / 4);
+}
+
+/**
+ * Model pricing table (cost per 1k tokens in USD).
+ * Used for estimating cost when providers don't return billing data.
+ */
+const MODEL_COST_PER_1K: Record<string, { input: number; output: number }> = {
+  'gpt-4o':               { input: 0.0025, output: 0.010  },
+  'gpt-4o-mini':          { input: 0.00015, output: 0.0006 },
+  'gpt-4-turbo':          { input: 0.01,   output: 0.03   },
+  'gpt-3.5-turbo':        { input: 0.0005, output: 0.0015 },
+  'claude-3-5-sonnet':    { input: 0.003,  output: 0.015  },
+  'claude-3-haiku':       { input: 0.00025, output: 0.00125 },
+  'gemini-1.5-flash':     { input: 0.000075, output: 0.0003 },
+  'gemini-1.5-pro':       { input: 0.00125, output: 0.005 },
+};
+
+function estimateCostUsd(
+  promptTokens: number,
+  completionTokens: number,
+  model?: string,
+): number | undefined {
+  if (!model) return undefined;
+  const key = Object.keys(MODEL_COST_PER_1K).find((k) => model.toLowerCase().includes(k));
+  if (!key) return undefined;
+  const prices = MODEL_COST_PER_1K[key];
+  return (promptTokens / 1000) * prices.input + (completionTokens / 1000) * prices.output;
+}
+
+// ─── Hallucination scorer ──────────────────────────────────────────────────────
+
+/**
+ * Ask the LLM to self-critique: rate how well the answer is supported by the context.
+ * Returns a score 0–1 and a brief reason. Silently returns undefined on failure.
+ */
+async function scoreHallucination(
+  llm: ILLMProvider,
+  answer: string,
+  context: string,
+): Promise<{ score: number; reason: string } | undefined> {
+  const maxContextChars = 3000;
+  const truncatedContext = context.length > maxContextChars
+    ? context.slice(0, maxContextChars) + '\n...[truncated]'
+    : context;
+
+  const prompt = `You are an AI quality checker. Given the CONTEXT retrieved from a knowledge base and the ANSWER generated from it, rate how well-grounded the answer is.
+
+CONTEXT:
+${truncatedContext}
+
+ANSWER:
+${answer}
+
+Return ONLY a valid JSON object with no markdown fences:
+{"score": <float 0-1>, "reason": "<one sentence>"}
+
+Where score 0 = fully grounded in context, score 1 = likely hallucinated or unsupported.`;
+
+  try {
+    const raw = await llm.chat(
+      [{ role: 'user', content: prompt }],
+      '',
+      { temperature: 0, maxTokens: 120 },
+    );
+
+    const jsonMatch = raw.match(/\{[\s\S]*\}/);
+    if (!jsonMatch) return undefined;
+
+    const parsed = JSON.parse(jsonMatch[0]);
+    if (typeof parsed.score === 'number' && typeof parsed.reason === 'string') {
+      return { score: Math.min(1, Math.max(0, parsed.score)), reason: parsed.reason };
+    }
+  } catch {
+    // Non-critical — fail silently
+  }
+  return undefined;
+}
+
 // ─── Pipeline ─────────────────────────────────────────────────────────────────
 
 /**
@@ -71,6 +159,7 @@ class LRUEmbeddingCache {
  * - Error recovery for transient failures
  * - Multi-tenancy support via namespacing
  * - LRU-bounded embedding cache (max 500 entries, prevents memory leaks)
+ * - Full observability tracing (latency, tokens, hallucination scoring)
  */
 export class Pipeline {
   private vectorDB!: BaseVectorProvider;
@@ -97,215 +186,36 @@ export class Pipeline {
     this.reranker = new Reranker();
   }
 
+  /**
+   * Expose the underlying LLM provider (set after initialize()).
+   * Used by the stream handler to pass to UITransformer.analyzeAndDecide().
+   */
+  getLLMProvider(): ILLMProvider | undefined {
+    return this.initialised ? this.llmProvider : undefined;
+  }
+
   async initialize(): Promise<void> {
     if (this.initialised) return;
 
-    // Augment system prompt with a Universal Visualization Protocol
-    // We use a unique marker to ensure this is only injected once but is ALWAYS present.
-    const CHART_MARKER = '<!-- UI_PROTOCOL_V7 -->';
-    const chartInstruction = `
-
-${CHART_MARKER}
-### UNIVERSAL UI PROTOCOL V7 (STRICT MODE)
-
-You are responsible for returning a SINGLE structured UI block when visualization is required.
-
----
-
-## 1. INTENT CLASSIFICATION (MANDATORY)
-
-Classify the query into ONE of the following intents:
-
-- "explore" → user wants to browse items (e.g., "show me products")
-- "analyze" → user wants aggregation/distribution (e.g., "distribution of products")
-- "compare" → user wants side-by-side comparison
-- "detail" → user wants structured/tabular data
-
----
-
-## 2. VIEW SELECTION (DYNAMIC)
+    const chartInstruction = `\
+You are a helpful product assistant. Use the provided context to answer questions accurately.
 
-Choose the best primary "view" ("carousel", "chart", or "table") based on the data:
-- TRENDS/GROWTH → "chart" (type: "line")
-- DISTRIBUTIONS/RATIOS → "chart" (type: "pie")
-- COMPARISONS → "chart" (type: "bar")
-- PRODUCT LISTS → "carousel"
-- RAW DATA/SPECS → "table"
-
-⚠️ Dynamic Overrides:
-- If query is "show products", ALWAYS use carousel.
-- If data is statistical, ALWAYS use a chart.
-- NEVER explain why you chose a view.
-
----
-
-## 3. RESPONSE FORMAT (STRICT JSON ONLY INSIDE BLOCK)
-
-\`\`\`ui
-{
-  "view": "carousel" | "chart" | "table",
-  "title": "string",
-  "description": "string",
-  "metadata": {
-    "intent": "explore" | "analyze" | "compare" | "detail",
-    "confidence": 0.0-1.0
-  },
-
-  // CHART ONLY
-  "chart": {
-    "type": "pie" | "bar" | "line",
-    "xKey": "label",
-    "yKeys": ["value"],
-    "data": []
-  },
-
-  // TABLE ONLY
-  "table": {
-    "columns": [],
-    "rows": []
-  },
-
-  // CAROUSEL ONLY
-  "carousel": {
-    "items": []
-  },
-
-  "insights": []
-}
-\`\`\`
+### UI STYLE RULES (CRITICAL):
+- NEVER generate markdown tables. If you do, the UI will break.
+- NEVER generate HTML tags like <figure>, <tbody>, <tr>, etc.
+- NEVER generate text-based charts or graphs.
+- ONLY use plain text and bullet points.
 
----
-
-## 4. DATA CONTRACT (CRITICAL)
-
-### 🔹 CAROUSEL ITEM FORMAT
-{
-  "id": "string",
-  "name": "string",
-  "brand": "string",
-  "price": number,
-  "image": "url",
-  "link": "url",
-  "inStock": boolean
-}
-
-### 🔹 CHART DATA FORMAT
-{
-  "label": "string",
-  "value": number,
-  "inStockCount": number (optional),
-  "outOfStockCount": number (optional)
-}
-
-### 🔹 TABLE FORMAT
-- columns MUST match row keys
-- rows MUST be flat objects
-
----
-
-## 5. DATA RULES
-
-- NEVER hallucinate fields or data.
-- ONLY use data provided in the retrieved context.
-- STRICT DATA ISOLATION: The labels and numbers in the EXAMPLES section below are placeholders. NEVER use them.
-- IF NO DATA IS FOUND: You MUST NOT generate a chart. You MUST NOT say "This example assumes" or provide hypothetical numbers. Simply state "No relevant data found in the database." and STOP.
-- NO META-COMMENTARY: NEVER talk about the task itself (e.g., "I've used JSON", "In a real-world app", "Dedicated library"). 
-- NEVER explain YOUR formatting choice.
-- Aggregate BEFORE rendering chart.
-- DYNAMIC INTELLIGENCE: Autonomously extract relevant categories/values from the context and calculate totals for the "data" array based on the user query.
-
----
-
-## 6. SMART RULES (IMPORTANT)
-
-- If user asks BOTH:
-  "show products AND distribution"
-  → PRIORITIZE "explore" → carousel
-
-- If dataset size > 20:
-  → aggregate → chart or table
-
----
-
-## 7. OUTPUT RULES
-
-- ALWAYS return EXACTLY ONE \`\`\`ui\`\`\` block.
-- JSON must be VALID.
-- FIELD MAPPING (MANDATORY): Always use "label" for the name/category and "value" for the numeric count.
-- NO xKey or yKeys: Do not include these fields; the UI uses "label" and "value" by default.
-- SILENT DATA: If a chart is provided, DO NOT say "Here is the JSON" or "Below is the data". Simply output the block.
-- TEXT OUTSIDE THE BLOCK: One brief sentence maximum.
-
----
-
-## 8. EXAMPLES
-
-User: "show me beauty products"
-
-\`\`\`ui
-{
-  "view": "carousel",
-  "title": "TITLE",
-  "description": "DESCRIPTION",
-  "metadata": {
-    "intent": "explore",
-    "confidence": 0.95
-  },
-  "carousel": {
-    "items": [
-      {
-        "id": "ID",
-        "name": "PRODUCT_NAME",
-        "brand": "BRAND",
-        "price": 0,
-        "image": "URL",
-        "link": "URL",
-        "inStock": true
-      }
-    ]
-  },
-  "insights": ["INSIGHT"]
-}
-\`\`\`
-
----
-
-
-
-\`\`\`ui
-{
-  "view": "chart",
-  "title": "TITLE",
-  "description": "DESCRIPTION",
-  "metadata": {
-    "intent": "analyze",
-    "confidence": 0.98
-  },
-  "chart": {
-    "type": "pie",
-    "data": [
-      { "label": "LABEL_A", "value": 0, "inStockCount": 0, "outOfStockCount": 0 },
-      { "label": "LABEL_B", "value": 0, "inStockCount": 0, "outOfStockCount": 0 }
-    ]
-  },
-  "insights": ["INSIGHT"]
-}
-\`\`\`
+### PRODUCT DISPLAY:
+- When recommending products, simply list their names, prices, and features in a friendly, conversational manner.
+- The UI will automatically detect these products and show high-quality product cards in a carousel below your message.
+- Do NOT try to format product lists as tables.
 `;
-    if (!this.config.llm.systemPrompt?.includes(CHART_MARKER)) {
-      // Clean up ANY existing protocol markers (V1-V7) to avoid prompt bloat
-      let cleanPrompt = this.config.llm.systemPrompt || '';
-      cleanPrompt = cleanPrompt.replace(/\n\n### UNIVERSAL UI PROTOCOL[\s\S]*?(?=\n\n|$)/g, '');
-      cleanPrompt = cleanPrompt.replace(/<!-- UI_PROTOCOL_V\d+ -->[\s\S]*?(?=\n\n|$)/g, '');
-      
-      this.config.llm.systemPrompt = cleanPrompt.trim() + chartInstruction;
-    }
 
-    console.log(`[Pipeline] Initializing with provider: ${this.config.vectorDb.provider}...`);
-    // Resolve vector DB provider
+    this.config.llm.systemPrompt = chartInstruction;
+
     this.vectorDB = await ProviderRegistry.createVectorProvider(this.config.vectorDb);
 
-    // Resolve LLM + embedding providers
     const { llmProvider, embeddingProvider } = await EmbeddingStrategyResolver.resolve(
       this.config.llm,
       this.config.embedding
@@ -322,7 +232,6 @@ User: "show me beauty products"
 
     await this.vectorDB.initialize();
 
-    // Initialize Agentic Layer if configured
     if (this.config.rag?.architecture === 'agentic') {
       this.agent = new LangChainAgent(this, this.config);
       await this.agent.initialize(this.llmProvider as unknown);
@@ -333,7 +242,6 @@ User: "show me beauty products"
 
   /**
    * Ingest documents with automatic chunking, embedding, and batch upsert.
-   * Handles retries for transient failures.
    */
   async ingest(
     documents: IngestDocument[],
@@ -357,10 +265,7 @@ User: "show me beauty products"
 
         const totalProcessed = await this.processUpserts(upsertDocs, ns);
 
-        results.push({
-          docId: doc.docId,
-          chunksIngested: totalProcessed,
-        });
+        results.push({ docId: doc.docId, chunksIngested: totalProcessed });
 
         if (this.graphDB && this.entityExtractor) {
           await this.processGraphIngestion(doc.docId, chunks);
@@ -374,23 +279,17 @@ User: "show me beauty products"
     return results;
   }
 
-  /**
-   * Step 1: Chunk the document content.
-   */
+  /** Step 1: Chunk the document content. */
   private async prepareChunks(doc: IngestDocument): Promise<Chunk[]> {
     if (this.llamaIngestor) {
-      return await this.llamaIngestor.chunk(doc.content, {
+      return this.llamaIngestor.chunk(doc.content, {
         docId: doc.docId,
         metadata: doc.metadata,
         chunkSize: this.config.rag?.chunkSize,
         chunkOverlap: this.config.rag?.chunkOverlap,
       });
     }
-
-    return this.chunker.chunk(doc.content, {
-      docId: doc.docId,
-      metadata: doc.metadata,
-    });
+    return this.chunker.chunk(doc.content, { docId: doc.docId, metadata: doc.metadata });
   }
 
   /**
@@ -410,7 +309,6 @@ User: "show me beauty products"
       embedBatchOptions
     );
 
-    // mapWithRetry filters out undefined on failure; guard against mismatch explicitly
     if (vectors.length !== chunks.length) {
       throw new Error(
         `[Pipeline] Embedding mismatch: got ${vectors.length} vectors for ${chunks.length} chunks. ` +
@@ -421,9 +319,7 @@ User: "show me beauty products"
     return vectors;
   }
 
-  /**
-   * Step 3: Upsert chunks to vector database with retry logic.
-   */
+  /** Step 3: Upsert chunks to vector database with retry logic. */
   private async processUpserts(upsertDocs: UpsertDocument[], namespace: string): Promise<number> {
     const upsertBatchOptions: BatchOptions = {
       batchSize: 100,
@@ -444,14 +340,10 @@ User: "show me beauty products"
     return upsertResult.totalProcessed;
   }
 
-  /**
-   * Step 4: Optional graph-based entity extraction and ingestion.
-   */
+  /** Step 4: Optional graph-based entity extraction and ingestion. */
   private async processGraphIngestion(docId: string | number, chunks: Chunk[]): Promise<void> {
-    console.log(`[Pipeline] Extracting entities for doc ${docId} (${chunks.length} chunks)...`);
-
     const extractionOptions: BatchOptions = {
-      batchSize: 2, // Low concurrency for LLM extraction
+      batchSize: 2,
       maxRetries: 1,
       initialDelayMs: 500,
     };
@@ -465,7 +357,7 @@ User: "show me beauty products"
             if (nodes.length > 0) await this.graphDB!.addNodes(nodes);
             if (edges.length > 0) await this.graphDB!.addEdges(edges);
           } catch (err) {
-            console.warn(`[Pipeline] Entity extraction failed for chunk:`, err);
+            console.warn(`[Pipeline] Entity extraction failed for chunk in doc ${docId}:`, err);
           }
         }
       },
@@ -477,7 +369,6 @@ User: "show me beauty products"
     await this.initialize();
 
     if (this.config.rag?.architecture === 'agentic' && this.agent) {
-      console.log('[Pipeline] 🤖 Executing in Agentic Mode...');
       const agentReply = await this.agent.run(question, history);
       return { reply: agentReply, sources: [] };
     }
@@ -486,58 +377,68 @@ User: "show me beauty products"
     let reply = '';
     let sources: VectorMatch[] = [];
     let graphData: GraphSearchResult | undefined;
+    let uiTransformation: unknown;
+    let trace: ObservabilityTrace | undefined;
 
-    let uiTransformation: any;
     for await (const chunk of stream) {
       if (typeof chunk === 'string') {
         reply += chunk;
       } else if (typeof chunk === 'object' && chunk !== null) {
-        if ('sources' in chunk) sources = chunk.sources;
-        if ('graphData' in chunk) graphData = chunk.graphData;
-        if ('ui_transformation' in chunk) uiTransformation = chunk.ui_transformation;
+        if ('sources' in chunk) sources = (chunk as ChatResponse).sources;
+        if ('graphData' in chunk) graphData = (chunk as ChatResponse).graphData;
+        if ('ui_transformation' in chunk) uiTransformation = (chunk as ChatResponse).ui_transformation;
+        if ('trace' in chunk) trace = (chunk as { trace: ObservabilityTrace }).trace;
       }
     }
 
-    return { reply, sources, graphData, ui_transformation: uiTransformation };
+    return { reply, sources, graphData, ui_transformation: uiTransformation, trace };
   }
 
   /**
    * High-performance streaming RAG flow.
-   * Yields text chunks first, then the retrieval metadata at the end.
+   * Yields text chunks first, then the retrieval metadata + observability trace at the end.
    */
   async *askStream(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;
     const scoreThreshold = this.config.rag?.scoreThreshold ?? 0.0;
+    const requestId = `req_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;
+    const requestStart = performance.now();
 
     try {
       let searchQuery = question;
+      let rewrittenQuery: string | undefined;
 
       // 1. Query Pre-processing
       if (this.config.rag?.useQueryTransformation) {
         searchQuery = await this.rewriteQuery(question, history);
+        rewrittenQuery = searchQuery !== question ? searchQuery : undefined;
       }
 
-      // 2. Parallel Retrieval
+      // 2. Parallel Retrieval (with latency tracking)
       const hints = QueryProcessor.extractQueryFieldHints(question, this.config.rag?.filterableFields);
       const filter = QueryProcessor.buildQueryFilter(question, hints);
 
-      console.log(`[Pipeline] 🧩 Extracted filters:`, JSON.stringify(filter));
-
+      const embedStart = performance.now();
       const { sources: rawSources, graphData } = await this.retrieve(searchQuery, {
         namespace: ns,
         topK: topK * 2,
         filter
       });
+      const retrieveEnd = performance.now();
+      const embedMs = retrieveEnd - embedStart; // embed + retrieve combined (cache may make embed ~0)
+      const retrieveMs = retrieveEnd - embedStart;
 
       // 3. Reranking
+      const rerankStart = performance.now();
       let sources = rawSources.filter((m) => m.score >= scoreThreshold);
       if (this.config.rag?.useReranking) {
         sources = await this.reranker.rerank(sources, question, topK);
       } else {
         sources = sources.slice(0, topK);
       }
+      const rerankMs = performance.now() - rerankStart;
 
       // 4. Context Augmentation
       let context = sources.length
@@ -551,8 +452,24 @@ User: "show me beauty products"
         context = `GRAPH KNOWLEDGE:\n${graphContext}\n\nVECTOR CONTEXT:\n${context}`;
       }
 
-      // 5. Generation (Streaming)
-      const messages: ChatMessage[] = [...history, { role: 'user', content: question }];
+      // 4.5 Schema Training (non-blocking)
+      const allMetadataKeys = Array.from(new Set(sources.flatMap(s => Object.keys(s.metadata || {}))));
+      const trainingPromise = allMetadataKeys.length > 0
+        ? SchemaMapper.train(this.llmProvider, ns, allMetadataKeys)
+        : Promise.resolve(undefined);
+
+      // 5. Generation (Streaming) — capture the full reply for observability
+      const restrictionSuffix = '\n\n(IMPORTANT: Use plain text only. NEVER generate tables, HTML figures, or text charts. If listing products, use simple bullet points.)';
+      const hardenedHistory = [...history];
+      const userQuestion = { role: 'user', content: question + restrictionSuffix } as ChatMessage;
+      const messages: ChatMessage[] = [...hardenedHistory, userQuestion];
+
+      // Build prompt strings for tracing
+      const systemPrompt = this.config.llm.systemPrompt ?? '';
+      const userPrompt = messages.map(m => `${m.role}: ${m.content}`).join('\n');
+
+      let fullReply = '';
+      const generateStart = performance.now();
 
       if (this.llmProvider.chatStream) {
         const stream = this.llmProvider.chatStream(messages, context);
@@ -560,25 +477,73 @@ User: "show me beauty products"
           throw new Error(`[Pipeline] ${this.config.llm.provider} chatStream returned undefined`);
         }
         for await (const chunk of stream) {
+          fullReply += chunk;
           yield chunk;
         }
       } else {
         const reply = await this.llmProvider.chat(messages, context);
+        fullReply = reply;
         yield reply;
       }
 
-      // 6. Automated UI Fallback (UITransformer)
-      // This provides a secondary, rule-based visualization frame in case the LLM 
-      // fails to adhere to the strict protocol or outputs a text fallback.
-      const uiTransformation = UITransformer.transform(question, sources);
-
-      // Yield retrieval metadata and automated UI transformation at the end
-      yield { 
-        reply: '', 
-        sources, 
-        graphData, 
-        ui_transformation: uiTransformation 
-      } as ChatResponse;
+      const generateMs = performance.now() - generateStart;
+      const totalMs = performance.now() - requestStart;
+
+      const latency: LatencyBreakdown = {
+        embedMs: Math.round(embedMs),
+        retrieveMs: Math.round(retrieveMs),
+        rerankMs: this.config.rag?.useReranking ? Math.round(rerankMs) : undefined,
+        generateMs: Math.round(generateMs),
+        totalMs: Math.round(totalMs),
+      };
+
+      // 6. Token estimation
+      const promptText = systemPrompt + '\n' + context + '\n' + userPrompt;
+      const promptTokens = estimateTokens(promptText);
+      const completionTokens = estimateTokens(fullReply);
+      const tokens: TokenUsage = {
+        promptTokens,
+        completionTokens,
+        totalTokens: promptTokens + completionTokens,
+        estimatedCostUsd: estimateCostUsd(promptTokens, completionTokens, this.config.llm.model),
+      };
+
+      // 7. Hallucination scoring (async, non-blocking on generation yield)
+      const hallucinationResult = await scoreHallucination(this.llmProvider, fullReply, context).catch(() => undefined);
+
+      // 8. Automated UI transformation
+      const trainedSchema = await trainingPromise;
+      const uiTransformation = await this.generateUiTransformation(question, sources, context, trainedSchema);
+
+      // 9. Build observability trace
+      const trace: ObservabilityTrace = {
+        requestId,
+        query: question,
+        rewrittenQuery,
+        systemPrompt,
+        userPrompt: question + restrictionSuffix,
+        chunks: sources.map((s) => ({
+          id: s.id,
+          score: s.score,
+          content: s.content,
+          metadata: s.metadata ?? {},
+          namespace: ns,
+        } satisfies RetrievedChunk)),
+        latency,
+        tokens,
+        hallucinationScore: hallucinationResult?.score,
+        hallucinationReason: hallucinationResult?.reason,
+        timestamp: new Date().toISOString(),
+      };
+
+      // Yield metadata + trace + UI transformation
+      yield {
+        reply: '',
+        sources,
+        graphData,
+        ui_transformation: uiTransformation,
+        trace,
+      } as ChatResponse & { trace: ObservabilityTrace };
 
     } catch (error) {
       throw new Error(`[Pipeline] Stream failed: ${error instanceof Error ? error.message : String(error)}`);
@@ -589,6 +554,24 @@ User: "show me beauty products"
    * Universal retrieval method combining all enabled providers.
    * Uses an LRU-bounded embedding cache to avoid re-embedding the same query.
    */
+  private async generateUiTransformation(
+    question: string,
+    sources: VectorMatch[],
+    _context: string,
+    trainedSchema?: SchemaMap
+  ): Promise<unknown> {
+    if (!sources || sources.length === 0) {
+      return UITransformer.transform(question, sources, this.config, trainedSchema);
+    }
+
+    try {
+      return await UITransformer.analyzeAndDecide(question, sources, this.llmProvider);
+    } catch (err) {
+      console.warn('[Pipeline] generateUiTransformation failed, using heuristic fallback:', err);
+      return UITransformer.transform(question, sources, this.config, trainedSchema);
+    }
+  }
+
   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;
@@ -613,12 +596,9 @@ User: "show me beauty products"
     return { sources, graphData };
   }
 
-  /**
-   * Rewrite the user query for better retrieval performance.
-   */
+  /** Rewrite the user query for better retrieval performance. */
   private async rewriteQuery(question: string, history: ChatMessage[]): Promise<string> {
-    const prompt = `
-Given the following conversation history and a new question, rewrite the question to be a better search query for a vector database. 
+    const prompt = `Given the following conversation history and a new question, rewrite the question to be a better search query for a vector database. 
 Focus on extracting the core intent and entities. Do not answer the question, just rewrite it.
 
 History:
@@ -639,29 +619,19 @@ Optimized Search Query:`;
     return rewrite.trim() || question;
   }
 
-  /**
-   * Generate 3-5 short, relevant questions based on the vector database content.
-   */
+  /** Generate 3-5 short, relevant questions based on the vector database content. */
   async getSuggestions(query: string, namespace?: string): Promise<string[]> {
-    if (!query || query.trim().length < 3) {
-      return [];
-    }
+    if (!query || query.trim().length < 3) return [];
 
     await this.initialize();
     const ns = namespace ?? this.config.projectId;
 
     try {
-      // 1. Retrieve relevant context (top 3 matches)
       const { sources } = await this.retrieve(query, { namespace: ns, topK: 3 });
+      if (sources.length === 0) return [];
 
-      if (sources.length === 0) {
-        return [];
-      }
-
-      // 2. Generate suggestions using LLM
       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 3 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"].
@@ -679,7 +649,6 @@ Suggestions:`;
         ''
       );
 
-      // Simple parsing of JSON array from LLM response
       const match = response.match(/\[[\s\S]*\]/);
       if (match) {
         const suggestions = JSON.parse(match[0]);