@retrivora-ai/rag-engine 1.7.9 → 1.8.1

package/src/config/RagConfig.ts

@@ -182,6 +182,11 @@ export interface RAGConfig {
    * Used to dynamically extract hints from natural language queries. 
    */
   filterableFields?: string[];
+  /** Optional mapping from database metadata fields to standard UI properties.
+   * Useful for databases with custom schemas.
+   * e.g., { "name": "ProductName", "price": "SalesPrice", "image": "ThumbnailUrl" }
+   */
+  uiMapping?: Record<string, string>;
 }
 
 // ---------------------------------------------------------------------------

package/src/config/serverConfig.ts

@@ -76,7 +76,9 @@ export function getEnvConfig(env: Record<string, string | undefined> = process.e
     vectorDbOptions.apiKey = readString(env, 'PINECONE_API_KEY') ?? (base?.vectorDb?.options?.apiKey as string) ?? '';
     vectorDbOptions.indexName = readString(env, 'PINECONE_INDEX') ?? (base?.vectorDb?.indexName as string) ?? (base?.vectorDb?.options?.indexName as string);
   } else if (vectorProvider === 'pgvector' || vectorProvider === 'postgresql') {
-    vectorDbOptions.connectionString = readString(env, 'PGVECTOR_CONNECTION_STRING') ?? (base?.vectorDb?.options?.connectionString as string) ?? '';
+    vectorDbOptions.connectionString = readString(env, 'PGVECTOR_CONNECTION_STRING') ?? readString(env, 'POSTGRES_URL') ?? (base?.vectorDb?.options?.connectionString as string) ?? '';
+    vectorDbOptions.tables = (readString(env, 'VECTOR_DB_TABLES') ?? readString(env, 'POSTGRES_TABLES'))?.split(',').map(t => t.trim()) ?? (base?.vectorDb?.options?.tables as string[]);
+    vectorDbOptions.searchFields = readString(env, 'POSTGRES_SEARCH_FIELDS')?.split(',').map(f => f.trim()) ?? (base?.vectorDb?.options?.searchFields as string[]);
     vectorDbOptions.dimensions = embeddingDimensions;
   } else if (vectorProvider === 'mongodb') {
     vectorDbOptions.uri = readString(env, 'MONGODB_URI') ?? (base?.vectorDb?.options?.uri as string) ?? '';

package/src/core/Pipeline.ts

@@ -14,6 +14,7 @@ import { EmbeddingStrategyResolver } from '../config/EmbeddingStrategy';
 import { QueryProcessor } from './QueryProcessor';
 import { IngestDocument, ChatResponse, VectorMatch, GraphSearchResult, RetrievalResult, UpsertDocument } from '../types';
 import { UITransformer } from '../utils/UITransformer';
+import { SchemaMapper, SchemaMap } from '../utils/SchemaMapper';
 
 // ─── LRU Embedding Cache ───────────────────────────────────────────────────────
 
@@ -97,209 +98,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 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 MAPPING (STRICT)
-
-| Intent   | View       |
-|----------|-----------|
-| explore  | carousel  |
-| analyze  | chart     |
-| compare  | table     |
-| detail   | table     |
-
-⚠️ Override Rules:
-- If query contains "distribution", "percentage", "ratio" → ALWAYS chart
-- If query contains "show products", "list items" → ALWAYS carousel
-- NEVER mix views
-
----
-
-## 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": []
-}
-\`\`\`
-
----
-
-## 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,
-  "inStock": boolean (optional)
-}
-
-### 🔹 TABLE FORMAT
-- columns MUST match row keys
-- rows MUST be flat objects
-
----
-
-## 5. DATA RULES
-
-- NEVER hallucinate fields or data
-- ONLY use retrieved vector DB data
-- NO placeholders ("...", "N/A")
-- If image/link missing → REMOVE FIELD
-- Aggregate BEFORE rendering chart
-- IF NO DATA IS FOUND in the retrieved context:
-  → DO NOT output a \`\`\`ui\`\`\` block
-  → Instead, state clearly that no relevant products or data were found.
-
----
-
-## 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
-- No explanation inside block
-- NEVER use ASCII charts, Markdown tables, or text-based drawings for data.
-- If data is tabular or statistical, ALWAYS use the \`\`\`ui\`\`\` block.
-
----
-
-## 8. EXAMPLES
-
-User: "show me beauty products"
-
-\`\`\`ui
-{
-  "view": "carousel",
-  "title": "Beauty Products",
-  "description": "Browse available beauty items",
-  "metadata": {
-    "intent": "explore",
-    "confidence": 0.95
-  },
-  "carousel": {
-    "items": [
-      {
-        "id": "1",
-        "name": "Mascara",
-        "brand": "Essence",
-        "price": 9.99,
-        "image": "https://example.com/mascara.jpg",
-        "link": "https://example.com/p1",
-        "inStock": true
-      }
-    ]
-  },
-  "insights": ["Most products are affordable"]
-}
-\`\`\`
-
----
-
-User: "distribution of products across categories"
-
-\`\`\`ui
-{
-  "view": "chart",
-  "title": "Product Distribution",
-  "description": "Category-wise product distribution",
-  "metadata": {
-    "intent": "analyze",
-    "confidence": 0.98
-  },
-  "chart": {
-    "type": "pie",
-    "xKey": "label",
-    "yKeys": ["value"],
-    "data": [
-      { "label": "Beauty", "value": 15 },
-      { "label": "Electronics", "value": 10 }
-    ]
-  },
-  "insights": ["Beauty dominates inventory"]
-}
-\`\`\`
+You are a helpful product assistant. Use the provided context to answer questions accurately.
+
+### 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.
+
+### 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;
-    }
+
+    this.config.llm.systemPrompt = chartInstruction;
 
     console.log(`[Pipeline] Initializing with provider: ${this.config.vectorDb.provider}...`);
     // Resolve vector DB provider
@@ -487,16 +315,18 @@ User: "distribution of products across categories"
     let sources: VectorMatch[] = [];
     let graphData: GraphSearchResult | undefined;
 
+    let uiTransformation: unknown;
     for await (const chunk of stream) {
       if (typeof chunk === 'string') {
         reply += chunk;
-      } else if ('sources' in chunk) {
-        sources = chunk.sources;
-        graphData = chunk.graphData;
+      } 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;
       }
     }
 
-    return { reply, sources, graphData };
+    return { reply, sources, graphData, ui_transformation: uiTransformation };
   }
 
   /**
@@ -549,8 +379,17 @@ User: "distribution of products across categories"
         context = `GRAPH KNOWLEDGE:\n${graphContext}\n\nVECTOR CONTEXT:\n${context}`;
       }
 
+      // 4.5 Schema Training (Dynamic Discovery) - Start training but don't block the stream yet
+      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)
-      const messages: ChatMessage[] = [...history, { role: 'user', content: question }];
+      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];
 
       if (this.llmProvider.chatStream) {
         const stream = this.llmProvider.chatStream(messages, context);
@@ -566,16 +405,17 @@ User: "distribution of products across categories"
       }
 
       // 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);
+      // Wait for training to finish if it hasn't already
+      const trainedSchema = await trainingPromise;
+
+      const uiTransformation = await this.generateUiTransformation(question, sources, context, trainedSchema);
 
       // Yield retrieval metadata and automated UI transformation at the end
-      yield { 
-        reply: '', 
-        sources, 
-        graphData, 
-        ui_transformation: uiTransformation 
+      yield {
+        reply: '',
+        sources,
+        graphData,
+        ui_transformation: uiTransformation
       } as ChatResponse;
 
     } catch (error) {
@@ -587,6 +427,27 @@ User: "distribution of products across categories"
    * 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 {
+      // Delegate to the centralized LLM-driven visualization decision.
+      // UITransformer.analyzeAndDecide() internally falls back to the keyword
+      // heuristic if the LLM call fails or returns unparseable JSON.
+      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;

package/src/core/ProviderRegistry.ts

@@ -77,8 +77,8 @@ export class ProviderRegistry {
       }
       case 'pgvector':
       case 'postgresql': {
-        const { PostgreSQLProvider } = await import('../providers/vectordb/PostgreSQLProvider');
-        return PostgreSQLProvider;
+        const { MultiTablePostgresProvider } = await import('../providers/vectordb/MultiTablePostgresProvider');
+        return MultiTablePostgresProvider as unknown as VectorProviderClass;
       }
       case 'mongodb': {
         const { MongoDBProvider } = await import('../providers/vectordb/MongoDBProvider');

package/src/core/VectorPlugin.ts

@@ -38,6 +38,15 @@ export class VectorPlugin {
     return this.config;
   }
 
+  /**
+   * Get the initialized LLM provider (available after the first request).
+   * Used by handlers to pass to UITransformer.analyzeAndDecide() for
+   * LLM-driven visualization decisions.
+   */
+  getLLMProvider() {
+    return this.pipeline.getLLMProvider();
+  }
+
   /**
    * Perform pre-flight health checks on all configured providers.
    * Useful to verify connectivity before running operations.

package/src/handlers/index.ts

@@ -27,7 +27,7 @@ function sseMetaFrame(meta: unknown): string {
 
 /** Encode the UI transformation result as an SSE frame. */
 function sseUIFrame(uiTransformation: unknown): string {
-  return `data: ${JSON.stringify({ type: 'ui_transformation', ...((uiTransformation as object) ?? {}) })}\n\n`;
+  return `data: ${JSON.stringify({ type: 'ui_transformation', data: uiTransformation })}\n\n`;
 }
 
 /** Encode a stream error as an SSE frame. */
@@ -152,15 +152,31 @@ export function createStreamHandler(configOrPlugin?: Partial<RagConfig> | Vector
               // Retrieval metadata object — always the final frame
               enqueue(sseMetaFrame(chunk));
 
-              // Transform retrieved sources to UI representation
-              const sources = (chunk as ChatResponse)?.sources || [];
-              if (sources && sources.length > 0) {
+              // ── LLM-driven visualization decision ─────────────────────────
+              // Sends question + RAG context to the LLM and asks it to choose
+              // the best visualization format (bar_chart, pie_chart, table…).
+              // Falls back to keyword heuristic if the LLM call fails.
+              const responseChunk = chunk as ChatResponse;
+              const sources = responseChunk?.sources || [];
+
+              if (sources.length > 0) {
                 try {
-                  const uiTransformation = UITransformer.transform(message, sources);
-                  enqueue(sseUIFrame(uiTransformation));
+                  const llmProvider = plugin.getLLMProvider?.();
+                  const uiTransformation = responseChunk?.ui_transformation ??
+                    (llmProvider
+                      ? await UITransformer.analyzeAndDecide(message, sources, llmProvider)
+                      : UITransformer.transform(message, sources, plugin.getConfig()));
+
+                  if (uiTransformation) {
+                    enqueue(sseUIFrame(uiTransformation));
+                  }
                 } catch (transformError) {
                   console.warn('[createStreamHandler] UI transformation warning:', transformError);
-                  // Don't fail the stream if transformation fails, just skip it
+                  // Last-resort fallback: heuristic transform
+                  try {
+                    const fallback = UITransformer.transform(message, sources, plugin.getConfig());
+                    if (fallback) enqueue(sseUIFrame(fallback));
+                  } catch { /* silent */ }
                 }
               }
             }

package/src/hooks/useRagChat.ts

@@ -155,7 +155,7 @@ export function useRagChat(
           buffer = buffer.slice(lastBoundary + 2);
 
           for (const frame of parseSseChunk(complete)) {
-            if (frame.type === 'text') {
+            if (frame.type === 'text' && frame.text) {
               assistantContent += frame.text;
             } else if (frame.type === 'metadata') {
               sources = frame.sources ?? [];
@@ -184,7 +184,7 @@ export function useRagChat(
         // Process any remaining buffered data
         if (buffer.trim()) {
           for (const frame of parseSseChunk(buffer)) {
-            if (frame.type === 'text') assistantContent += frame.text;
+            if (frame.type === 'text' && frame.text) assistantContent += frame.text;
             else if (frame.type === 'metadata') sources = frame.sources ?? [];
           }
         }

package/src/llm/LLMFactory.ts

@@ -12,10 +12,50 @@ interface LLMProviderStatic {
   getHealthChecker?: () => IProviderHealthChecker;
 }
 
+/** Registry for custom provider factories registered at runtime. */
+const customProviders = new Map<string, (config: LLMConfig) => ILLMProvider>();
+
 /**
  * LLMFactory — instantiates the correct ILLMProvider based on llmConfig.provider.
  */
 export class LLMFactory {
+  /**
+   * Register a custom LLM provider factory at runtime.
+   *
+   * Use this to add support for any LLM backend (Azure OpenAI, Cohere, Mistral,
+   * Bedrock, etc.) without modifying this library's source code.
+   *
+   * @example
+   * // In your Next.js app initialization:
+   * import { LLMFactory } from '@retrivora-ai/rag-engine/server';
+   * import { MyCustomProvider } from './providers/MyCustomProvider';
+   *
+   * LLMFactory.register('my-provider', (config) => new MyCustomProvider(config));
+   *
+   * // Then set in your .env.local:
+   * // LLM_PROVIDER=my-provider
+   */
+  static register(name: string, factory: (config: LLMConfig) => ILLMProvider): void {
+    customProviders.set(name.toLowerCase(), factory);
+    console.log(`[LLMFactory] Registered custom provider: "${name}"`);
+  }
+
+  /**
+   * Unregister a previously registered custom provider.
+   */
+  static unregister(name: string): void {
+    customProviders.delete(name.toLowerCase());
+  }
+
+  /**
+   * List all registered provider names (built-in + custom).
+   */
+  static listProviders(): string[] {
+    return [
+      'openai', 'anthropic', 'ollama', 'gemini', 'rest', 'universal_rest', 'custom',
+      ...Array.from(customProviders.keys()),
+    ];
+  }
   static create(llmConfig: LLMConfig, embeddingConfig?: EmbeddingConfig): ILLMProvider {
     switch (llmConfig.provider) {
       case 'openai':
@@ -30,11 +70,23 @@ export class LLMFactory {
       case 'universal_rest':
       case 'custom':
         return new UniversalLLMAdapter(llmConfig);
-      default:
+      default: {
+        // Check the runtime custom provider registry first
+        const providerName = String(llmConfig.provider ?? '').toLowerCase();
+        const customFactory = customProviders.get(providerName);
+        if (customFactory) {
+          return customFactory(llmConfig);
+        }
+        // Fall back to UniversalLLMAdapter if a baseUrl is provided
         if (llmConfig.baseUrl || (llmConfig.options as Record<string, unknown>)?.baseUrl) {
           return new UniversalLLMAdapter(llmConfig);
         }
-        throw new Error(`[LLMFactory] Unknown provider "${llmConfig.provider}"`);
+        throw new Error(
+          `[LLMFactory] Unknown provider "${llmConfig.provider}". ` +
+          `Built-in providers: ${LLMFactory.listProviders().join(', ')}. ` +
+          `Register a custom provider with LLMFactory.register().`
+        );
+      }
     }
   }
 

package/src/llm/providers/AnthropicProvider.ts

@@ -75,13 +75,15 @@ export class AnthropicProvider implements ILLMProvider {
   }
 
   async chat(messages: ChatMessage[], context: string, options?: ChatOptions): Promise<string> {
-    const systemPrompt =
+    const resolvedPrompt = options?.systemPrompt ??
       this.llmConfig.systemPrompt ??
       `You are a helpful assistant. Use the context below to answer the user's question accurately.\n\nContext:\n${context}`;
 
