test_live_space.py

"""
Simulate judge evaluation against the live HF Space with LLM agent.
Connects to the deployed Space and runs all tasks using OpenAI API.
"""
import asyncio
import json
import os
import sys
import time
from typing import List, Optional

from dotenv import load_dotenv
from openai import OpenAI

from client import InvoiceGuardEnv
from models import (
    ActionType, DecisionType, ExceptionType, InvoiceGuardAction,
)

load_dotenv()

SPACE_URL = "https://piyush-mk-invoice-guard.hf.space"
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4.1-mini")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY") or ""

TASKS = [
    "task_1_clean_match",
    "task_2_partial_receipt",
    "task_3_price_variance",
    "task_4_duplicate_invoice",
    "task_5_mixed_discrepancy",
    "task_6_false_positive_duplicate",
    "task_7_retroactive_price",
    "task_8_split_invoice_pattern",
    "task_9_clean_from_risky_vendor",
    "task_10_rounding_false_alarm",
    "task_11_authorized_overship",
    "task_12_corrected_resubmission",
]

_MODELS_MAX_COMPLETION_TOKENS = {"gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano", "gpt-5", "gpt-5-mini", "gpt-5.1"}

def _tok_kwarg(limit=512):
    for p in _MODELS_MAX_COMPLETION_TOKENS:
        if MODEL_NAME.startswith(p):
            return {"max_completion_tokens": limit}
    return {"max_tokens": limit}

SYSTEM_PROMPT = """You are a senior accounts payable analyst. You will be given an invoice case to investigate and resolve.

The environment tells you your goal, available actions, and decision options. Read the goal carefully.

WORKFLOW:
1. Investigate: inspect documents (PO, GRN, vendor profile, policy rules), run comparisons (quantity, price, totals), check for duplicates.
2. Resolve: submit_final_resolution with your decision, exception type, evidence references, and explanation.

Complete a thorough investigation before resolving. Inspect at least: purchase order, goods receipt note, compare quantity, compare price, policy rules, duplicate check, and vendor profile.

RESPONSE FORMAT:
- Respond with ONLY a valid JSON object. No markdown, no commentary.
- Investigation example: {"action_type": "inspect_purchase_order"}
- Resolution example: {"action_type": "submit_final_resolution", "final_decision": "approve_for_payment", "exception_type": "clean_match", "evidence_references": ["inspect_purchase_order", "compare_quantity"], "explanation": "All documents match within tolerance."}

RULES:
- Pay close attention to POLICY findings -- they tell you when escalation is required.
- When multiple issues exist, escalation takes priority over hold.
- Check PO references carefully before concluding an invoice is a duplicate.
- Include all investigation actions you performed in evidence_references.
- Cite specific numbers in your explanation.
- NEVER repeat an action you already took.
- When remaining_steps is 3 or fewer, submit immediately with what you have.
"""


def build_observation_prompt(obs, is_first=False):
    parts = [
        f"Case: {obs.case_id} | Difficulty: {obs.difficulty} | Steps remaining: {obs.remaining_steps}",
        f"Invoice: {obs.invoice_summary}",
    ]
    if is_first and obs.goal:
        parts.append(f"\n{obs.goal}")
    if obs.revealed_documents:
        parts.append(f"Documents reviewed: {', '.join(obs.revealed_documents)}")
    if obs.findings:
        parts.append("Findings:")
        for i, f in enumerate(obs.findings, 1):
            parts.append(f"  {i}. {f}")
    if obs.last_action_result:
        parts.append(f"Last result: {obs.last_action_result}")
    if obs.warnings:
        parts.append(f"Warnings: {'; '.join(obs.warnings)}")
    if obs.remaining_steps <= 2:
        parts.append(">>> YOU MUST submit_final_resolution NOW. Decide based on what you have. <<<")
    return "\n".join(parts)


def parse_llm_response(text):
    text = text.strip()
    if "```json" in text:
        text = text.split("```json")[1].split("```")[0].strip()
    elif "```" in text:
        text = text.split("```")[1].split("```")[0].strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    for line in text.split("\n"):
        line = line.strip()
        if line.startswith("{"):
            try:
                return json.loads(line)
            except json.JSONDecodeError:
                continue
    return {"action_type": "summarize_findings"}


