@retrivora-ai/rag-engine 1.8.0 → 1.8.2

package/src/components/MessageBubble.tsx

@@ -1,650 +1,26 @@
 'use client';
 
 import React from 'react';
-import { Bot, User, ChevronDown, ChevronRight } from 'lucide-react';
+import { Bot, User, ChevronDown, ChevronRight, Activity } from 'lucide-react';
 import ReactMarkdown from 'react-markdown';
-import remarkGfm from 'remark-gfm';
-import { MessageBubbleProps, Product } from '../types';
-import { ChatViewportSize } from '../types/props';
+import { MessageBubbleProps, Product, UITransformationResponse } from '../types';
 import { SourceCard } from './SourceCard';
 import { ProductCarousel } from './ProductCarousel';
-import { DynamicChart } from './DynamicChart';
 import { VisualizationRenderer } from './VisualizationRenderer';
-import { UITransformationResponse } from '../utils/UITransformer';
-
-// ─── JSON sanitization ────────────────────────────────────────────────────────
-// Order matters: each step is a targeted fix, no step undoes a previous one.
-
-function sanitizeJson(raw: string): string {
-  let s = raw.trim();
-
-  // 1. Evaluate simple arithmetic expressions in values  e.g.  "v": 495 / 1000
-  s = s.replace(
-    /:\s*(-?\d+(?:\.\d+)?)\s*([+\-*/])\s*(-?\d+(?:\.\d+)?)/g,
-    (_, a, op, b) => {
-      const n1 = parseFloat(a);
-      const n2 = parseFloat(b);
-      const result =
-        op === '+' ? n1 + n2
-          : op === '-' ? n1 - n2
-            : op === '*' ? n1 * n2
-              : n2 !== 0 ? n1 / n2
-                : n1;
-      return `: ${result}`;
-    },
-  );
-
-  // 2. Strip JS-style comments while preserving strings (e.g. URLs)
-  s = s.replace(/("(?:\\.|[^\\"])*")|\/\/[^\n]*|\/\*[\s\S]*?\*\//g, (m, g1) => g1 || '');
-
-  // 3. Quote bare (unquoted) object keys
-  s = s.replace(/([{,]\s*)([A-Za-z_$][A-Za-z0-9_$]*)\s*:/g, '$1"$2":');
-
-  // 4. Replace single-quoted strings with double-quoted ones
-  s = s.replace(/'([^'\\]*(?:\\.[^'\\]*)*)'/g, '"$1"');
-
-  // 5. Remove trailing commas before ] or }
-  s = s.replace(/,\s*([}\]])/g, '$1');
-
-  // 6. Remove completely empty key-value pairs left by prior cleanup
-  s = s.replace(/([{,])\s*""\s*(?=[,}])/g, '$1');
-
-  // 7. Stack-based structure balancer.
-  //    Walk every character, maintain a stack of openers so we close each
-  //    unclosed structure with its exact matching closer — in the right order.
-  //    A flat counter approach appends all ]'s then all }'s, which corrupts
-  //    nested structures like {..., [..., {...}]}.
-  {
-    const stack: ('{' | '[')[] = [];
-    let inString = false;
-    let escape = false;
-
-    for (const ch of s) {
-      if (escape) { escape = false; continue; }
-      if (ch === '\\' && inString) { escape = true; continue; }
-      if (ch === '"') { inString = !inString; continue; }
-      if (inString) continue;
-
-      if (ch === '{' || ch === '[') {
-        stack.push(ch);
-      } else if (ch === '}') {
-        if (stack[stack.length - 1] === '{') stack.pop();
-      } else if (ch === ']') {
-        if (stack[stack.length - 1] === '[') stack.pop();
-      }
-    }
-
-    // Close any unterminated string literal first
-    if (inString) s += '"';
-
-    // Close every unclosed opener in reverse (innermost first)
-    for (let i = stack.length - 1; i >= 0; i--) {
-      s += stack[i] === '{' ? '}' : ']';
-    }
-  }
-
-  // 8. Trim trailing garbage that appears after the outermost closing char.
-  //    Only fires when something non-whitespace follows the last } or ].
-  const lastClose = Math.max(s.lastIndexOf('}'), s.lastIndexOf(']'));
-  if (lastClose !== -1 && lastClose < s.length - 1 && /\S/.test(s.slice(lastClose + 1))) {
-    s = s.substring(0, lastClose + 1);
-  }
-
-  return s;
-}
-
-function extractLikelyJsonBlock(raw: string): string {
-  const trimmed = raw.trim();
-  const objectStart = trimmed.indexOf('{');
-  const arrayStart = trimmed.indexOf('[');
-
-  const starts = [objectStart, arrayStart].filter((index) => index >= 0);
-  if (starts.length === 0) return trimmed;
-
-  return trimmed.slice(Math.min(...starts));
-}
-
-function looksLikeStructuredPayload(raw: string): boolean {
-  return /"?(view|chartType|type|data|items|rows|products|results)"?\s*:/.test(raw);
-}
-
-function stripMarkdownTables(raw: string): string {
-  return raw.replace(
-    /(?:^|\n)\|[^\n]+\|\n\|(?:\s*:?-+:?\s*\|)+\n(?:\|[^\n]*\|\n?)*/g,
-    '\n\n',
-  );
-}
-
-function stripStructuredUiLabels(raw: string): string {
-  return raw.replace(
-    /(?:^|\n)\*\*(?:Chart Data|Table Data|Visualization|Visual Representation)\*\:\s*(?=\n|$)/gi,
-    '\n',
-  ).replace(
-    /(?:^|\n)(?:Chart Data|Table Data|Visualization|Visual Representation)\s*\:\s*(?=\n|$)/gi,
-    '\n',
-  );
-}
-
-function extractBareJsonObject(raw: string): string | null {
-  const start = raw.indexOf('{');
-  if (start === -1) return null;
-
-  let depth = 0;
-  let inString = false;
-  let escape = false;
-
-  for (let i = start; i < raw.length; i += 1) {
-    const ch = raw[i];
-
-    if (escape) {
-      escape = false;
-      continue;
-    }
-
-    if (ch === '\\' && inString) {
-      escape = true;
-      continue;
-    }
-
-    if (ch === '"') {
-      inString = !inString;
-      continue;
-    }
-
-    if (inString) continue;
-
-    if (ch === '{') depth += 1;
-    if (ch === '}') {
-      depth -= 1;
-      if (depth === 0) {
-        return raw.slice(start, i + 1);
-      }
-    }
-  }
-
-  return null;
-}
-
-function extractStructuredPayload(raw: string): { payload: string | null; text: string } {
-  let working = raw;
-  const fenceRegex = /```(?:[a-z]+)?\s*([\s\S]*?)```/gi;
-
-  for (const match of raw.matchAll(fenceRegex)) {
-    const body = match[1]?.trim() ?? '';
-    if (!looksLikeStructuredPayload(body)) continue;
-
-    working = working.replace(match[0], '');
-    return {
-      payload: body,
-      text: stripStructuredUiLabels(working).replace(/\n{3,}/g, '\n\n').trim(),
-    };
-  }
-
-  const bareJson = extractBareJsonObject(raw);
-  if (bareJson && looksLikeStructuredPayload(bareJson)) {
-    working = working.replace(bareJson, '');
-    return {
-      payload: bareJson,
-      text: stripStructuredUiLabels(working).replace(/\n{3,}/g, '\n\n').trim(),
-    };
-  }
-
-  return {
-    payload: null,
-    text: stripStructuredUiLabels(raw).replace(/\n{3,}/g, '\n\n').trim(),
-  };
-}
-
-function isLikelyUiOnlyMessage(text: string): boolean {
-  const trimmed = text.trim();
-  return (
-    trimmed.length === 0 ||
-    /^(chart data|table data|visualization|visual representation|product catalog|product recommendations|catalog summary)\s*:?\s*$/i.test(trimmed)
-  );
-}
-
-// ─── Tiny helpers ─────────────────────────────────────────────────────────────
-
-function resolveImage(data: Record<string, unknown>): string | undefined {
-  const isPlaceholder = (val: unknown) => 
-    typeof val === 'string' && (val === '...' || val === '' || val.toLowerCase() === 'null' || val.toLowerCase() === 'undefined');
-
-  for (const key of ['image', 'img', 'thumbnail']) {
-    const v = data[key];
-    if (typeof v === 'string' && v && !isPlaceholder(v)) return v;
-  }
-  if (Array.isArray(data.images) && typeof data.images[0] === 'string' && !isPlaceholder(data.images[0])) {
-    return data.images[0];
-  }
-  return undefined;
-}
-
-/** Safely render table cell children that might be plain objects. */
-function normaliseChild(children: React.ReactNode): React.ReactNode {
-  if (
-    typeof children === 'object' &&
-    children !== null &&
-    !Array.isArray(children) &&
-    !React.isValidElement(children)
-  ) {
-    return JSON.stringify(children);
-  }
-  return children;
-}
-
-function normaliseFieldName(value: string): string {
-  return value.toLowerCase().replace(/[^a-z0-9]/g, '');
-}
-
-function formatColumnLabel(value: string): string {
-  return value
-    .replace(/([a-z0-9])([A-Z])/g, '$1 $2')
-    .replace(/[_-]+/g, ' ')
-    .replace(/\s+/g, ' ')
-    .trim()
-    .replace(/\b\w/g, (char) => char.toUpperCase());
-}
-
-// ─── Inline data table renderer ───────────────────────────────────────────────
-// Handles chart configs where type === "table". The LLM may emit any of:
-//   columns, dataKeys (+optional xAxisKey), or no key hints (auto-detect from row).
-
-interface TableConfig {
-  type: 'table';
-  title?: string;
-  description?: string;
-  columns?: string[];   // preferred: explicit ordered column labels
-  dataKeys?: string[];  // fallback column keys (may omit xAxisKey)
-  xAxisKey?: string;    // row-label key to prepend when using dataKeys
-  data: Record<string, unknown>[];
-}
-
-function DataTable({ config, viewportSize = 'large' }: { config: TableConfig; viewportSize?: ChatViewportSize }) {
-  const isCompact = viewportSize === 'compact';
-  const isMedium = viewportSize === 'medium';
-  const columns = React.useMemo<Array<{ label: string; accessor: string }>>(() => {
-    const rowKeys = Array.from(
-      new Set(
-        config.data.flatMap((row) => Object.keys(row)),
-      ),
-    );
-
-    const findAccessor = (label: string) => {
-      const exact = rowKeys.find((key) => key === label);
-      if (exact) return exact;
-
-      const normalisedLabel = normaliseFieldName(label);
-      const normalisedMatch = rowKeys.find((key) => normaliseFieldName(key) === normalisedLabel);
-      if (normalisedMatch) return normalisedMatch;
-
-      return label;
-    };
-
-    if (Array.isArray(config.columns) && config.columns.length) {
-      return config.columns.map((label) => ({
-        label: formatColumnLabel(label),
-        accessor: findAccessor(label),
-      }));
-    }
-
-    if (Array.isArray(config.dataKeys) && config.dataKeys.length) {
-      const keys = config.xAxisKey && !config.dataKeys.includes(config.xAxisKey)
-        ? [config.xAxisKey, ...config.dataKeys]
-        : config.dataKeys;
-
-      return keys.map((key) => ({
-        label: formatColumnLabel(key),
-        accessor: findAccessor(key),
-      }));
-    }
-
-    return rowKeys.map((key) => ({ label: formatColumnLabel(key), accessor: key }));
-  }, [config]);
-
-  if (!config.data.length || !columns.length) {
-    return <p className="text-xs text-slate-400 italic my-2">No data to display.</p>;
-  }
-
-  const formatCell = (val: unknown): string => {
-    if (val === null || val === undefined) return '—';
-    if (typeof val === 'boolean') return val ? '✓' : '✗';
-    return String(val);
-  };
-
-  return (
-    <div className="my-4 rounded-xl border border-slate-200 dark:border-white/10 shadow-sm overflow-hidden">
-      {config.title && (
-        <div className="px-4 py-2.5 bg-slate-50 dark:bg-white/5 border-b border-slate-200 dark:border-white/10">
-          <p className="text-xs font-semibold text-slate-600 dark:text-white/70 uppercase tracking-wide">
-            {config.title}
-          </p>
-        </div>
-      )}
-
-      <div className="overflow-x-auto">
-        <table className={`w-full text-left border-collapse ${isCompact ? 'min-w-[240px] text-xs' : isMedium ? 'min-w-[280px] text-[13px]' : 'min-w-[320px] text-sm'}`}>
-          <thead className="bg-slate-50 dark:bg-white/5 border-b border-slate-200 dark:border-white/10">
-            <tr>
-              {columns.map((column) => (
-                <th
-                  key={column.label}
-                  className={`${isCompact ? 'px-2.5 py-2' : 'px-4 py-3'} font-semibold text-slate-700 dark:text-white/90 whitespace-nowrap`}
-                >
-                  {column.label}
-                </th>
-              ))}
-            </tr>
-          </thead>
-          <tbody>
-            {config.data.map((row, ri) => (
-              <tr
-                key={ri}
-                className="border-b border-slate-100 dark:border-white/5 last:border-0 hover:bg-slate-50/60 dark:hover:bg-white/5 transition-colors"
-              >
-                {columns.map((column) => (
-                  <td
-                    key={column.label}
-                    className={`${isCompact ? 'px-2.5 py-2' : 'px-4 py-3'} text-slate-600 dark:text-white/70 ${column.accessor.toLowerCase().includes('description') ? 'whitespace-normal break-words' : 'whitespace-nowrap'}`}
-                  >
-                    {formatCell(row[column.accessor])}
-                  </td>
-                ))}
-              </tr>
-            ))}
-          </tbody>
-        </table>
-      </div>
-
-      <div className="px-4 py-2 bg-slate-50 dark:bg-white/5 border-t border-slate-100 dark:border-white/5">
-        <p className="text-[11px] text-slate-400 dark:text-white/30">
-          {config.data.length} row{config.data.length !== 1 ? 's' : ''}
-        </p>
-      </div>
-    </div>
-  );
-}
-
-// ─── Universal UI Dispatcher ────────────────────────────────────────────────
-// Handles charts, product carousels, and data tables via a single protocol.
-
-interface UIConfig {
-  view: 'chart' | 'carousel' | 'table';
-  chartType?: 'pie' | 'bar' | 'line';
-  xAxisKey?: string;
-  dataKeys?: string[];
-  columns?: string[];
-  colors?: string[];
-  data: Record<string, unknown>[];
-  title?: string;
-  description?: string;
-  insights?: string[];
-  type?: string;
-  items?: Record<string, unknown>[];
-  rows?: Record<string, unknown>[];
-  products?: Record<string, unknown>[];
-  results?: Record<string, unknown>[];
-
-  // V7 Protocol Fields
-  carousel?: {
-    items: Record<string, unknown>[];
-  };
-  chart?: {
-    type: 'pie' | 'bar' | 'line';
-    xKey: string;
-    yKeys: string[];
-    data: Record<string, unknown>[];
-  };
-  table?: {
-    columns: string[];
-    rows?: Record<string, unknown>[];
-    data?: Record<string, unknown>[];
-  };
-}
-
-interface RawUIConfig extends Omit<UIConfig, 'data' | 'items' | 'rows' | 'products' | 'results' | 'carousel' | 'chart' | 'table'> {
-  data: Array<Record<string, unknown> | unknown[]>;
-  items?: Array<Record<string, unknown> | unknown[]>;
-  rows?: Array<Record<string, unknown> | unknown[]>;
-  products?: Array<Record<string, unknown> | unknown[]>;
-  results?: Array<Record<string, unknown> | unknown[]>;
-
-  // V7 Protocol Fields (Raw)
-  carousel?: {
-    items: Array<Record<string, unknown> | unknown[]>;
-  };
-  chart?: {
-    type: 'pie' | 'bar' | 'line';
-    xKey: string;
-    yKeys: string[];
-    data: Array<Record<string, unknown> | unknown[]>;
-  };
-  table?: {
-    columns: string[];
-    rows?: Array<Record<string, unknown> | unknown[]>;
-    data?: Array<Record<string, unknown> | unknown[]>;
-  };
-}
-
-function normalizeTabularRows(
-  rows: Array<Record<string, unknown> | unknown[]> | undefined,
-  columns?: string[],
-): Record<string, unknown>[] {
-  if (!Array.isArray(rows)) return [];
-
-  return rows
-    .map((row) => {
-      if (Array.isArray(row)) {
-        const keys = Array.isArray(columns) && columns.length > 0
-          ? columns
-          : row.map((_, index) => `column_${index + 1}`);
-
-        return keys.reduce<Record<string, unknown>>((acc, key, index) => {
-          acc[key] = row[index];
-          return acc;
-        }, {});
-      }
-
-      if (row && typeof row === 'object') {
-        return row as Record<string, unknown>;
-      }
-
-      return { value: row };
-    })
-    .filter((row) => Object.keys(row).length > 0);
-}
-
-function resolveStructuredRows(config: Partial<RawUIConfig> & Record<string, unknown>) {
-  // V7 Nested Structures
-  if (config.carousel?.items && Array.isArray(config.carousel.items)) return config.carousel.items;
-  if (config.chart?.data && Array.isArray(config.chart.data)) return config.chart.data;
-  if (config.table?.rows && Array.isArray(config.table.rows)) return config.table.rows;
-  if (config.table?.data && Array.isArray(config.table.data)) return config.table.data;
-
-  // V6 Flat Structures
-  if (Array.isArray(config.data) && config.data.length > 0) return config.data;
-  if (Array.isArray(config.items) && config.items.length > 0) return config.items;
-  if (Array.isArray(config.products) && config.products.length > 0) return config.products;
-  if (Array.isArray(config.rows) && config.rows.length > 0) return config.rows;
-  if (Array.isArray(config.results) && config.results.length > 0) return config.results;
-  return [];
-}
-
-function UIDispatcher({ rawContent, primaryColor, accentColor, isStreaming, onAddToCart, viewportSize = 'large' }: {
-  rawContent: string;
-  primaryColor?: string;
-  accentColor?: string;
-  isStreaming?: boolean;
-  onAddToCart?: (product: Product) => void;
-  viewportSize?: ChatViewportSize;
-}) {
-  const result = React.useMemo<{ config: UIConfig } | { error: string } | { loading: true }>(() => {
-    if (isStreaming) return { loading: true };
-    try {
-      const clean = rawContent
-        .replace(/^```[a-z]*\s*/i, '')
-        .replace(/```\s*$/, '')
-        .trim();
-      const parsed = JSON.parse(sanitizeJson(extractLikelyJsonBlock(clean)));
-      const config = { ...parsed } as Partial<RawUIConfig> & Record<string, unknown>;
-      
-      // V7 Field Mapping
-      if (config.chart) {
-        if (!config.chartType) config.chartType = config.chart.type;
-        if (!config.xAxisKey) config.xAxisKey = config.chart.xKey;
-        if (!config.dataKeys) config.dataKeys = config.chart.yKeys;
-      }
-      if (config.table) {
-        if (!config.columns) config.columns = config.table.columns;
-      }
-
-      const resolvedRows = resolveStructuredRows(config);
-
-      // ─── Intent Healing & Auto-Detection ───────────────────────────────────
-
-      const hasProductLikeData = resolvedRows.length > 0 &&
-        typeof resolvedRows[0] === 'object' && resolvedRows[0] !== null &&
-        !Array.isArray(resolvedRows[0]) &&
-        (['price', 'brand', 'image', 'link', 'thumbnail'].some(key => key in resolvedRows[0]) ||
-          Object.keys(resolvedRows[0]).length >= 4);
-
-      const hasAggregationLikeData = resolvedRows.length > 0 &&
-        typeof resolvedRows[0] === 'object' && resolvedRows[0] !== null &&
-        !Array.isArray(resolvedRows[0]) &&
-        (['value', 'count', 'total', 'percentage'].some(key => key in resolvedRows[0]) ||
-          (Object.keys(resolvedRows[0]).length <= 3 && Object.values(resolvedRows[0]).some(v => typeof v === 'number')));
-
-      // 1. Healing: Override explicit view if it contradicts the data shape
-      if (config.view === 'chart' && hasProductLikeData && !hasAggregationLikeData) {
-        config.view = 'carousel';
-      } else if (config.view === 'carousel' && hasAggregationLikeData && !hasProductLikeData) {
-        config.view = 'chart';
-        config.chartType = (config.chartType as 'pie' | 'bar' | 'line') || 'bar';
-      }
-
-      // 2. Auto-Detection: Fill in missing or dynamic view
-      if (!config.view || ['pie', 'bar', 'line', 'pie_chart', 'bar_chart', 'line_chart'].includes(String(config.view))) {
-        const viewStr = String(config.view || '').toLowerCase();
-        
-        if (viewStr.includes('carousel') || config.type === 'products' || Array.isArray(config.items) || hasProductLikeData) {
-          config.view = 'carousel';
-        } else if (
-          viewStr.includes('chart') ||
-          viewStr.includes('pie') || viewStr.includes('bar') || viewStr.includes('line') ||
-          hasAggregationLikeData ||
-          (typeof config.type === 'string' && ['pie', 'bar', 'line'].includes(config.type)) ||
-          (typeof config.chartType === 'string' && ['pie', 'bar', 'line'].includes(config.chartType))
-        ) {
-          config.view = 'chart';
-          config.chartType = (config.chartType as 'pie' | 'bar' | 'line') ||
-            (viewStr.replace('_chart', '') as 'pie' | 'bar' | 'line') ||
-            (config.type as 'pie' | 'bar' | 'line') || 'bar';
-        } else if (viewStr.includes('table')) {
-          config.view = 'table';
-        } else {
-          config.view = 'table';
-        }
-      }
-
-      const normalizedConfig: UIConfig = {
-        ...(config as Omit<UIConfig, 'data'>),
-        data: normalizeTabularRows(resolvedRows, config.columns),
-      };
-
-      return { config: normalizedConfig };
-    } catch (err) {
-      return { error: String(err) };
-    }
-  }, [rawContent, isStreaming]);
-
-  if ('loading' in result) {
-    return (
-      <div className="my-4 p-8 bg-slate-50 dark:bg-white/5 rounded-xl border border-dashed border-slate-300 dark:border-white/10 flex flex-col items-center justify-center gap-3 animate-pulse">
-        <div className="w-5 h-5 border-2 border-indigo-500 border-t-transparent rounded-full animate-spin" />
-        <p className="text-xs text-slate-500 font-medium italic">Preparing view...</p>
-      </div>
-    );
-  }
-
-  if ('error' in result) {
-    return <pre className="p-4 my-2 bg-slate-900 text-slate-100 rounded-lg text-[10px] overflow-auto max-h-40"><code>{rawContent}</code></pre>;
-  }
-
-  const { config } = result;
-  const isCompact = viewportSize === 'compact';
-  const hasInsights = Array.isArray(config.insights) && config.insights.length > 0;
-  const hasDescription = typeof config.description === 'string' && config.description.trim().length > 0;
-
-  switch (config.view) {
-    case 'chart':
-      return (
-        <div className={`${isCompact ? 'my-4 p-3' : 'my-6 p-4'} bg-white dark:bg-slate-900 rounded-2xl border border-slate-200 dark:border-white/10 shadow-sm overflow-hidden`}>
-          {config.title && <h4 className={`${isCompact ? 'text-[11px] mb-3' : 'text-xs mb-4'} font-semibold text-slate-500 px-2`}>{config.title}</h4>}
-          {hasDescription && (
-            <p className={`px-2 ${isCompact ? 'text-xs' : 'text-sm'} text-slate-500 dark:text-white/60`}>{config.description}</p>
-          )}
-          <DynamicChart
-            config={{
-              type: (config.chartType || 'bar') as 'bar' | 'line' | 'pie',
-              data: config.data as unknown as Record<string, string | number>[],
-              xAxisKey: config.xAxisKey,
-              dataKeys: config.dataKeys,
-              colors: config.colors,
-            }}
-            primaryColor={primaryColor}
-            accentColor={accentColor}
-            viewportSize={viewportSize}
-          />
-          {hasInsights && (
-            <div className="mt-4 flex flex-wrap gap-2 px-2">
-              {config.insights?.map((insight, i) => (
-                <span
-                  key={`insight-${i}`}
-                  className={`rounded-full border border-emerald-200 bg-emerald-50 ${isCompact ? 'px-2.5 py-1 text-[10px]' : 'px-3 py-1 text-[11px]'} font-medium text-emerald-700 dark:border-emerald-500/20 dark:bg-emerald-500/10 dark:text-emerald-300`}
-                >
-                  {String(insight)}
-                </span>
-              ))}
-            </div>
-          )}
-        </div>
-      );
-    case 'carousel':
-      return (
-        <div className="my-4">
-          {config.title && <h4 className={`${isCompact ? 'mb-1.5 text-[11px]' : 'mb-2 text-xs'} font-semibold text-slate-500`}>{config.title}</h4>}
-          {hasDescription && (
-            <p className={`mb-3 ${isCompact ? 'text-xs' : 'text-sm'} text-slate-500 dark:text-white/60`}>{config.description}</p>
-          )}
-          <ProductCarousel products={(config.data as unknown as Product[]).map(item => ({
-            ...item,
-            image: item.image ?? (typeof resolveImage === 'function' ? resolveImage(item as unknown as Record<string, unknown>) : undefined)
-          }))} primaryColor={primaryColor} onAddToCart={onAddToCart} />
-        </div>
-      );
-    case 'table':
-      return (
-        <div className={`${isCompact ? 'my-3 p-3 max-h-[320px]' : 'my-4 p-4 max-h-[400px]'} bg-white dark:bg-slate-900 rounded-xl border border-slate-200 dark:border-white/10 shadow-sm overflow-auto`}>
-          {config.title && <h4 className={`${isCompact ? 'text-[11px]' : 'text-xs'} font-semibold text-slate-500 mb-2`}>{config.title}</h4>}
-          {hasDescription && (
-            <p className={`mb-3 ${isCompact ? 'text-xs' : 'text-sm'} text-slate-500 dark:text-white/60`}>{config.description}</p>
-          )}
-          <DataTable config={{
-            type: 'table',
-            title: config.title,
-            description: config.description,
-            columns: config.columns,
-            dataKeys: config.dataKeys,
-            xAxisKey: config.xAxisKey,
-            data: config.data,
-          }} viewportSize={viewportSize} />
-        </div>
-      );
-    default:
-      return null;
-  }
-}
-
-
+import { ObservabilityPanel } from './ObservabilityPanel';
+import HourglassLoader from './HourglassLoader';
+import {
+  extractStructuredPayload,
+  isLikelyUiOnlyMessage,
+  extractProductsFromSources,
+  extractProductsFromContent,
+  deduplicateProducts,
+  looksLikeStructuredPayload,
+  stripStructuredUiLabels,
+  stripMarkdownTables
+} from '../utils/ProductExtractor';
+import { createMarkdownComponents } from './MarkdownComponents';
+import { UIDispatcher } from './UIDispatcher';
 
 // ─── Main component ───────────────────────────────────────────────────────────
 
