core/categorize.py

"""Understanding agent — the text model reasons over an extracted bill (Phase v2).

After the vision model extracts a bill, this step (MiniCPM4.1-8B via core.inference)
*understands* it: classifies the vendor, assigns a category to each line item and
an overall category, and writes a one-line summary. Deterministic repair fills a
missing total. Returns a clean, save-ready transaction + `_uncertain` flags so the
UI highlights only what (if anything) needs a human glance.
"""

from __future__ import annotations

import json
from typing import Any

from core import inference
from core.extract import reconcile, _coerce_number, _extract_json_object, _strip_fences

MAX_NEW_TOKENS = 384

CATEGORIES = [
    "Groceries", "Dining", "Cafe", "Transport", "Fuel", "Utilities", "Rent",
    "Shopping", "Clothing", "Electronics", "Health", "Pharmacy", "Personal Care",
    "Entertainment", "Subscriptions", "Education", "Travel", "Telecom",
    "Insurance", "Household", "Fees & Charges", "Gifts", "Other",
]
DEFAULT_CATEGORY = "Other"
_LOOKUP = {c.lower(): c for c in CATEGORIES}

SYSTEM_PROMPT = (
    "You are a meticulous bookkeeping assistant for a personal budget tracker. "
    "You are given a bill (vendor, line items, charges, total). Understand it and "
    "categorise it. Use ONLY these categories:\n"
    f"{', '.join(CATEGORIES)}.\n"
    "Guidance: a restaurant / dhaba / food court bill is Dining; a coffee/tea shop "
    "is Cafe; a supermarket / kirana / grocery store is Groceries; a chemist is "
    "Pharmacy; petrol/diesel is Fuel; cab/bus/metro is Transport; a tailor/clothes "
    "shop is Clothing. IMPORTANT: judge each item by the VENDOR's nature — at a "
    "restaurant, dishes/drinks like 'Misal Pav', '2 Course', 'House Wine' are ALL "
    "Dining (never Utilities/Bills/Health). Only use Utilities/Bills/Telecom/Rent "
    "for actual utility, telecom, rent or bill-payment vendors. When unsure, match "
    "the item to the overall category. Return ONLY one JSON object, no prose:\n"
    "{\n"
    '  "category": "<overall category for the whole bill>",\n'
    '  "item_categories": ["<one category per line item, same order>"],\n'
    '  "summary": "<one short sentence, <=14 words, describing the bill>"\n'
    "}\n"
    "item_categories MUST have exactly one entry per line item, in order."
)

RETRY_SUFFIX = (
    "\n\nReturn ONLY the JSON object: {\"category\": ..., \"item_categories\": [...], "
    "\"summary\": ...} with one item category per line item, chosen from the allowed list."
)


def _normalize(value: Any) -> str:
    if not isinstance(value, str):
        return DEFAULT_CATEGORY
    return _LOOKUP.get(value.strip().lower(), DEFAULT_CATEGORY)


def _build_prompt(record: dict[str, Any], items: list[dict[str, Any]]) -> str:
    cur = record.get("currency", "") or ""
    lines = [f"Vendor: {record.get('vendor','') or '(unknown)'}",
             f"Total: {record.get('total', 0)} {cur}".strip(), ""]
    lines.append("Line items:")
    if items:
        for i, it in enumerate(items, 1):
            qty = it.get("qty", 1)
            lines.append(f"{i}. {it.get('name','')} x{qty} = {it.get('amount',0)}")
    else:
        lines.append("(none — single payment)")
    charges = record.get("charges") or []
    if charges:
        lines.append("Charges: " + ", ".join(
            f"{c.get('label','')} {c.get('amount',0)}" for c in charges))
    lines.append("")
    lines.append(f"Give category + {len(items)} item_categories (in order) + summary.")
    return "\n".join(lines)


def _run_model(prompt: str) -> str:
    return inference.text_generate(
        [{"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": prompt}],
        max_new_tokens=MAX_NEW_TOKENS,
    )


def _parse(text: str, n_items: int) -> dict[str, Any]:
    data = json.loads(_extract_json_object(_strip_fences(text)))
    if not isinstance(data, dict):
        raise ValueError("not an object")
    cats = [_normalize(c) for c in (data.get("item_categories") or [])]
    if len(cats) < n_items:
        cats += [DEFAULT_CATEGORY] * (n_items - len(cats))
    cats = cats[:n_items]
    return {
        "category": _normalize(data.get("category")),
        "item_categories": cats,
        "summary": str(data.get("summary", "") or "").strip()[:160],
    }


def understand(record: dict[str, Any]) -> dict[str, Any]:
    """Reason over an extracted bill → categorised, summarised, repaired transaction."""
    rec = dict(record)
    rec.setdefault("charges", [])
    items = rec.get("line_items") or []

    # Deterministic repair: compute a missing total from items + charges.
    if items and _coerce_number(rec.get("total", 0)) == 0:
        rec["total"] = reconcile(rec)["expected_total"]

    parsed = None
    base = _build_prompt(rec, items)
    for prompt in (base, base + RETRY_SUFFIX):
        try:
            parsed = _parse(_run_model(prompt), len(items))
            break
        except Exception as e:  # pragma: no cover - model dependent
            print(f"[understand] parse failed: {e}")
    if parsed is None:
        parsed = {"category": DEFAULT_CATEGORY,
                  "item_categories": [DEFAULT_CATEGORY] * len(items), "summary": ""}

    rec["category"] = parsed["category"]
    rec["receipt_category"] = parsed["category"]  # back-compat for storage/analytics
    rec["line_items"] = [dict(it, category=c) for it, c in zip(items, parsed["item_categories"])]
    rec["understanding"] = parsed["summary"]

    recon = reconcile(rec)
    unc: list[str] = []
    if not str(rec.get("vendor", "")).strip():
        unc.append("vendor")
    if not str(rec.get("date", "")).strip():
        unc.append("date")
    if _coerce_number(rec.get("total", 0)) == 0:
        unc.append("total")
    if items and not recon["ok"]:
        unc.append("total")
    rec["_uncertain"] = sorted(set(unc))
    return rec