backend/app/rag.py

"""RAG / Knowledge Base layer.

The reference docs call for a vendor-master / business-rules knowledge base that an
agent retrieves from at runtime — instead of stuffing every rule into the prompt.
This module is the open-source analogue of UiPath's master-data lookups, built on
the LlamaIndex/pgvector pattern but kept dependency-light for the prototype:

  • A small in-memory KB of vendor records + extraction/validation rules.
  • Vector retrieval via sentence-transformers when installed; otherwise a
    token-overlap (BM25-ish) fallback so it works with zero extra deps.
  • Used at two points in the pipeline (the canonical RAG insertion points):
      1. enrich   — retrieve vendor context, fill vendor_id, add layout hints.
      2. validate — confirm the vendor exists in master data (vendor_known check).

In production this is swapped for LlamaIndex + pgvector over the customer's real
vendor master and a versioned rules KB — same interface, bigger corpus.
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field

# --- the seeded knowledge base ------------------------------------------------
# Each record mixes master data (vendor_id, currency, terms) with free-text
# layout/validation hints an LLM can use — exactly what the reference docs store.

VENDOR_KB: list[dict] = [
    {
        "vendor_name": "Acme Industrial Supplies", "vendor_id": "V-ACME-001",
        "currency": "USD", "payment_terms": "Net 30", "category": "industrial",
        "hint": "Vendor ACME: invoice number format INV-XXXX in the top-left; "
                "tax is always 10%; totals appear on the last line.",
    },
    {
        "vendor_name": "Acme Industrial", "vendor_id": "V-ACME-001",
        "currency": "USD", "payment_terms": "Net 30", "category": "industrial",
        "hint": "Acme purchase orders use PO-1004xx numbering; ship-to is a US address.",
    },
    {
        "vendor_name": "GlobalParts GmbH", "vendor_id": "V-GLOB-014",
        "currency": "EUR", "payment_terms": "Net 45", "category": "components",
        "hint": "GlobalParts invoices are in EUR with European number format "
                "(1.234,56); VAT 19%; line items start at row 5.",
    },
    {
        "vendor_name": "Initech Supplies", "vendor_id": "V-INIT-007",
        "currency": "USD", "payment_terms": "Net 15", "category": "office",
        "hint": "Initech multi-page POs; totals only on the final page.",
    },
    {
        "vendor_name": "Northwind Traders", "vendor_id": "V-NORT-022",
        "currency": "USD", "payment_terms": "Net 30", "category": "furniture",
        "hint": "Northwind scanned invoices; OCR often needed; INV-77xx numbering.",
    },
    {
        "vendor_name": "Wayne Enterprises", "vendor_id": "V-WAYN-003",
        "currency": "USD", "payment_terms": "Net 60", "category": "industrial",
        "hint": "Wayne invoices are table-dense; reconcile line-item sum to subtotal.",
    },
    {
        "vendor_name": "Stark Components", "vendor_id": "V-STAR-011",
        "currency": "USD", "payment_terms": "Net 30", "category": "electronics",
        "hint": "Stark invoices sometimes omit the explicit total — compute it.",
    },
]


def _norm(s: str) -> str:
    return re.sub(r"[^a-z0-9 ]", " ", (s or "").lower())


def _tokens(s: str) -> set[str]:
    return {t for t in _norm(s).split() if len(t) > 2}


def _doc_text(r: dict) -> str:
    return f"{r['vendor_name']} {r['category']} {r['currency']} {r['hint']}"


@dataclass
class KnowledgeBase:
    """Vendor-master KB backed by the persistent VectorStore (app/rag_store.py).

    Same public API as the original prototype (retrieve / match_vendor / backend),
    but now reads from a real on-disk vector DB. A deterministic name match runs
    first so exact vendor lookups are reliable regardless of embedding backend.
    """

    records: list[dict] = field(default_factory=lambda: list(VENDOR_KB))
    store: object = None  # VectorStore

    def __post_init__(self):
        if self.store is None:
            from .config import get_settings
            from .rag_store import VectorStore

            self.store = VectorStore(get_settings().rag_db_path)
        # idempotently seed the vendor-master collection
        self.store.seed("vendor_master", self.records,
                        text_key=_doc_text, ref_key="vendor_id")

    def backend(self) -> str:
        return self.store.backend

    def retrieve(self, query: str, k: int = 2) -> list[dict]:
        if not query:
            return []
        hits = self.store.search(query, k=max(k, 4), collection="vendor_master")
        out = []
        for h in hits:
            rec = h.get("metadata") or {}
            # name-match boost for direct hits
            score = h["score"]
            if _norm(rec.get("vendor_name", "")) in _norm(query) or \
                    _norm(query) in _norm(rec.get("vendor_name", "")):
                score += 1.0
            out.append({**rec, "_score": round(float(score), 3)})
        out.sort(key=lambda r: r["_score"], reverse=True)
        return [r for r in out[:k] if r["_score"] > 0]

    def match_vendor(self, vendor_name: str) -> dict | None:
        if not vendor_name:
            return None
        # 1) deterministic exact/substring match (robust)
        q = _norm(vendor_name)
        for r in self.records:
            n = _norm(r["vendor_name"])
            if n == q or n in q or q in n:
                return {**r, "_score": 1.0}
        # 2) vector fallback
        hits = self.retrieve(vendor_name, k=1)
        if hits and hits[0]["_score"] >= 0.5:
            return hits[0]
        return None


# module-level singleton
KB = KnowledgeBase()