@@ -660,118 +36,57 @@ export function MessageBubble({
   const isUser = message.role === 'user';
   const isCompact = viewportSize === 'compact';
   const isMedium = viewportSize === 'medium';
+
   const [showSources, setShowSources] = React.useState(false);
-  const hasStructuredProductBlock = React.useMemo(
-    () => /```(?:ui|json|chart)\s*[\s\S]*?"?(?:view|type|data|products)"?\s*:\s*"?(?:carousel|products|table)"?/i.test(message.content) ||
-      (/\{[\s\S]*?"?(?:view|type|data|products)"?\s*:\s*"?(?:carousel|products)"?[\s\S]*\}/i.test(message.content) && looksLikeStructuredPayload(message.content)),
-    [message.content],
-  );
+  const [showTrace, setShowTrace] = React.useState(false);
+
   const structuredContent = React.useMemo(
     () => isUser ? { payload: null, text: message.content } : extractStructuredPayload(message.content),
     [isUser, message.content],
   );
-  const shouldRenderStructuredOnly = !isUser && Boolean(structuredContent.payload) && isLikelyUiOnlyMessage(structuredContent.text);
 
-  // ── Products from sources ──────────────────────────────────────────────────
+  const hasRichUI = React.useMemo(() => {
+    if (message.uiTransformation) {
+      const type = (message.uiTransformation as UITransformationResponse).type;
+      if (type && !['text', 'table'].includes(type)) return true;
+    }
+    if (structuredContent.payload) {
+      return /"?(?:view|type|chartType)"?\s*:\s*"?(?:chart|carousel|pie|bar|line|product_carousel)"?/i.test(structuredContent.payload);
+    }
+    return false;
+  }, [message.uiTransformation, structuredContent.payload]);
+
+  const shouldRenderStructuredOnly = !isUser && hasRichUI && isLikelyUiOnlyMessage(structuredContent.text);
+
+  // ── Extraction ─────────────────────────────────────────────────────────────
+
   const productsFromSources = React.useMemo<Product[]>(() => {
-    if (isUser || !sources) return [];
-    return sources
-      .filter(s => {
-        const m = s.metadata ?? {};
-        return m.price || m.image || m.img || m.thumbnail || m.images || m.brand || m.type === 'product';
-      })
-      .map(s => {
-        const m = (s.metadata ?? {}) as Record<string, unknown>;
-        return {
-          id: s.id,
-          name: (m.name ?? m.title ?? s.content.split('\n')[0] ?? 'Unknown Product') as string,
-          brand: m.brand as string,
-          price: m.price as string | number,
-          image: resolveImage(m),
-          link: m.link as string,
-          description: s.content,
-        };
-      });
+    return extractProductsFromSources(sources, isUser);
   }, [sources, isUser]);
 
-  // ── Products + cleaned content from structured JSON blocks in the message ──
   const { productsFromContent, cleanContent } = React.useMemo(() => {
     if (isUser) {
       return { productsFromContent: [] as Product[], cleanContent: message.content };
     }
 
-    const jsonRegex = /```(?:json|ui)?\s*([\s\S]*?)\s*```/g;
-    const products: Product[] = [];
-    let content = structuredContent.text;
-
-    const payloadCandidates = [structuredContent.payload, ...content.matchAll(jsonRegex).map((match) => match[1])].filter(Boolean) as string[];
+    const payloadCandidates = [
+      structuredContent.payload,
+      ...Array.from(structuredContent.text.matchAll(/```(?:json|ui)?\s*([\s\S]*?)\s*```/g)).map(m => m[1])
+    ].filter(Boolean) as string[];
 
-    for (const candidate of payloadCandidates) {
-      try {
-        const data = JSON.parse(sanitizeJson(candidate));
-        const itemSet =
-          Array.isArray(data.data) ? data.data
-            : Array.isArray(data.items) ? data.items
-              : Array.isArray(data.products) ? data.products
-                : null;
-
-        if (
-          (data.type === 'products' || data.view === 'carousel' || data.view === 'table' || Array.isArray(data.products)) &&
-          itemSet
-        ) {
-          products.push(
-            ...itemSet.map((item: Record<string, unknown>) => ({
-              ...item,
-              image: (item.image as string) ?? resolveImage(item),
-            })) as Product[],
-          );
-        }
-      } catch {
-        // Ignore malformed product payloads
-      }
-    }
-
-    if (content.includes('"type": "products"') || content.includes('"type":"products"')) {
-      for (const match of content.matchAll(jsonRegex)) {
-        try {
-          const data = JSON.parse(sanitizeJson(match[1]));
-          if (data.type === 'products' && Array.isArray(data.items)) {
-            products.push(
-              ...data.items.map((item: Record<string, unknown>) => ({
-                ...item,
-                image: (item.image as string) ?? resolveImage(item),
-              })) as Product[],
-            );
-            content = content.replace(match[0], '');
-          }
-        } catch {
-          // Malformed product block — skip silently
-        }
-      }
-    }
-
-    return { productsFromContent: products, cleanContent: content.trim() };
-  }, [message.content, isUser, structuredContent]);
+    return extractProductsFromContent(
+      message.content,
+      structuredContent.text,
+      payloadCandidates,
+      productsFromSources
+    );
+  }, [message.content, isUser, structuredContent, productsFromSources]);
 
-  // ── Merge & deduplicate products ───────────────────────────────────────────
   const allProducts = React.useMemo<Product[]>(() => {
-    if (productsFromContent.length > 0) return productsFromContent;
-
-    const seen = new Set<string>();
-    const contentLower = message.content.toLowerCase();
-
-    return productsFromSources.filter(p => {
-      const id = String(p.id ?? p.name);
-      if (seen.has(id)) return false;
-      const mentioned =
-        contentLower.includes(p.name.toLowerCase()) ||
-        (p.brand ? contentLower.includes(p.brand.toLowerCase()) : false);
-      if (mentioned) { seen.add(id); return true; }
-      return false;
-    });
+    return deduplicateProducts(productsFromSources, productsFromContent, message.content);
   }, [productsFromSources, productsFromContent, message.content]);
 
-  const processedMarkdown = React.useMemo(() => {
+  const processedMarkdown = React.useMemo<string>(() => {
     let raw = structuredContent.payload ? structuredContent.text : (cleanContent || message.content);
     const hasStructuredUiBlock = Boolean(structuredContent.payload) || /```(?:ui|json|chart)\s*[\s\S]*?"(?:view|chartType|type|data|items|rows|products|results)"\s*:/i.test(raw);
 
@@ -779,12 +94,9 @@ export function MessageBubble({
     raw = raw.replace(/([^\n])\n\|/g, '$1\n\n|') // Ensure blank line before table
       .replace(/(\|[^\n]*\|)\n\n+(?=\|)/g, '$1\n'); // Remove blank lines between rows
 
-    // 2. HEALING: Detect naked JSON charts/products and wrap them if markers are missing
-    // This handles models (like Llama 3.2) that often omit the ```chart markers.
+    // 2. HEALING: Wrap naked JSON blocks
     if (!isStreaming) {
-      // Look for objects that look like chart or product data
       raw = raw.replace(/(?:^|\n)\s*\{[\s\S]*\}\s*(?:\n|$)/g, (match) => {
-        // Only wrap if not already in a code block
         if (match.includes('```')) return match;
         if (!looksLikeStructuredPayload(match)) return match;
         const type = match.includes('"view"') || match.includes('"chartType"') || match.includes('"type": "pie"') || match.includes('"type":"pie"') || match.includes('"type": "bar"') || match.includes('"type":"bar"') || match.includes('"type": "line"') || match.includes('"type":"line"')
@@ -794,156 +106,19 @@ export function MessageBubble({
       });
     }
 
-    // 2b. If the model already provided a structured UI block, suppress duplicated
-    // markdown tables that often follow as a fallback explanation.
+    // 2b. Suppress duplicated markdown tables if a UI block is present
     if (hasStructuredUiBlock) {
       raw = stripStructuredUiLabels(stripMarkdownTables(raw))
         .replace(/\n{3,}/g, '\n\n')
         .trim();
     }
 
-    // 3. HIJACKING: If streaming, hide table-like content to prevent raw markdown flashing.
-    if (isStreaming) {
-      raw = raw.replace(/((?:^|\n)\|.*)+/g, (match) => {
-        return `\n\n\`\`\`table-loading\n${match.trim()}\n\`\`\`\n\n`;
-      });
-    }
-
     return raw;
   }, [cleanContent, message.content, isStreaming, structuredContent]);
 
   // ── Markdown component overrides ───────────────────────────────────────────
   const markdownComponents = React.useMemo(
-    () => ({
-      // Wrap in not-prose so Tailwind Typography resets don't clobber our styles.
-      // prose applies display:block and padding:0 to table/th/td — not-prose opts out.
-      table: ({ ...props }: React.HTMLAttributes<HTMLTableElement>) => {
-        if (isStreaming) {
-          return (
-            <div className="my-5 p-6 bg-slate-50 dark:bg-white/5 rounded-xl border border-dashed border-slate-300 dark:border-white/10 flex items-center justify-center gap-3 select-none">
-              <div className="w-4 h-4 border-2 border-indigo-500 border-t-transparent rounded-full animate-spin" />
-              <p className="text-xs text-slate-500 font-medium italic animate-pulse">
-                Generating data table...
-              </p>
-            </div>
-          );
-        }
-
-        return (
-          <div className="not-prose overflow-hidden my-5 rounded-xl border border-slate-200 dark:border-white/10 shadow-sm bg-white dark:bg-slate-900/50">
-            <div className="overflow-x-auto">
-              <table
-                className="!table w-full text-left border-collapse min-w-[400px] text-sm"
-                {...props}
-              />
-            </div>
-          </div>
-        );
-      },
-      thead: ({ ...props }: React.HTMLAttributes<HTMLTableSectionElement>) => (
-        <thead
-          className="!table-header-group bg-slate-50 dark:bg-white/5 border-b border-slate-200 dark:border-white/10"
-          {...props}
-        />
-      ),
-      tbody: ({ ...props }: React.HTMLAttributes<HTMLTableSectionElement>) => (
-        <tbody className="!table-row-group divide-y divide-slate-100 dark:divide-white/5" {...props} />
-      ),
-      tr: ({ ...props }: React.HTMLAttributes<HTMLTableRowElement>) => (
-        <tr
-          className="!table-row hover:bg-slate-50/70 dark:hover:bg-white/5 transition-colors"
-          {...props}
-        />
-      ),
-      th: ({ children, ...props }: React.ThHTMLAttributes<HTMLTableCellElement>) => (
-        <th
-          className="!table-cell px-4 py-3 font-bold text-slate-700 dark:text-white/90 whitespace-nowrap text-[11px] uppercase tracking-wider text-left"
-          {...props}
-        >
-          {normaliseChild(children)}
-        </th>
-      ),
-      td: ({ children, ...props }: React.TdHTMLAttributes<HTMLTableCellElement>) => (
-        <td
-          className="!table-cell px-4 py-3 text-slate-600 dark:text-white/70 whitespace-nowrap text-sm"
-          {...props}
-        >
-          {normaliseChild(children)}
-        </td>
-      ),
-      code({
-        inline,
-        className,
-        children,
-        ...props
-      }: React.HTMLAttributes<HTMLElement> & { inline?: boolean }) {
-        const lang = /language-(\w+)/.exec(className ?? '')?.[1];
-
-        if (!inline && (lang === 'ui' || lang === 'chart')) {
-          return (
-            <UIDispatcher
-              rawContent={String(children ?? '').trim()}
-              primaryColor={primaryColor}
-              accentColor={accentColor}
-              isStreaming={isStreaming}
-              onAddToCart={onAddToCart}
-              viewportSize={viewportSize}
-            />
-          );
-        }
-
-        if (!inline && lang === 'json') {
-          const content = String(children ?? '').trim();
-          if (looksLikeStructuredPayload(content)) {
-            return (
-              <UIDispatcher
-                rawContent={content}
-                primaryColor={primaryColor}
-                accentColor={accentColor}
-                isStreaming={isStreaming}
-                onAddToCart={onAddToCart}
-                viewportSize={viewportSize}
-              />
-            );
-          }
-        }
-
-        if (!inline) {
-          const content = String(children ?? '').trim();
-          if (looksLikeStructuredPayload(content)) {
-            return (
-              <UIDispatcher
-                rawContent={content}
-                primaryColor={primaryColor}
-                accentColor={accentColor}
-                isStreaming={isStreaming}
-                onAddToCart={onAddToCart}
-                viewportSize={viewportSize}
-              />
-            );
-          }
-        }
-
-        if (!inline && lang === 'table-loading') {
-          return (
-            <div className="my-5 p-6 bg-slate-50 dark:bg-white/5 rounded-xl border border-dashed border-slate-300 dark:border-white/10 flex items-center justify-center gap-3 select-none">
-              <div className="w-4 h-4 border-2 border-indigo-500 border-t-transparent rounded-full animate-spin" />
-              <p className="text-xs text-slate-500 font-medium italic animate-pulse">
-                Generating data table...
-              </p>
-            </div>
-          );
-        }
-
-        return (
-          <code className={className} {...props}>
-            {typeof children === 'string' || typeof children === 'number'
-              ? children
-              : JSON.stringify(children)}
-          </code>
-        );
-      },
-    }),
+    () => createMarkdownComponents({ primaryColor, accentColor, isStreaming, onAddToCart, viewportSize }),
     [primaryColor, accentColor, isStreaming, onAddToCart, viewportSize],
   );
 
@@ -961,7 +136,7 @@ export function MessageBubble({
         {isUser ? <User className={`${isCompact ? 'w-3.5 h-3.5' : 'w-4 h-4'} text-white`} /> : <Bot className={`${isCompact ? 'w-3.5 h-3.5' : 'w-4 h-4'}`} />}
       </div>
 
-      <div className={`flex flex-col gap-1 ${isCompact ? 'max-w-[94%]' : isMedium ? 'max-w-[92%]' : 'max-w-[90%]'} ${isUser ? 'items-end' : 'items-start'}`}>
+      <div className={`flex flex-col gap-1 min-w-0 ${isCompact ? 'max-w-[94%]' : isMedium ? 'max-w-[92%]' : 'max-w-[90%]'} ${isUser ? 'items-end' : 'items-start'}`}>
         {/* Bubble */}
         <div
           className={`relative ${isCompact ? 'px-3 py-2.5 text-[13px]' : 'px-4 py-3 text-sm'} rounded-2xl leading-relaxed shadow-sm dark:shadow-lg ${isUser
@@ -992,27 +167,52 @@ export function MessageBubble({
               )}
 
               {!shouldRenderStructuredOnly && (
-                <ReactMarkdown remarkPlugins={[remarkGfm]} components={markdownComponents}>
+                <ReactMarkdown components={markdownComponents}>
                   {processedMarkdown}
                 </ReactMarkdown>
               )}
 
-              {/* Automated UI Fallback (if LLM failed to output structured block but UITransformer found data) */}
-              {!isUser && !structuredContent.payload && !!message.uiTransformation && (
-                <div className="mt-4 border-t border-slate-100 dark:border-white/5 pt-4">
-                  <VisualizationRenderer data={message.uiTransformation as UITransformationResponse} />
-                </div>
-              )}
-
-              {isStreaming && message.content && (
-                <span className="inline-block w-1.5 h-3.5 ml-1 bg-emerald-500 animate-pulse align-middle" />
+              {/* Fix: Always show streaming cursor when isStreaming is true */}
+              {isStreaming && (
+                <span className="inline-flex items-center gap-1.5 ml-2 text-emerald-500 animate-pulse align-middle text-xs font-medium">
+                  <HourglassLoader
+                    size={18}
+                    primaryColor={primaryColor}
+                    accentColor={accentColor}
+                    glow={false} />
+                  Thinking...
+                </span>
               )}
             </div>
           )}
         </div>
 
+        {/* Auto-generated visualization fallback */}
+        {(() => {
+          if (isUser || structuredContent.payload || !message.uiTransformation) return null;
+          const ui = message.uiTransformation as UITransformationResponse;
+          const textContent = (ui.data as { content?: string })?.content ?? '';
+          const shouldShow =
+            (ui.type === 'table' && allProducts.length === 0) ||
+            !['text', 'table'].includes(ui.type) ||
+            (ui.type === 'text' && !hasRichUI && allProducts.length === 0 && textContent &&
+              !processedMarkdown.toLowerCase().includes(textContent.trim().toLowerCase()));
+          if (!shouldShow) return null;
+          return (
+            <div className="w-full mt-3">
+              <VisualizationRenderer
+                data={ui}
+                primaryColor={primaryColor}
+                onAddToCart={onAddToCart}
+                className="rounded-3xl border border-slate-200/60 dark:border-white/10 bg-white/90 dark:bg-slate-900/80 p-4"
+                isStreaming={isStreaming}
+              />
+            </div>
+          );
+        })()}
+
         {/* Product Carousel */}
-        {!isUser && !hasStructuredProductBlock && allProducts.length > 0 && (
+        {!isUser && !hasRichUI && allProducts.length > 0 && (
           <div className="w-full mt-1">
             <ProductCarousel
               products={allProducts}
@@ -1022,26 +222,46 @@ export function MessageBubble({
           </div>
         )}
 
-        {/* Sources */}
-        {!isUser && sources && sources.length > 0 && (
-          <div className="w-full">
-            <button
-              onClick={() => setShowSources(s => !s)}
-              className="text-[11px] text-indigo-400 hover:text-indigo-300 transition-colors flex items-center gap-1 mt-1"
-            >
-              {showSources ? <ChevronDown className="w-3 h-3" /> : <ChevronRight className="w-3 h-3" />}
-              {sources.length} source{sources.length !== 1 ? 's' : ''} used
-            </button>
+        {/* Footer actions: Sources and Trace */}
+        {!isUser && (
+          <div className="flex items-center gap-4 mt-1 w-full flex-wrap">
+            {sources && sources.length > 0 && (
+              <button
+                onClick={() => setShowSources(s => !s)}
+                className="text-[11px] text-indigo-400 hover:text-indigo-300 transition-colors flex items-center gap-1"
+              >
+                {showSources ? <ChevronDown className="w-3 h-3" /> : <ChevronRight className="w-3 h-3" />}
+                {sources.length} source{sources.length !== 1 ? 's' : ''} used
+              </button>
+            )}
 
-            {showSources && (
-              <div className="mt-2 flex flex-col gap-2">
-                {sources.map((src, i) => (
-                  <SourceCard key={src.id} source={src} index={i} />
-                ))}
-              </div>
+            {message.trace && (
+              <button
+                onClick={() => setShowTrace(t => !t)}
+                className="text-[11px] text-emerald-500 hover:text-emerald-400 transition-colors flex items-center gap-1"
+              >
+                <Activity className="w-3 h-3" />
+                {showTrace ? 'Hide Trace' : 'Trace'}
+              </button>
             )}
           </div>
         )}
+
+        {/* Expandable areas */}
+        <div className="w-full flex flex-col gap-2 min-w-0">
+          {showSources && sources && sources.length > 0 && (
+            <div className="mt-1 flex flex-col gap-2">
+              {sources.map((src, i) => (
+                <SourceCard key={src.id} source={src} index={i} />
+              ))}
+            </div>
+          )}
+
+          {showTrace && message.trace && (
+            <ObservabilityPanel trace={message.trace} primaryColor={primaryColor} />
+          )}
+        </div>
+
       </div>
     </div>
   );