backend/eval/depth.py

"""Depth tests for Create Direct PO — drive each variant from utterance #1 through full slot
fill, variant lock, and API fire. Assert the payload field-by-field.

These are stricter than the general eval suite — they catch regressions in the entire pipeline.
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

from backend.agent.kb import KnowledgeBase
from backend.agent.loop import AgentLoop
from backend.agent.resolver import resolve


@dataclass
class DepthCase:
    id: str
    description: str
    turns: list[str]
    expect_slots: dict[str, Any] = field(default_factory=dict)
    expect_apis_fired: list[str] = field(default_factory=list)
    expect_payload_fields: dict[str, Any] = field(default_factory=dict)
    expect_terminated: bool = False
    expect_variant_locked: bool = False
    expect_violations: Optional[bool] = None
    # A list of substrings — at end of run, NO violation's error_message may contain any of these.
    # Used to verify a specific forbidden-rule was correctly cleared after correction.
    forbidden_violations_must_not_contain: list[str] = field(default_factory=list)


@dataclass
class DepthResult:
    case_id: str
    passed: bool
    failures: list[str]
    final_slots: dict[str, Any]
    fired_apis: list[dict]
    terminated: bool
    variant_locked: bool


def _flat_payload(call: dict) -> dict:
    """Flatten {schema: {prop: value}} → {prop: value} for assertions."""
    out: dict[str, Any] = {}
    for schema_block in (call.get("payload") or {}).values():
        if isinstance(schema_block, dict):
            out.update(schema_block)
    return out


def _value_equiv(a, b) -> bool:
    sa = str(a).strip().rstrip(".,;:").strip().upper()
    sb = str(b).strip().rstrip(".,;:").strip().upper()
    if sa == sb:
        return True
    try:
        return float(sa) == float(sb)
    except (ValueError, TypeError):
        return False


def run_case(kb: KnowledgeBase, case: DepthCase) -> DepthResult:
    loop = AgentLoop(kb)
    failures: list[str] = []
    for i, utter in enumerate(case.turns, 1):
        try:
            loop.handle_turn(utter)
        except Exception as e:
            failures.append(f"turn {i} raised: {type(e).__name__}: {e}")
            break

    final_slots = dict(loop.state.slots)

    # Slot assertions
    for k, v in case.expect_slots.items():
        actual = final_slots.get(k.lower())
        if actual is None:
            failures.append(f"slot '{k}' missing (expected {v!r})")
        elif not _value_equiv(actual, v):
            failures.append(f"slot '{k}': expected {v!r}, got {actual!r}")

    # Terminated / locked
    if case.expect_terminated and not loop.state.terminated:
        failures.append(f"expected terminated=True, got {loop.state.terminated}")
    if case.expect_variant_locked and not loop.state.variant_locked:
        failures.append(f"expected variant_locked=True, got {loop.state.variant_locked}")

    # Violations
    journey_kb = kb.journey(loop.state.current_journey) if loop.state.current_journey else None
    if journey_kb is not None:
        r = resolve(journey_kb, loop.state.slots)
        if case.expect_violations is False and r.active_violations:
            v_summary = [v.error_message or v.rule_kind for v in r.active_violations[:3]]
            failures.append(f"unexpected violations at end: {v_summary}")
        # Specific-rule-must-be-gone assertions (e.g. "Adhoc Item class should be None for ...")
        for needle in case.forbidden_violations_must_not_contain:
            for v in r.active_violations:
                msg = v.error_message or ""
                if needle.lower() in msg.lower():
                    failures.append(f"forbidden-violation still firing: matched '{needle}' in '{msg[:120]}'")

    # APIs fired
    all_calls: list[dict] = []
    for fire in loop.state.fired_apis:
        all_calls.extend(fire.get("calls") or [])
    api_files = [c["api_file"] for c in all_calls]
    for expected_api in case.expect_apis_fired:
        if not any(expected_api in af for af in api_files):
            failures.append(f"API '{expected_api}' not fired (fired: {api_files})")

    # Payload field assertions on the main CreatePO call
    if case.expect_payload_fields:
        main_call = next((c for c in all_calls if "CreatePO" in c["api_file"]), None)
        if not main_call:
            failures.append("expected CreatePO call but none fired")
        else:
            flat = _flat_payload(main_call)
            for prop, expected in case.expect_payload_fields.items():
                if prop not in flat:
                    failures.append(f"payload missing property '{prop}' (expected {expected!r})")
                elif not _value_equiv(flat[prop], expected):
                    failures.append(f"payload[{prop}]: expected {expected!r}, got {flat[prop]!r}")

    return DepthResult(
        case_id=case.id,
        passed=not failures,
        failures=failures,
        final_slots=final_slots,
        fired_apis=[{"api_file": c["api_file"], "success": c.get("success")} for c in all_calls],
        terminated=loop.state.terminated,
        variant_locked=loop.state.variant_locked,
    )


# ============================================================================
# Depth case library
# ============================================================================

# Minimum slots we typically need to lock + fire a General PO:
#   potypeenum=1, supplier_code, buyerhdr, currencycode, item_codeml, qtyml, cost,
#   adhocitemclassml=NONE, podate
# Plus various defaults the system auto-fills.

CASES: list[DepthCase] = [
    DepthCase(
        id="depth_001_general_po_minimal",
        description="General PO with the minimum slot set; verify slot extraction over 4 turns.",
        turns=[
            "Create a general PO from supplier SUP-001, buyer Maria, currency USD",
            "Item ITEM-A, 10 units at $50, adhoc class NONE",
            "PO date is 2026-06-01",
            "Go ahead and create it",
        ],
        expect_slots={
            "potypeenum": "1",
            "supplier_code": "SUP-001",
            "buyerhdr": "Maria",
            "currencycode": "USD",
            "item_codeml": "ITEM-A",
            "qtyml": "10",
            "cost": "50",
            "adhocitemclassml": "NONE",
        },
        # Note: real master-data validation rules (e.g. "Item code does not exist in master") will
        # fire on synthetic data. We only assert that the slots are extracted correctly here.
    ),
    DepthCase(
        id="depth_002_capital_po_full_extraction",
        description="Capital PO: verify potype=2 + supplier + buyer + currency + qty extracted.",
        turns=[
            "Create a capital PO for CAPCO-99, buyer John, in EUR",
            "Item CNC-X, 1 unit at 85000",
            "PO date 2026-06-15, adhoc class NONE",
        ],
        expect_slots={
            "potypeenum": "2",
            "supplier_code": "CAPCO-99",
            "buyerhdr": "John",
            "currencycode": "EUR",
            "item_codeml": "CNC-X",
            "qtyml": "1",
            "cost": "85000",
            "adhocitemclassml": "NONE",
        },
    ),
    DepthCase(
        id="depth_003_dropship_po_with_customer",
        description="Dropship PO requires customer; agent must elicit it and extract correctly.",
        turns=[
            "Dropship PO from supplier DROP-1 to customer CUST-MX-44",
            "Item ITEM-DROP, 5 units at 200 USD",
            "Buyer is Diego",
            "PO date 2026-07-10",
        ],
        expect_slots={
            "potypeenum": "3",
            "supplier_code": "DROP-1",
            "customercode_ml": "CUST-MX-44",
            "buyerhdr": "Diego",
            "currencycode": "USD",
            "qtyml": "5",
            "cost": "200",
        },
    ),
    DepthCase(
        id="depth_004_forbidden_combo_then_correct",
        description="User triggers consignment+adhoc; user corrects → the specific forbidden-combo rule must clear.",
        turns=[
            "Create a consignment PO for SUP-Z with adhoc class ML01",
            "Actually make the adhoc class NONE instead",
        ],
        expect_slots={
            "potypeenum": "4",
            "supplier_code": "SUP-Z",
            "adhocitemclassml": "NONE",
        },
        # The specific capital/consignment + adhoc rule must NOT be firing after the correction.
        forbidden_violations_must_not_contain=["Adhoc Item class should be None"],
    ),
    DepthCase(
        id="depth_005_correction_buyer",
        description="User changes buyer mid-flow.",
        turns=[
            "Create a PO for SUP-A, buyer Maria, currency USD",
            "Actually change buyer to John",
        ],
        expect_slots={
            "supplier_code": "SUP-A",
            "buyerhdr": "John",
            "currencycode": "USD",
        },
    ),
    DepthCase(
        id="depth_006_pivot_then_resume",
        description="User pivots away, then resumes the create flow.",
        turns=[
            "Create a PO for SUP-RESUME, buyer Sarah",
            "Wait — actually I need to approve something else first",
            "Never mind, let's go back to creating that PO",
        ],
        # After resume, original supplier + buyer should still be in state.
        expect_slots={
            "supplier_code": "SUP-RESUME",
            "buyerhdr": "Sarah",
        },
    ),
    DepthCase(
        id="depth_007_mid_flow_question",
        description="User asks a question mid-flow; state is preserved.",
        turns=[
            "Create a PO for supplier SUP-Q, buyer Maria, currency USD",
            "What does capital PO type mean exactly?",
            "OK — let's stick with general type then. Item ITEM-Q, 10 units at 25",
        ],
        expect_slots={
            "supplier_code": "SUP-Q",
            "buyerhdr": "Maria",
            "currencycode": "USD",
            "item_codeml": "ITEM-Q",
            "qtyml": "10",
            "cost": "25",
        },
    ),
    DepthCase(
        id="depth_008_multi_slot_dense",
        description="Heavy multi-slot extraction in one utterance.",
        turns=[
            "Create a general PO: supplier ACME-9, buyer Maria, USD currency, item ITEM-DENSE, "
            "qty 25, cost 12.50, PO date 2026-08-01, adhoc class NONE",
        ],
        expect_slots={
            "potypeenum": "1",
            "supplier_code": "ACME-9",
            "buyerhdr": "Maria",
            "currencycode": "USD",
            "item_codeml": "ITEM-DENSE",
            "qtyml": "25",
            "cost": "12.50",
            "adhocitemclassml": "NONE",
        },
    ),
    DepthCase(
        id="depth_009_no_slot_duplication",
        description="Regression: verify intent classifier doesn't leak `supplier_name`-style raw keys.",
        turns=[
            "Create a purchase order for Stationery Inc, buyer Maria, 100 reams of paper",
        ],
        expect_slots={
            "supplier_code": "Stationery Inc",
            "buyerhdr": "Maria",
            "qtyml": "100",
        },
    ),
    DepthCase(
        id="depth_010_pr_reference_extraction",
        description="User mentions a PR/SO number; verify it lands in refsono (the kb-canonical key).",
        turns=[
            "Create a PO from PR/2026/12345 for supplier SUP-PR",
        ],
        expect_slots={
            "supplier_code": "SUP-PR",
            "refsono": "PR/2026/12345",
        },
    ),
]


def run_all() -> list[DepthResult]:
    from concurrent.futures import ThreadPoolExecutor, as_completed
    kb = KnowledgeBase("PO")
    out: list[DepthResult] = []
    with ThreadPoolExecutor(max_workers=10) as ex:
        futs = {ex.submit(run_case, kb, c): c for c in CASES}
        for f in as_completed(futs):
            out.append(f.result())
    return sorted(out, key=lambda r: r.case_id)


def main() -> None:
    results = run_all()
    passed = sum(1 for r in results if r.passed)
    print(f"\n=== Depth tests: {passed}/{len(results)} passed ===\n")
    for r in results:
        status = "PASS" if r.passed else "FAIL"
        print(f"[{status}] {r.case_id}")
        if not r.passed:
            for fl in r.failures:
                print(f"    - {fl}")
        if r.fired_apis:
            api_files = ", ".join(c["api_file"] for c in r.fired_apis)
            print(f"    APIs: {api_files}")
        print(f"    terminated={r.terminated}, locked={r.variant_locked}")
        print()
    # Write report
    report = {
        "passed": passed,
        "total": len(results),
        "pass_rate": passed / max(len(results), 1),
        "cases": [
            {
                "case_id": r.case_id,
                "passed": r.passed,
                "failures": r.failures,
                "final_slots": r.final_slots,
                "fired_apis": r.fired_apis,
                "terminated": r.terminated,
                "variant_locked": r.variant_locked,
            }
            for r in results
        ],
    }
    Path("eval_cases/depth_results.json").write_text(json.dumps(report, indent=2))
    print(f"Report written to eval_cases/depth_results.json")


if __name__ == "__main__":
    main()