def build_action(params):
    action_type = params.get("action_type", "summarize_findings")
    try:
        ActionType(action_type)
    except ValueError:
        action_type = "summarize_findings"
    kwargs = {"action_type": action_type}
    if params.get("final_decision"):
        try:
            kwargs["final_decision"] = DecisionType(params["final_decision"])
        except ValueError:
            pass
    if params.get("exception_type"):
        try:
            kwargs["exception_type"] = ExceptionType(params["exception_type"])
        except ValueError:
            pass
    if params.get("evidence_references"):
        kwargs["evidence_references"] = list(params["evidence_references"])
    if params.get("explanation"):
        kwargs["explanation"] = str(params["explanation"])
    return InvoiceGuardAction(**kwargs)


async def run_task(env, llm, task_id):
    result = await env.reset(task_id=task_id)
    obs = result.observation
    obs.done = result.done
    obs.reward = result.reward
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    steps = 0
    rewards = []

    print(f"  [RESET] case={obs.case_id} difficulty={obs.difficulty} steps={obs.remaining_steps}")

    while not obs.done:
        user_msg = build_observation_prompt(obs, is_first=(steps == 0))
        messages.append({"role": "user", "content": user_msg})

        try:
            api_kwargs = {"model": MODEL_NAME, "messages": messages, "temperature": 0.0, **_tok_kwarg()}
            try:
                api_kwargs["response_format"] = {"type": "json_object"}
                response = llm.chat.completions.create(**api_kwargs)
            except Exception:
                del api_kwargs["response_format"]
                response = llm.chat.completions.create(**api_kwargs)
            assistant_msg = response.choices[0].message.content or ""
        except Exception as e:
            print(f"  [LLM ERROR] {e}")
            assistant_msg = '{"action_type": "summarize_findings"}'

        messages.append({"role": "assistant", "content": assistant_msg})

        params = parse_llm_response(assistant_msg)
        action = build_action(params)

        result = await env.step(action)
        obs = result.observation
        obs.done = result.done
        obs.reward = result.reward
        reward = result.reward or 0.0
        rewards.append(reward)
        steps += 1

        status = "ERR" if obs.last_action_error else "ok"
        print(f"  [STEP {steps:2d}] {action.action_type.value:35s} reward={reward:+.2f} remain={obs.remaining_steps} {status}")

    grader = getattr(obs, "grader_result", {}) or {}
    score = grader.get("score", 0.0) if isinstance(grader, dict) else 0.0
    decision = params.get("final_decision", "none")

    return {"task_id": task_id, "steps": steps, "score": score, "decision": decision, "rewards": rewards}


async def main():
    print("=" * 65)
    print("  JUDGE SIMULATION -- Live HF Space + LLM Agent")
    print(f"  Space:  {SPACE_URL}")
    print(f"  Model:  {MODEL_NAME}")
    print(f"  API:    {API_BASE_URL}")
    print("=" * 65)

    llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
    env = InvoiceGuardEnv(base_url=SPACE_URL)

    results = []
    try:
        async with env:
            for task_id in TASKS:
                print(f"\n--- {task_id} ---")
                start = time.time()
                try:
                    r = await run_task(env, llm, task_id)
                    r["time"] = time.time() - start
                    results.append(r)
                    print(f"  >> score={r['score']:.4f}  decision={r['decision']}  time={r['time']:.1f}s")
                except Exception as e:
                    print(f"  >> FAILED: {e}")
                    results.append({"task_id": task_id, "score": 0.0, "steps": 0, "error": str(e)})
    except Exception as e:
        print(f"\nConnection error: {e}")
        sys.exit(1)

    print(f"\n\n{'='*65}")
    print("  RESULTS SUMMARY")
    print(f"{'='*65}")
    scores = []
    for r in results:
        s = r.get("score", 0.0)
        scores.append(s)
        dec = r.get("decision", "n/a")
        print(f"  {r['task_id']:40s} score={s:.4f}  decision={dec}")
    avg = sum(scores) / len(scores) if scores else 0.0
    print(f"\n  Average score:  {avg:.4f}")
    print(f"  Tasks passed:   {sum(1 for s in scores if s >= 0.5)}/{len(scores)}")
    print(f"{'='*65}")


if __name__ == "__main__":
    asyncio.run(main())