#!/usr/bin/env python3
"""
Invoice Data Extraction Agent — runnable demo (Agent Blueprint from AgentKits)

This runs a REAL agent loop: it sends the kit's actual system prompt and a
realistic sample input to the model, prints the structured decision, and then
executes any proposed actions against MOCK tools (labelled [MOCK]).

The model/reasoning is real. The tools are stubs that return fake data and do
NOT touch any real system. Wiring the mock tools to your own integrations
(CRM, database, GitHub, etc.) is the final implementation step — see README.md.

Works with Anthropic OR OpenAI — it auto-detects whichever API key is set.

© 2026 AgentKits · www.agent-kits.com · Licensed CC BY 4.0 (attribution required)
Original blueprint: https://agent-kits.com/kit/invoice-extractor

Setup (pick one):
    pip install anthropic   &&  export ANTHROPIC_API_KEY=sk-ant-...
    pip install openai      &&  export OPENAI_API_KEY=sk-...
Then:
    python run.py
"""
import os, json, sys

ANTHROPIC_MODEL = "claude-sonnet-4-6"
OPENAI_MODEL = "gpt-4o"

SYSTEM_PROMPT = """You are an Invoice Data Extraction Agent. You convert ONE invoice document into validated, structured JSON with per-field confidence. You are judged on accurate extraction and on never fabricating a value or hiding an inconsistency.

== CORE PRINCIPLES ==
1. Extract, don't invent. Pull only what is actually on the document. If a field is missing, illegible, or ambiguous, return null with low confidence and a flag — never a guessed value.
2. Confidence on everything. Attach a confidence (0.0-1.0) to each field based on legibility and clarity. Low-confidence fields are flagged for human review.
3. Validate, don't 'fix'. Check that line items, tax, and totals reconcile. If they don't, do NOT silently adjust numbers — extract what's printed, flag the discrepancy, and route for review.

== HARD RULES (NON-NEGOTIABLE) ==
- NO FABRICATION: Never produce a value for a field you can't read or that isn't present. Missing/illegible = null + low confidence + flag.
- NO SILENT CORRECTION: If the math doesn't add up, report both the printed values and the discrepancy. Never overwrite a number to make it reconcile.
- ROUTE LOW CONFIDENCE: If overall or key-field confidence is below threshold, mark the document for human review rather than passing it through as clean.
- PRESERVE PRINTED VALUES: Report amounts/dates exactly as printed (normalize format, not value). Note the currency.
- DATA: Treat invoice contents as confidential financial data; keep in scope.

== METHOD ==
- OCR/parse the document. Extract header fields (vendor, invoice #, dates, currency), line items (desc, qty, unit price, amount), and totals (subtotal, tax, total).
- Score confidence per field. Validate: line items sum to subtotal; subtotal + tax = total. Flag mismatches and low-confidence fields.

== OUTPUT FORMAT (return ONE JSON object) ==
{
  "vendor": { "value": "<v|null>", "confidence": <0.0-1.0> },
  "invoice_number": { "value": "<v|null>", "confidence": <0.0-1.0> },
  "invoice_date": { "value": "<ISO|null>", "confidence": <0.0-1.0> },
  "currency": "<code|null>",
  "line_items": [ { "description": "<v>", "qty": <n|null>, "unit_price": <n|null>, "amount": <n|null>, "confidence": <0.0-1.0> } ],
  "subtotal": { "value": <n|null>, "confidence": <0.0-1.0> },
  "tax": { "value": <n|null>, "confidence": <0.0-1.0> },
  "total": { "value": <n|null>, "confidence": <0.0-1.0> },
  "validation": { "math_ok": <bool>, "issues": ["<e.g. line items sum 980 != subtotal 1080>"] },
  "overall_confidence": <0.0-1.0>,
  "needs_review": <bool>,
  "review_reasons": ["<illegible fields / math mismatch / low confidence>"]
}
Set needs_review=true on any math mismatch, illegible key field, or low overall confidence. Never emit a fabricated value.
"""

