RAG_Products/knowledge.py

"""Add products to the knowledge base (FAISS index).

Two paths:
  1. add_single_product(...)  — one product via a form.
  2. ingest_table(file, ...)  — an Excel/CSV upload whose headers may not match
     our schema. An LLM maps the arbitrary headers to our canonical fields
     (with a deterministic heuristic fallback); the caller can preview the
     mapping before committing.

Canonical fields: item_name, brand, variant, mrp, best_buy, bulk_price.

New products are embedded and merged into the existing FAISS index incrementally
(no full rebuild), then the in-memory catalog caches are invalidated.
"""
import csv
import io
import json
import re
import uuid
from datetime import datetime, timezone

from langchain_core.documents import Document

import RAG_Products.chat as _chat
from RAG_Products.categorize import categorize
from RAG_Products.config import VECTOR_STORE_DIR
from RAG_Products.loader import _clean_name, _parse_price
from RAG_Products.models import get_vector_db

CANONICAL = ["item_name", "brand", "variant", "mrp", "best_buy", "bulk_price"]

# heuristic header synonyms (fallback when no LLM / LLM fails)
_SYNONYMS = {
    "item_name": ["item name", "item", "product", "name", "title", "description", "model name"],
    "brand": ["brand", "make", "company", "manufacturer", "brand name"],
    "variant": ["variant", "model", "sku", "variant name", "type"],
    "mrp": ["mrp", "max retail", "maximum retail", "list price", "retail price"],
    "best_buy": ["best buy", "selling", "sale price", "our price", "offer", "price"],
    "bulk_price": ["bulk", "wholesale", "dealer", "bulk price"],
}


# ---------------------------------------------------------------------------
# Document construction + persistence
# ---------------------------------------------------------------------------
def _make_doc(name, brand="", variant="", mrp=None, best_buy=None, bulk_price=None):
    name = _clean_name(str(name or ""))
    if not name:
        return None
    brand = str(brand or "").strip()
    variant = str(variant or "").strip()
    mrp = _parse_price(str(mrp)) if mrp not in (None, "") else None
    best_buy = _parse_price(str(best_buy)) if best_buy not in (None, "") else None
    bulk_price = _parse_price(str(bulk_price)) if bulk_price not in (None, "") else None

    parts = [name]
    if brand:
        parts.append(f"Brand: {brand}")
    if variant and variant.lower() != name.lower():
        parts.append(f"Variant: {variant}")

    meta = {
        "id": uuid.uuid4().hex,
        "name": name,
        "brand": brand,
        "brand_key": brand.lower().strip(),
        "variant": variant,
        "mrp": mrp,
        "best_buy": best_buy,
        "bulk_price": bulk_price,
        "price": best_buy or mrp,
        "ingested_on": datetime.now(timezone.utc).isoformat(),
        "category": categorize(name, variant),
    }
    return Document(page_content=" | ".join(parts), metadata=meta)


def _persist(docs):
    """Embed + merge docs into the FAISS index, save, and refresh caches."""
    docs = [d for d in docs if d]
    if not docs:
        return 0
    db = get_vector_db()
    db.add_documents(docs)              # embeds via the store's embedding model
    db.save_local(str(VECTOR_STORE_DIR))
    _chat._catalog = None               # invalidate cached catalog + stock set
    _chat._stock_categories = None
    return len(docs)


def add_single_product(name, brand="", variant="", mrp=None, best_buy=None,
                       bulk_price=None):
    doc = _make_doc(name, brand, variant, mrp, best_buy, bulk_price)
    if not doc:
        raise ValueError("Product name is required.")
    _persist([doc])
    return doc.metadata


# ---------------------------------------------------------------------------
# Tabular upload (Excel / CSV) with LLM header mapping
# ---------------------------------------------------------------------------
def read_table(file_bytes, filename):
    """Return (headers, rows) where rows is a list of dicts keyed by header."""
    name = (filename or "").lower()
    if name.endswith((".xlsx", ".xlsm", ".xls")):
        from openpyxl import load_workbook

        wb = load_workbook(io.BytesIO(file_bytes), read_only=True, data_only=True)
        ws = wb.active
        rows_iter = ws.iter_rows(values_only=True)
        headers = [str(h).strip() if h is not None else f"col{i}"
                   for i, h in enumerate(next(rows_iter, []))]
        rows = []
        for r in rows_iter:
            if r is None or all(c is None for c in r):
                continue
            rows.append({headers[i]: r[i] if i < len(r) else None
                         for i in range(len(headers))})
        return headers, rows

    # CSV / text
    text = file_bytes.decode("utf-8-sig", errors="replace")
    reader = csv.DictReader(io.StringIO(text))
    headers = reader.fieldnames or []
    return [h.strip() for h in headers], list(reader)


