agents.py

"""
OmniParse AI — Agentic AI Pipeline
Multi-stage extraction with specialized sub-agents:
  1. OCR Agent     — reads raw pixels → text
  2. Extractor     — structured extraction via LLM
  3. Validator     — cross-field validation, duplicate check
  4. Chat Agent    — natural language queries over invoice corpus
All sub-agents degrade gracefully if their dependencies are missing.
"""

import re, json, time, uuid
from datetime import datetime
from config import GROQ_API_KEY, HF_TOKEN

# ── Lazy Groq client ────────────────────────────────────────────────────────
_groq = None
def _groq_client():
    global _groq
    if _groq is None and GROQ_API_KEY:
        try:
            from groq import Groq
            _groq = Groq(api_key=GROQ_API_KEY)
        except ImportError: pass
    return _groq

# ── Agent 1: OCR (external, handled by ocr.py) ──────────────────────────────
# ── Agent 2: Extraction ─────────────────────────────────────────────────────

EXTRACTION_SYSTEM = (
    "You are an invoice data extraction engine. Extract structured data from the "
    "raw OCR text of an invoice. Return ONLY a valid JSON object, no markdown, "
    "with exactly these keys: vendor (string), invoice_number (string), "
    "invoice_date (string, YYYY-MM-DD if possible), due_date (string, "
    "YYYY-MM-DD if possible), amount (number, subtotal before tax), "
    "vat_amount (number), total (number), currency (3-letter code), "
    "line_items (array of {description, quantity, unit_price, total}). "
    "Use null for any field you cannot determine. Never invent data."
)

def _safe_json(text: str) -> dict | None:
    if not text: return None
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if not m: return None
    try: return json.loads(m.group(0))
    except json.JSONDecodeError: return None

def extractor_groq(ocr_text: str) -> dict | None:
    """Sub-agent: Groq Llama-based structured extraction."""
    gc = _groq_client()
    if not gc: return None
    try:
        r = gc.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=[
                {"role":"system","content":EXTRACTION_SYSTEM},
                {"role":"user","content":ocr_text[:3000]}
            ],
            max_tokens=512, temperature=0.05, timeout=15,
        )
        return _safe_json(r.choices[0].message.content)
    except Exception as e:
        print(f"[AGENT] Groq extraction failed: {e}")
        return None

def extractor_hf(ocr_text: str) -> dict | None:
    """Sub-agent: HuggingFace Inference fallback."""
    if not HF_TOKEN: return None
    try:
        import requests
        url = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
        h = {"Authorization": f"Bearer {HF_TOKEN}"}
        prompt = f"<s>[INST] {EXTRACTION_SYSTEM}\n\n{ocr_text[:3000]} [/INST]"
        payload = {"inputs":prompt,"parameters":{"max_new_tokens":512,"temperature":0.05}}
        r = requests.post(url, headers=h, json=payload, timeout=45)
        if r.status_code == 503:
            time.sleep(25)
            r = requests.post(url, headers=h, json=payload, timeout=45)
        r.raise_for_status()
        d = r.json()
        text = d[0]["generated_text"] if isinstance(d,list) else str(d)
        return _safe_json(text)
    except Exception as e:
        print(f"[AGENT] HF extraction failed: {e}")
        return None

def extractor_regex(ocr_text: str) -> dict:
    """Sub-agent: Pure regex fallback when no LLM available."""
    def _re(pat, txt, g=1):
        m = re.search(pat, txt, re.IGNORECASE)
        return m.group(g) if m else None

    inv = _re(r"(?:invoice|inv)[#:\s]+([A-Z0-9\-]{4,24})", ocr_text)
    dates = re.findall(r"\d{1,2}[/.\-]\d{1,2}[/.\-]\d{4}", ocr_text)
    tm = _re(r"(?:total|amount due)[\s:$]+([0-9,\.]+)", ocr_text)
    vendor = None
    for line in ocr_text.splitlines():
        s = line.strip()
        if s and not s.lower().startswith(
            ("total","invoice","date","due","amount","subtotal","tax","vat")
        ):
            vendor = s; break
    try: tv = float(tm.replace(",","")) if tm else None
    except: tv = None
    return {
        "vendor":vendor, "invoice_number":inv,
        "invoice_date":dates[0] if dates else None,
        "due_date":dates[1] if len(dates)>1 else None,
        "amount":None, "vat_amount":None, "total":tv,
        "currency":"USD", "line_items":[],
    }

# ── Agent 3: Validator ──────────────────────────────────────────────────────

def _parse_date_any(s: str):
    for fmt in ("%Y-%m-%d","%d/%m/%Y","%m/%d/%Y","%d.%m.%Y","%d-%m-%Y"):
        try: return datetime.strptime(str(s), fmt)
        except: continue
    return None