-    const system = systemPrompt.includes('{{context}}')
-      ? systemPrompt.replace('{{context}}', context)
-      : `${systemPrompt}\n\nContext:\n${context}`;
+    const system = resolvedPrompt.includes('{{context}}')
+      ? resolvedPrompt.replace('{{context}}', context)
+      : options?.systemPrompt
+        ? resolvedPrompt // per-call: context is already in the user message
+        : `${resolvedPrompt}\n\nContext:\n${context}`;
 
     const anthropicMessages: Anthropic.MessageParam[] = messages.map((m) => ({
       role: m.role === 'assistant' ? 'assistant' : 'user',
@@ -100,13 +102,15 @@ export class AnthropicProvider implements ILLMProvider {
   }
 
   async *chatStream(messages: ChatMessage[], context: string, options?: ChatOptions): AsyncIterable<string> {
-    const systemPrompt =
+    const resolvedPrompt = options?.systemPrompt ??
       this.llmConfig.systemPrompt ??
       `You are a helpful assistant. Use the context below to answer the user's question accurately.\n\nContext:\n${context}`;
 
-    const system = systemPrompt.includes('{{context}}')
-      ? systemPrompt.replace('{{context}}', context)
-      : `${systemPrompt}\n\nContext:\n${context}`;
+    const system = resolvedPrompt.includes('{{context}}')
+      ? resolvedPrompt.replace('{{context}}', context)
+      : options?.systemPrompt
+        ? resolvedPrompt
+        : `${resolvedPrompt}\n\nContext:\n${context}`;
 
     const anthropicMessages: Anthropic.MessageParam[] = messages.map((m) => ({
       role: m.role === 'assistant' ? 'assistant' : 'user',