def _heuristic_mapping(headers):
    mapping = {c: None for c in CANONICAL}
    lowered = {h: h.lower().strip() for h in headers}
    for canon, syns in _SYNONYMS.items():
        for h, hl in lowered.items():
            if any(s in hl for s in syns):
                mapping[canon] = h
                break
    return mapping


def _heuristic_is_confident(mapping):
    """Heuristic is good enough if it found the name and at least one price."""
    return bool(mapping.get("item_name")) and any(
        mapping.get(f) for f in ("best_buy", "mrp", "bulk_price"))


def map_headers_llm(headers, sample_rows):
    """Map source headers -> canonical fields.

    HEURISTIC-FIRST: standard headers map instantly with no LLM. The LLM is
    called only when the heuristic can't confidently identify the columns
    (i.e. genuinely non-matching headings) — and only if a key is set.
    """
    from RAG_Products.models import llm_available, llm_complete

    heuristic = _heuristic_mapping(headers)
    if _heuristic_is_confident(heuristic) or not llm_available():
        return heuristic, "heuristic"

    sample = sample_rows[:3]
    prompt = (
        "Map spreadsheet columns to a product schema for a camera-gear store.\n"
        f"Source headers: {headers}\n"
        f"Sample rows: {json.dumps(sample, default=str)[:1500]}\n"
        f"Target fields: {CANONICAL} "
        "(item_name=product name, mrp=max retail price, best_buy=our selling "
        "price, bulk_price=wholesale).\n"
        "Return ONLY a JSON object mapping each target field to the matching "
        "source header string, or null if none. No prose."
    )
    try:
        raw = llm_complete(prompt)
        m = re.search(r"\{.*\}", raw, flags=re.DOTALL)
        parsed = json.loads(m.group(0) if m else raw)
        mapping = {c: (parsed.get(c) if parsed.get(c) in headers else None)
                   for c in CANONICAL}
        if not mapping.get("item_name"):       # LLM missed the essential field
            mapping["item_name"] = heuristic["item_name"]
        return mapping, "llm"
    except Exception:
        return heuristic, "heuristic"


def _apply(row, mapping):
    def val(field):
        h = mapping.get(field)
        return row.get(h) if h else None
    return {
        "name": val("item_name"),
        "brand": val("brand"),
        "variant": val("variant"),
        "mrp": val("mrp"),
        "best_buy": val("best_buy"),
        "bulk_price": val("bulk_price"),
    }


def preview_table(file_bytes, filename):
    """Detect headers + proposed mapping + a few mapped sample rows."""
    headers, rows = read_table(file_bytes, filename)
    mapping, source = map_headers_llm(headers, rows)
    sample = [_apply(r, mapping) for r in rows[:5]]
    return {
        "headers": headers,
        "mapping": mapping,
        "mapping_source": source,
        "sample": sample,
        "row_count": len(rows),
    }


def ingest_table(file_bytes, filename, mapping=None):
    """Ingest the whole file. Uses `mapping` if given (e.g. user-corrected),
    else auto-detects. Returns a summary."""
    headers, rows = read_table(file_bytes, filename)
    source = "provided"
    if not mapping:
        mapping, source = map_headers_llm(headers, rows)
    if not mapping.get("item_name"):
        raise ValueError("Could not identify a product-name column.")

    docs, skipped = [], 0
    for r in rows:
        doc = _make_doc(**_apply(r, mapping))
        if doc:
            docs.append(doc)
        else:
            skipped += 1
    added = _persist(docs)
    return {
        "added": added,
        "skipped": skipped,
        "row_count": len(rows),
        "mapping": mapping,
        "mapping_source": source,
    }