def validator_crossfield(data: dict) -> list[str]:
    """Check arithmetic consistency, date logic, sanity."""
    warnings = []
    try:
        a = data.get("amount"); v = data.get("vat_amount"); t = data.get("total")
        if all(x is not None for x in (a, v, t)):
            if abs((float(a)+float(v))-float(t)) > 0.10:
                warnings.append("Subtotal + VAT does not match total.")
    except (TypeError,ValueError): pass
    try:
        d1 = data.get("invoice_date"); d2 = data.get("due_date")
        if d1 and d2:
            p1 = _parse_date_any(str(d1)); p2 = _parse_date_any(str(d2))
            if p1 and p2 and p2 < p1:
                warnings.append("Due date is before invoice date.")
    except (TypeError,ValueError): pass
    try:
        if data.get("total") is not None and float(data["total"]) < 0:
            warnings.append("Total amount cannot be negative.")
    except: pass
    vendor = data.get("vendor")
    if vendor and len(str(vendor).strip()) < 2:
        warnings.append("Vendor name appears invalid.")
    return warnings

# ── Agent 4: Chat Agent ─────────────────────────────────────────────────────

CHAT_SYSTEM = (
    "You are a financial assistant answering questions about the user's invoices. "
    "Here is their invoice data as JSON: {context}. "
    "Answer concisely based ONLY on this data. Never reveal the raw JSON. "
    "If asked about totals, compute carefully. Respond in plain English."
)

def chat_agent(message: str, invoices_json: str) -> str:
    """Answer natural-language questions about invoice data."""
    gc = _groq_client()
    if not gc:
        return "AI Chat is not configured. Please set GROQ_API_KEY."

    msg = message.strip()[:2000]
    ctx = invoices_json[:6000]

    try:
        r = gc.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=[
                {"role":"system","content":CHAT_SYSTEM.format(context=ctx)},
                {"role":"user","content":msg},
            ],
            max_tokens=400, temperature=0.2, timeout=15,
        )
        return r.choices[0].message.content
    except Exception as e:
        return f"AI is temporarily unavailable. Please try again later. ({e})"

# ── Agent 5: Duplicate Detector ─────────────────────────────────────────────

def duplicate_detector(
    vendor: str, total: float, existing_invoices: list[dict]
) -> bool:
    """
    Check if a vendor+total combination already exists in this month's invoices.
    """
    if not vendor or total is None:
        return False
    now = datetime.now()
    vl = vendor.strip().lower()
    for inv in existing_invoices:
        try:
            created = datetime.fromisoformat(inv.get("created_at",""))
        except (ValueError,KeyError):
            continue
        if created.year == now.year and created.month == now.month:
            iv = (inv.get("vendor") or "").strip().lower()
            it = inv.get("total")
            if iv == vl and it is not None and abs(float(it)-float(total)) < 0.01:
                return True
    return False

# ── Orchestrator: Full pipeline ─────────────────────────────────────────────

def run_extraction_pipeline(ocr_text: str, filename: str) -> dict:
    """
    Main agent orchestrator. Runs extraction → validation → sanitization.
    Returns a complete, sanitized invoice data dict.
    """
    # Stage 1: Extract with fallback chain
    if ocr_text.strip():
        data = extractor_groq(ocr_text) or extractor_hf(ocr_text) or extractor_regex(ocr_text)
        confidence = 0.95
    else:
        data = {
            "vendor":"Demo Vendor Inc.",
            "invoice_number":f"DEMO-{uuid.uuid4().hex[:6].upper()}",
            "invoice_date":datetime.now().strftime("%Y-%m-%d"),
            "due_date":None, "amount":100.0, "vat_amount":21.0,
            "total":121.0, "currency":"USD", "line_items":[],
        }
        confidence = 0.3

    data["filename"] = filename
    data["confidence"] = confidence

    # Stage 2: Validate
    warnings = validator_crossfield(data)
    data["warnings"] = warnings
    data["status"] = "review" if warnings else "done"

    # Stage 3: Sanitize (strip HTML, control chars, truncate)
    data = _sanitize_invoice(data)

    return data

# ── Sanitization ────────────────────────────────────────────────────────────

def _sanitize_string(v, max_len=500) -> str:
    if not isinstance(v, str): return ""
    clean = re.sub(r"<[^>]*>", "", v)
    clean = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]", "", clean)
    clean = re.sub(r"\s+", " ", clean).strip()
    return clean[:max_len]

def _sanitize_invoice(data: dict) -> dict:
    string_fields = (
        "vendor","invoice_number","invoice_date","due_date",
        "currency","filename","status",
    )
    for f in string_fields:
        if f in data and data[f] is not None:
            data[f] = _sanitize_string(str(data[f]))
    if "line_items" in data and isinstance(data["line_items"], list):
        for item in data["line_items"]:
            if isinstance(item, dict) and "description" in item:
                if item["description"] is not None:
                    item["description"] = _sanitize_string(str(item["description"]), 300)
    return data