@retrivora-ai/rag-engine 2.2.4 → 2.2.6

package/src/utils/SchemaMapper.ts

@@ -13,13 +13,18 @@ export interface SchemaMap {
 /**
  * SchemaMapper — Uses AI to "train" the plugin on unknown database schemas.
  * It maps arbitrary metadata keys to standard UI properties.
+ *
+ * Design notes:
+ * - Bypasses ILLMProvider.chat() entirely to avoid injecting the full chat
+ *   system prompt (which can exceed Groq's free-tier TPM limit).
+ * - Deduplicates concurrent in-flight training calls for the same schema.
+ * - Falls back to heuristics on any failure (non-blocking).
  */
 export class SchemaMapper {
   private static cache = new Map<string, SchemaMap>();
+  /** In-flight deduplication — prevents concurrent identical LLM calls. */
+  private static pending = new Map<string, Promise<SchemaMap>>();
 
-  /**
-   * Descriptions of standard UI properties to help the AI map fields accurately.
-   */
   private static readonly TARGET_PROPERTIES: Record<string, string> = {
     name: 'The primary title, name, or label of the item',
     price: 'The numeric cost, price, MSRP, or amount',
@@ -32,74 +37,95 @@ export class SchemaMapper {
 
   /**
    * Trains the plugin on a set of keys.
-   * This is done once per schema and cached.
+   * Results are cached in-process; concurrent calls for the same key set are
+   * deduplicated to a single LLM request.
    */
   static async train(
     llm: ILLMProvider,
     projectId: string,
     keys: string[]
   ): Promise<SchemaMap> {
-    const cacheKey = `${projectId}:${keys.sort().join(',')}`;
+    const cacheKey = `${projectId}:${keys.slice().sort().join(',')}`;
+
+    // 1. Return from in-process cache (fast path)
     if (this.cache.has(cacheKey)) {
       return this.cache.get(cacheKey)!;
     }
 
-    console.log(`[SchemaMapper] 🧠 Training AI on new schema keys: ${keys.join(', ')}`);
+    // 2. Deduplicate concurrent in-flight requests for the same schema
+    if (this.pending.has(cacheKey)) {
+      return this.pending.get(cacheKey)!;
+    }
+
+    const promise = this._doTrain(llm, cacheKey, keys);
+    this.pending.set(cacheKey, promise);
+    promise.finally(() => this.pending.delete(cacheKey));
+    return promise;
+  }
+
+  private static async _doTrain(
+    llm: ILLMProvider,
+    cacheKey: string,
+    keys: string[]
+  ): Promise<SchemaMap> {
+    console.log(`[SchemaMapper] 🧠 Training on new schema keys: ${keys.join(', ')}`);
 
-    // Dynamically build the property list for the prompt
     const propertyList = Object.entries(this.TARGET_PROPERTIES)
       .map(([prop, desc]) => `- ${prop} (${desc})`)
       .join('\n');
 
-    const prompt = `
-Given these metadata keys from a database: [${keys.join(', ')}]
-
-Identify which keys best correspond to these standard UI properties:
-${propertyList}
-
-Return ONLY a valid JSON object where the keys are the UI properties and the values are the EXACT matching database keys from the list above. 
-If no good match is found for a property, omit it.
-
-Example:
-{
-  "name": "Title",
-  "price": "Variant Price",
-  "brand": "Vendor",
-  "image": "Image Src",
-  "stock": "Variant Inventory Qty"
-}
-`;
+    // Compact prompt — no chat system prompt injected.
+    const messages = [
+      { role: 'system', content: 'You are a database schema expert. Respond ONLY with valid JSON.' },
+      {
+        role: 'user',
+        content:
+          `Keys: [${keys.join(', ')}]\n\nMap to:\n${propertyList}\n\n` +
+          `Return a JSON object like {"name":"Title","price":"Price",...}. Omit unmapped properties.`,
+      },
+    ];
 
     try {
-      const response = await llm.chat(
-        [
-          { role: 'system', content: 'You are a database schema expert. You ONLY respond with valid JSON.' },
-          { role: 'user', content: prompt }
-        ],
-        ''
-      );
-
-      // Extract the first balanced JSON object from the response
-      const startIdx = response.indexOf('{');
-      if (startIdx !== -1) {
-        let braceCount = 0;
-        let endIdx = -1;
-
-        for (let i = startIdx; i < response.length; i++) {
-          if (response[i] === '{') braceCount++;
-          else if (response[i] === '}') {
-            braceCount--;
-            if (braceCount === 0) {
-              endIdx = i;
-              break;
-            }
-          }
+      // Use a direct fetch against the resolved base URL + auth headers,
+      // bypassing ILLMProvider.chat() which injects the full chat system prompt.
+      const baseUrl: string = (llm as any)?.baseUrl ?? '';
+      const apiKey: string = (llm as any)?.apiKey ?? '';
+      const model: string = (llm as any)?.model ?? 'llama-3.1-8b-instant';
+
+      let responseText: string | undefined;
+
+      if (baseUrl) {
+        const endpoint = `${baseUrl.replace(/\/$/, '')}/chat/completions`;
+        const headers: Record<string, string> = { 'Content-Type': 'application/json' };
+        if (apiKey) headers['Authorization'] = `Bearer ${apiKey}`;
+
+        const res = await fetch(endpoint, {
+          method: 'POST',
+          headers,
+          body: JSON.stringify({
+            model,
+            messages,
+            max_tokens: 256,    // Only needs a small JSON object back
+            temperature: 0,
+            stream: false,
+          }),
+          signal: AbortSignal.timeout(8000),
+        });
+
+        if (res.ok) {
+          const data = await res.json();
+          responseText =
+            data?.choices?.[0]?.message?.content ??
+            data?.choices?.[0]?.delta?.content ??
+            undefined;
+        } else {
+          console.warn(`[SchemaMapper] LLM returned ${res.status}, falling back to heuristics.`);
         }
+      }
 
-        if (endIdx !== -1) {
-          const jsonContent = response.substring(startIdx, endIdx + 1);
-          const cleanJson = this.sanitizeJson(jsonContent);
-          const mapping = JSON.parse(cleanJson) as SchemaMap;
+      if (responseText) {
+        const mapping = this._parseJson(responseText);
+        if (mapping) {
           this.cache.set(cacheKey, mapping);
           return mapping;
         }
@@ -108,24 +134,43 @@ Example:
       console.warn('[SchemaMapper] AI training failed, falling back to heuristics:', error);
     }
 
+    // Store empty result so we don't keep retrying on the same schema
+    this.cache.set(cacheKey, {});
     return {};
   }
 
-  /**
-   * Forgiving JSON parser that fixes common AI formatting mistakes.
-   */
-  private static sanitizeJson(s: string): string {
+  private static _parseJson(text: string): SchemaMap | null {
+    const startIdx = text.indexOf('{');
+    if (startIdx === -1) return null;
+    let braceCount = 0;
+    let endIdx = -1;
+    for (let i = startIdx; i < text.length; i++) {
+      if (text[i] === '{') braceCount++;
+      else if (text[i] === '}') {
+        braceCount--;
+        if (braceCount === 0) { endIdx = i; break; }
+      }
+    }
+    if (endIdx === -1) return null;
+    try {
+      return JSON.parse(this._sanitizeJson(text.substring(startIdx, endIdx + 1))) as SchemaMap;
+    } catch {
+      return null;
+    }
+  }
+
+  private static _sanitizeJson(s: string): string {
     return s
-      .replace(/\/\*[\s\S]*?\*\/|([^:]|^)\/\/.*$/gm, '') // Remove comments
-      .replace(/[\u201C\u201D\u2018\u2019]/g, '"') // Replace smart quotes
-      .replace(/'/g, '"') // Replace single quotes
-      .replace(/([{,]\s*)([a-zA-Z0-9_]+?)\s*:/g, '$1"$2":') // Add quotes to keys
-      .replace(/,\s*([}\]])/g, '$1') // Remove trailing commas
+      .replace(/\/\*[\s\S]*?\*\/|([^:]|^)\/\/.*$/gm, '')
+      .replace(/[\u201C\u201D\u2018\u2019]/g, '"')
+      .replace(/'/g, '"')
+      .replace(/([{,]\s*)([a-zA-Z0-9_]+?)\s*:/g, '$1"$2":')
+      .replace(/,\s*([}\]])/g, '$1')
       .trim();
   }
 
   static getCached(projectId: string, keys: string[]): SchemaMap | undefined {
-    const cacheKey = `${projectId}:${keys.sort().join(',')}`;
+    const cacheKey = `${projectId}:${keys.slice().sort().join(',')}`;
     return this.cache.get(cacheKey);
   }
 }

package/src/utils/UITransformer.ts

@@ -128,16 +128,16 @@ export class UITransformer {
       return this.createTextResponse('No data available', 'No relevant data found for your query.');
     }
 
-    console.log('[UITransformer.transform] Processing retrievedData:', {
-      userQuery,
-      retrievedCount: retrievedData.length,
-      sample: retrievedData.slice(0, 2).map((item) => ({
-        id: item?.id,
-        contentType: typeof item?.content,
-        contentSnippet: String(item?.content ?? '').substring(0, 150),
-        metadataKeys: Object.keys(item?.metadata || {}),
-      })),
-    });
+    // console.log('[UITransformer.transform] Processing retrievedData:', {
+    //   userQuery,
+    //   retrievedCount: retrievedData.length,
+    //   sample: retrievedData.slice(0, 2).map((item) => ({
+    //     id: item?.id,
+    //     contentType: typeof item?.content,
+    //     contentSnippet: String(item?.content ?? '').substring(0, 150),
+    //     metadataKeys: Object.keys(item?.metadata || {}),
+    //   })),
+    // });
 
     const resolvedIntent = intent ?? this.detectIntentHeuristic(userQuery);
 
@@ -642,16 +642,16 @@ RULES:
     const fallbackData: BarChartData[] = barData.length > 0
       ? barData
       : data.slice(0, 12).map((item, index) => {
-          const meta = item.metadata || {};
-          const label = String(
-            this.getDynamicVal(meta, 'name') ??
-            meta.label ??
-            meta.title ??
-            `Result ${index + 1}`,
-          );
-          const value = this.extractNumericValue(meta) ?? item.score ?? 0;
-          return { category: label, value: Number(value) };
-        });
+        const meta = item.metadata || {};
+        const label = String(
+          this.getDynamicVal(meta, 'name') ??
+          meta.label ??
+          meta.title ??
+          `Result ${index + 1}`,
+        );
+        const value = this.extractNumericValue(meta) ?? item.score ?? 0;
+        return { category: label, value: Number(value) };
+      });
 
     return {
       type: horizontal ? 'horizontal_bar' : 'bar_chart',
@@ -724,8 +724,8 @@ RULES:
     const numericField = this.selectNumericField(profile, query);
     const values = numericField
       ? profile.records
-          .map(record => this.toFiniteNumber(record.fields[numericField.key]))
-          .filter((value): value is number => value !== null)
+        .map(record => this.toFiniteNumber(record.fields[numericField.key]))
+        .filter((value): value is number => value !== null)
       : [];
 
     const value = operation === 'count' || values.length === 0