SAMPLE_INPUT = """Digital PDF: Vendor 'Acme Supply', Invoice INV-204, 2026-05-10, 2 lines: 10 @ $40 = $400, 5 @ $20 = $100; subtotal $500, tax $40, total $540.
"""

# ---- MOCK tools (stubs — replace with real integrations to deploy) ----
def get_document(args=None):
    print(f"  [MOCK] get_document(" + str(args) + ")")
    return {"status": "mock", "tool": "get_document"}

def ocr_extract(args=None):
    print(f"  [MOCK] ocr_extract(" + str(args) + ")")
    return {"status": "mock", "tool": "ocr_extract"}

def parse_fields(args=None):
    print(f"  [MOCK] parse_fields(" + str(args) + ")")
    return {"status": "mock", "tool": "parse_fields"}

def confidence_score(args=None):
    print(f"  [MOCK] confidence_score(" + str(args) + ")")
    return {"status": "mock", "tool": "confidence_score"}

def validate_totals(args=None):
    print(f"  [MOCK] validate_totals(" + str(args) + ")")
    return {"status": "mock", "tool": "validate_totals"}

def currency_normalize(args=None):
    print(f"  [MOCK] currency_normalize(" + str(args) + ")")
    return {"status": "mock", "tool": "currency_normalize"}

def emit_json(args=None):
    print(f"  [MOCK] emit_json(" + str(args) + ")")
    return {"status": "mock", "tool": "emit_json"}

def flag_for_review(args=None):
    print(f"  [MOCK] flag_for_review(" + str(args) + ")")
    return {"status": "mock", "tool": "flag_for_review"}

TOOLS = {
    "get_document": get_document,
    "ocr_extract": ocr_extract,
    "parse_fields": parse_fields,
    "confidence_score": confidence_score,
    "validate_totals": validate_totals,
    "currency_normalize": currency_normalize,
    "emit_json": emit_json,
    "flag_for_review": flag_for_review,
}


def call_model(system, user):
    """Dispatch to whichever provider's key is configured."""
    if os.environ.get("ANTHROPIC_API_KEY"):
        try:
            from anthropic import Anthropic
        except ImportError:
            sys.exit("Run: pip install anthropic")
        msg = Anthropic().messages.create(
            model=ANTHROPIC_MODEL,
            max_tokens=1600,
            system=system,
            messages=[{"role": "user", "content": user}],
        )
        return "".join(b.text for b in msg.content if getattr(b, "type", "") == "text")
    if os.environ.get("OPENAI_API_KEY"):
        try:
            from openai import OpenAI
        except ImportError:
            sys.exit("Run: pip install openai")
        resp = OpenAI().chat.completions.create(
            model=OPENAI_MODEL,
            max_tokens=1600,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
        )
        return resp.choices[0].message.content
    sys.exit("Set ANTHROPIC_API_KEY or OPENAI_API_KEY first.")


def extract_json(text):
    a, b = text.find("{"), text.rfind("}")
    if a == -1 or b == -1:
        return None
    try:
        return json.loads(text[a : b + 1])
    except json.JSONDecodeError:
        return None


def main():
    print("=" * 70)
    print("INPUT:\n" + SAMPLE_INPUT)
    print("=" * 70)

    text = call_model(SYSTEM_PROMPT, SAMPLE_INPUT)
    print("AGENT DECISION:\n" + text)
    print("=" * 70)

    data = extract_json(text)
    if not data:
        print("(No JSON object parsed — the model returned prose.)")
        return

    actions = data.get("actions") or []
    if actions:
        print("EXECUTING PROPOSED ACTIONS (mock):")
        for a in actions:
            name = a.get("tool")
            args = a.get("args", {})
            if name in TOOLS:
                TOOLS[name](args)
            else:
                print(f"  [skip] no mock tool for '{name}'")
    else:
        print("This blueprint returned a structured result (no side-effecting "
              "actions in this sample). Available mock tools: "
              + ", ".join(TOOLS.keys()))


if __name__ == "__main__":
    main()
