llmplus.py

"""High-accuracy LLM logic: multi-pass extraction, field voting, math validation."""

from llm import DEFAULT_MODEL, run_single_extraction


# ── Internal helpers ───────────────────────────────────────────────────────────

def _safe_float(v) -> float:
    try:
        return float(v) if v is not None else 0.0
    except (TypeError, ValueError):
        return 0.0


def _vote_fields(candidates: list[dict]) -> dict:
    """
    Merge N extraction candidates into one authoritative result.
      - Numeric fields  : median across candidates (robust to outliers)
      - String fields   : majority vote
      - Nested objects  : most-complete candidate (fewest null values)
      - line_items      : candidate with the most items
    """
    from collections import Counter

    def vote_scalar(field: str):
        vals = [c.get(field) for c in candidates if c.get(field) is not None]
        if not vals:
            return None
        if all(isinstance(v, (int, float)) for v in vals):
            return sorted(vals)[len(vals) // 2]
        return Counter(vals).most_common(1)[0][0]

    result: dict = {}
    for f in ("payment_no", "date", "due_date", "currency",
              "subtotal", "tax", "shipping", "discount", "total", "notes"):
        result[f] = vote_scalar(f)

    for nested in ("vendor", "bill_to"):
        best = max(
            (c.get(nested) or {} for c in candidates),
            key=lambda d: sum(1 for v in (d or {}).values() if v is not None),
            default={},
        )
        result[nested] = best or None

    all_item_lists = [c.get("line_items") or [] for c in candidates]
    result["line_items"] = max(all_item_lists, key=len, default=[])

    return result


def _validate_math(data: dict) -> dict:
    """
    Check: sum(line_item.total) + tax + shipping - discount == grand total.

    If totals are missing entirely (e.g. last page was a blank footer and
    extract_totals returned nothing), derive subtotal and total from line items
    so the UI always has something meaningful to display.

    If totals exist but the absolute gap exceeds $0.01, append a reconciliation
    line item so the table always sums to the stated grand total.
    """
    items    = data.get("line_items") or []
    line_sum = sum(_safe_float(it.get("total")) for it in items)
    tax      = _safe_float(data.get("tax"))
    shipping = _safe_float(data.get("shipping"))
    discount = _safe_float(data.get("discount"))
    total    = _safe_float(data.get("total"))

    # Totals missing but line items present — derive from items
    if total == 0.0 and line_sum > 0:
        data["subtotal"] = round(line_sum, 2)
        data["total"]    = round(line_sum + tax + shipping - discount, 2)
        return data

    if total == 0.0:
        return data

    computed = line_sum + tax + shipping - discount
    gap = round(total - computed, 2)

    if abs(gap) > 0.01:
        label = "Missing / unaccounted items" if gap > 0 else "Over-credited / duplicate items"
        data["line_items"].append({
            "description": f"{label}  [reconciliation gap: {'+' if gap >= 0 else ''}{gap}]",
            "code":        None,
            "qty":         None,
            "unit":        None,
            "unit_price":  None,
            "total":       gap,
        })
        data["_reconciliation_gap"] = gap

    return data


# ── Public API ─────────────────────────────────────────────────────────────────

def extract_fields_high(
    text: str,
    model: str = DEFAULT_MODEL,
    runs: int = 3,
) -> dict:
    """
    High-accuracy extraction:
      1. Run the LLM `runs` times with a 40 k-char window.
         Temperature increases slightly each pass to encourage diversity.
      2. Vote across all successful candidates (median / majority / most-complete).
      3. Python math validation: if totals don't balance, add a reconciliation
         line item so the line-items table always ties to the grand total.

    The 40 k window covers most multi-page invoice PDFs (Docling already
    concatenates every page into one markdown string before this is called).
    """
    candidates: list[dict] = []

    for i in range(runs):
        try:
            temp = round(0.05 + i * 0.05, 2)   # 0.05 / 0.10 / 0.15
            candidates.append(
                run_single_extraction(text, model, truncate=40_000, temperature=temp, num_predict=12288)
            )
        except Exception:
            pass

    if not candidates:
        raise RuntimeError("All high-accuracy extraction passes failed.")

    merged = _vote_fields(candidates) if len(candidates) > 1 else candidates[0]
    return _validate_math(merged)