How to Test an AI Agent (End to End)

"""Agent tracing wrapper.

Wraps an agent so that every LLM call, tool call, and decision it makes is
captured into a structured, ordered trace. The trace is a plain list of step
dicts (JSON-serializable) that downstream tests can replay and assert against.

Design goals:
  * No external dependencies (standard library only).
  * Deterministic: steps are recorded in call order, no timestamps, so two runs
    of the same agent logic produce byte-identical traces.
  * Non-invasive: you decorate the functions the agent already calls.

Run it directly to see a small stub agent loop traced end to end:

    python src/tracing.py

It prints a JSON trace to stdout.
"""

from __future__ import annotations

import functools
import json
from typing import Any, Callable, Dict, List


class Tracer:
    """Collects an ordered list of steps an agent takes.

    Each step is a dict of the form::

        {"type": "llm" | "tool" | "decision",
         "name": <str>,
         "args": <json-serializable>,
         "output": <json-serializable>}

    Ordering is implicit in list position, so we never record timestamps. That
    keeps traces reproducible and diffable in tests.
    """

    def __init__(self) -> None:
        self.steps: List[Dict[str, Any]] = []

    def record(self, step_type: str, name: str, args: Any, output: Any) -> None:
        self.steps.append(
            {"type": step_type, "name": name, "args": args, "output": output}
        )

    def decision(self, name: str, chosen: Any, options: Any = None) -> None:
        """Record a branching decision the agent made (e.g. which tool to use)."""
        self.record("decision", name, {"options": options}, chosen)

    def wrap(self, step_type: str, fn: Callable) -> Callable:
        """Return a wrapped version of ``fn`` that records every invocation.

        ``step_type`` is "llm" for model calls and "tool" for tool calls.
        """

        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            output = fn(*args, **kwargs)
            # Positional args are captured as a list; kwargs merged in so the
            # recorded args are a single JSON object regardless of call style.
            recorded_args: Dict[str, Any] = dict(kwargs)
            if args:
                recorded_args["_positional"] = list(args)
            self.record(step_type, fn.__name__, recorded_args, output)
            return output

        return wrapper

    def to_json(self) -> str:
        return json.dumps(self.steps, indent=2, sort_keys=False)


# --- A tiny stub agent so the module is runnable on its own -----------------


def _fake_llm(prompt: str) -> str:
    """Stand-in for a real model call. Deterministic for reproducibility.

    Two intents: a planning turn (decide whether to use a tool) and a
    finalizing turn (write the user-facing reply once results are in).
    """
    lowered = prompt.lower()
    if "write the final reply" in lowered:
        # Finalizing turn: produce the user-facing answer.
        return "final:It is currently 18C in Paris."
    if "weather" in lowered:
        # Planning turn: this needs the weather tool.
        return "call_tool:get_weather"
    return "final:I can help with weather questions."


def _get_weather(city: str) -> Dict[str, Any]:
    """Stand-in tool. Returns a fixed payload so runs are deterministic."""
    table = {"Paris": 18, "Tokyo": 24}
    return {"city": city, "temp_c": table.get(city, 20), "unit": "celsius"}


def run_stub_agent(tracer: Tracer, user_message: str) -> str:
    """A minimal plan/act agent loop, fully traced.

    1. Ask the LLM what to do.
    2. If it wants a tool, record the decision and call the (wrapped) tool.
    3. Ask the LLM again to produce the final answer.
    """
    llm = tracer.wrap("llm", _fake_llm)
    get_weather = tracer.wrap("tool", _get_weather)

    plan = llm(f"User asks: {user_message}. What should I do?")

    if plan.startswith("call_tool:"):
        tool_name = plan.split(":", 1)[1]
        tracer.decision("route", chosen=tool_name, options=["get_weather", "answer_directly"])
        # In this stub we only support one tool.
        weather = get_weather("Paris")
        final = llm(f"Weather is {weather['temp_c']}C. Write the final reply.")
    else:
        tracer.decision("route", chosen="answer_directly", options=["get_weather", "answer_directly"])
        final = plan

    answer = final.split(":", 1)[1] if ":" in final else final
    tracer.record("decision", "final_answer", args={}, output=answer)
    return answer


def main() -> None:
    tracer = Tracer()
    run_stub_agent(tracer, "What's the weather in Paris?")
    print(tracer.to_json())


if __name__ == "__main__":
    main()

"""Deterministic evaluation of an agent's *structured* output.

When an agent is asked to emit machine-readable output (JSON with an enum, a
classification, a set of extracted fields), you do not need an LLM to grade it.
You assert on it like any other function's return value:

  * exact-match assertions on individual fields, and
  * jsonschema validation of the whole object (shape, types, allowed enum
    values, required keys).

This is the cheapest, most reliable layer of an agent test suite: fast,
deterministic, and free. Run with::

    pytest tests/test_structured_output.py
"""

import jsonschema
import pytest
from jsonschema import Draft7Validator

# The contract the agent's structured output must satisfy.
TRIAGE_SCHEMA = {
    "type": "object",
    "additionalProperties": False,
    "required": ["intent", "priority", "needs_human", "entities"],
    "properties": {
        "intent": {
            "type": "string",
            "enum": ["billing", "bug_report", "feature_request", "other"],
        },
        "priority": {"type": "string", "enum": ["low", "medium", "high"]},
        "needs_human": {"type": "boolean"},
        "entities": {
            "type": "array",
            "items": {"type": "string"},
        },
    },
}


@pytest.fixture
def agent_output():
    """A sample structured response, as an agent would return it.

    In a real suite this comes from invoking your agent with a fixed prompt.
    Here it is inlined so the test is self-contained and deterministic.
    """
    return {
        "intent": "billing",
        "priority": "high",
        "needs_human": True,
        "entities": ["invoice #4471", "double charge"],
    }


def test_output_matches_schema(agent_output):
    """The whole object conforms to the JSON schema (shape + enums + types)."""
    # validate() raises on failure with a readable error path.
    jsonschema.validate(instance=agent_output, schema=TRIAGE_SCHEMA)


def test_no_extra_or_missing_fields(agent_output):
    """additionalProperties=False means a hallucinated extra key fails."""
    polluted = dict(agent_output)
    polluted["hallucinated_field"] = "oops"
    errors = list(Draft7Validator(TRIAGE_SCHEMA).iter_errors(polluted))
    assert errors, "expected schema to reject an unexpected extra field"


def test_exact_field_values(agent_output):
    """Exact-match assertions on the fields that must be deterministic."""
    assert agent_output["intent"] == "billing"
    assert agent_output["priority"] == "high"
    assert agent_output["needs_human"] is True


def test_enum_violation_is_rejected(agent_output):
    """An out-of-enum value must fail validation, not slip through."""
    bad = dict(agent_output, priority="URGENT")  # not in the enum
    with pytest.raises(jsonschema.ValidationError):
        jsonschema.validate(instance=bad, schema=TRIAGE_SCHEMA)

"""LLM-as-judge harness for grading free-text agent responses.

Not everything an agent produces can be checked with an exact-match assertion.
For open-ended answers you grade against an explicit, multi-criteria rubric and
have a *judge model* return a numeric score. The test then asserts the score
clears a threshold.

Keys to making this reliable:
  * The rubric is explicit and enumerated. Vague rubrics produce noisy scores.
  * The judge is asked to return machine-readable JSON, not prose.
  * Temperature is 0 so the judge is as reproducible as the provider allows.

This file is safe to collect and run without an API key: it skips gracefully
when OPENAI_API_KEY is unset. With a key set::

    pytest tests/test_llm_judge.py
"""

import json
import os

import pytest

# The rubric. Each criterion is scored 0-2 by the judge; total is 0-10.
RUBRIC = """\
You are a strict evaluator. Score the ASSISTANT RESPONSE against the USER
QUESTION on each criterion below. For each, award an integer 0, 1, or 2:
  0 = fails the criterion, 1 = partially meets it, 2 = fully meets it.

Criteria:
  1. RELEVANCE   - directly answers the user's actual question.
  2. CORRECTNESS - factually accurate, no fabricated details.
  3. COMPLETENESS- covers the key points the question requires.
  4. GROUNDING   - claims are supported, no unsupported speculation.
  5. CLARITY     - well structured and unambiguous.

Return ONLY minified JSON of the form:
{"relevance":N,"correctness":N,"completeness":N,"grounding":N,"clarity":N}
"""

# Below this total (out of 10) the response is considered a regression.
SCORE_THRESHOLD = 8


def build_judge_messages(question: str, answer: str):
    """Construct the chat messages sent to the judge model."""
    return [
        {"role": "system", "content": RUBRIC},
        {
            "role": "user",
            "content": (
                f"USER QUESTION:\n{question}\n\n"
                f"ASSISTANT RESPONSE:\n{answer}\n\n"
                "Grade now and return only the JSON object."
            ),
        },
    ]


def judge_score(question: str, answer: str, model: str = "gpt-4o-mini") -> int:
    """Call the judge model and return a total score in the range 0-10.

    Requires OPENAI_API_KEY. Raises if the judge returns unparseable output so
    a broken judge fails loudly rather than silently passing.
    """
    from openai import OpenAI  # imported lazily so import errors don't break collection

    client = OpenAI()  # reads OPENAI_API_KEY from the environment
    resp = client.chat.completions.create(
        model=model,
        temperature=0,
        response_format={"type": "json_object"},
        messages=build_judge_messages(question, answer),
    )
    raw = resp.choices[0].message.content
    scores = json.loads(raw)
    return sum(int(scores[k]) for k in
               ("relevance", "correctness", "completeness", "grounding", "clarity"))


@pytest.mark.skipif(
    not os.getenv("OPENAI_API_KEY"),
    reason="OPENAI_API_KEY not set; skipping live LLM-judge test.",
)
def test_agent_answer_clears_threshold():
    question = "What HTTP status code should a health-check endpoint return when healthy?"
    # A good answer we expect the judge to score highly.
    answer = (
        "A healthy health-check endpoint should return HTTP 200 OK. When a "
        "dependency is down, return 503 Service Unavailable so load balancers "
        "and orchestrators can route traffic away from the instance."
    )
    score = judge_score(question, answer)
    assert score >= SCORE_THRESHOLD, (
        f"judge scored {score}/10, below threshold {SCORE_THRESHOLD}"
    )

"""Trajectory (trace replay) tests.

A response-level eval only sees the agent's final answer. A trajectory test
inspects *how* the agent got there: which tools it called, in what order, with
what arguments, and how it recovered from a failing tool.

We replay a captured trace (the same structure emitted by ``src/tracing.py``)
and assert on the sequence of steps. This catches regressions that a
final-answer check misses, e.g. the agent silently skipping a required tool, or
calling tools in the wrong order, or not retrying after a transient failure.

Run with::

    pytest tests/test_trajectory.py
"""

import pytest

# A sample recorded trace, shaped exactly like Tracer.steps in src/tracing.py.
# This one includes a tool failure followed by a successful retry (recovery).
RECORDED_TRACE = [
    {"type": "llm", "name": "_fake_llm",
     "args": {"_positional": ["User asks: refund order 4471. What should I do?"]},
     "output": "call_tool:lookup_order"},
    {"type": "decision", "name": "route",
     "args": {"options": ["lookup_order", "answer_directly"]},
     "output": "lookup_order"},
    # First tool call fails (simulated transient error).
    {"type": "tool", "name": "lookup_order",
     "args": {"order_id": "4471"},
     "output": {"ok": False, "error": "timeout"}},
    # Agent decides to retry rather than give up.
    {"type": "decision", "name": "retry",
     "args": {"reason": "timeout"},
     "output": "retry_lookup_order"},
    # Retry succeeds.
    {"type": "tool", "name": "lookup_order",
     "args": {"order_id": "4471"},
     "output": {"ok": True, "status": "shipped", "total": 42.0}},
    {"type": "tool", "name": "issue_refund",
     "args": {"order_id": "4471", "amount": 42.0},
     "output": {"ok": True, "refund_id": "rf_991"}},
    {"type": "decision", "name": "final_answer", "args": {},
     "output": "Your refund of $42.00 has been issued."},
]


def steps_of_type(trace, step_type):
    return [s for s in trace if s["type"] == step_type]


def tool_names(trace):
    return [s["name"] for s in trace if s["type"] == "tool"]


def test_tools_called_in_expected_order():
    """The agent must look up the order before issuing a refund."""
    names = tool_names(RECORDED_TRACE)
    assert names == ["lookup_order", "lookup_order", "issue_refund"]
    # lookup must precede refund (never refund an order we never looked up).
    assert names.index("issue_refund") > names.index("lookup_order")


def test_refund_amount_matches_looked_up_total():
    """The refund argument must be derived from the tool result, not invented."""
    tools = steps_of_type(RECORDED_TRACE, "tool")
    successful_lookup = next(
        s for s in tools if s["name"] == "lookup_order" and s["output"]["ok"]
    )
    refund = next(s for s in tools if s["name"] == "issue_refund")
    assert refund["args"]["amount"] == successful_lookup["output"]["total"]


def test_failure_then_retry_recovery():
    """A transient tool failure must be followed by a retry decision + success."""
    lookups = [s for s in RECORDED_TRACE
               if s["type"] == "tool" and s["name"] == "lookup_order"]
    assert len(lookups) == 2, "expected one failed lookup and one retry"
    assert lookups[0]["output"]["ok"] is False
    assert lookups[1]["output"]["ok"] is True

    decisions = [s["output"] for s in steps_of_type(RECORDED_TRACE, "decision")]
    assert "retry_lookup_order" in decisions, "agent should decide to retry on failure"


def test_no_refund_before_successful_lookup():
    """Guard against the dangerous ordering: refund issued off a failed lookup."""
    seen_successful_lookup = False
    for step in RECORDED_TRACE:
        if step["type"] == "tool" and step["name"] == "lookup_order" and step["output"]["ok"]:
            seen_successful_lookup = True
        if step["type"] == "tool" and step["name"] == "issue_refund":
            assert seen_successful_lookup, "refund issued before a successful lookup"

"""Majority-vote helper for non-deterministic agent runs.

Agents are stochastic: the same prompt can produce different phrasings, and an
occasional run is simply wrong. Asserting on a single run makes a test flaky.
Instead, run the agent N times and require at least K of N runs to clear a
similarity threshold against a reference answer. That converts a flaky signal
into a stable pass/fail without pretending the agent is deterministic.

Similarity here uses ``difflib.SequenceMatcher`` (standard library) so the
helper runs with no extra dependencies. Swap in a semantic embedding similarity
for production use; the K-of-N structure stays the same.

Run it directly to see a 4-of-5 majority pass::

    python lib/majority_vote.py
"""

from __future__ import annotations

import difflib
from typing import Callable, List


def similarity(candidate: str, reference: str) -> float:
    """Return a 0..1 similarity ratio between two strings.

    Uses a token-set Jaccard blended with difflib's sequence ratio so that both
    word overlap and ordering contribute. Deterministic and dependency-free.
    """
    a_tokens = set(candidate.lower().split())
    b_tokens = set(reference.lower().split())
    if not a_tokens or not b_tokens:
        jaccard = 0.0
    else:
        jaccard = len(a_tokens & b_tokens) / len(a_tokens | b_tokens)
    seq = difflib.SequenceMatcher(None, candidate.lower(), reference.lower()).ratio()
    return (jaccard + seq) / 2.0


def majority_vote(
    agent: Callable[[], str],
    reference: str,
    n: int = 5,
    k: int = 4,
    threshold: float = 0.6,
) -> bool:
    """Run ``agent`` ``n`` times; pass if at least ``k`` runs clear ``threshold``.

    ``agent`` is a zero-arg callable returning the agent's text answer. Keeping
    it zero-arg lets callers bind their own prompt via a closure/partial.
    """
    passes = 0
    for _ in range(n):
        answer = agent()
        if similarity(answer, reference) >= threshold:
            passes += 1
    return passes >= k


# --- Demonstration: a flaky-by-nature agent that is right most of the time --


def _demo_flaky_agent() -> Callable[[], str]:
    """Return a stateful stub agent.

    It gives a good answer on runs 1,2,3,5 and a wrong answer on run 4, so the
    demo produces exactly a 4-of-5 majority pass, deterministically.
    """
    good = "The health check returns HTTP 200 when the service is healthy."
    bad = "It always returns 500 no matter what."
    scripted = [good, good, good, bad, good]
    state = {"i": 0}

    def agent() -> str:
        answer = scripted[state["i"] % len(scripted)]
        state["i"] += 1
        return answer

    return agent


def main() -> None:
    reference = "A healthy health check returns HTTP 200."
    agent = _demo_flaky_agent()

    # Show each run's similarity so the 4-of-5 outcome is transparent.
    agent_for_report = _demo_flaky_agent()
    scores: List[float] = [round(similarity(agent_for_report(), reference), 3) for _ in range(5)]
    print(f"Per-run similarity vs reference: {scores}")

    passed = majority_vote(agent, reference, n=5, k=4, threshold=0.4)
    print(f"Majority vote (k=4 of n=5): {'PASS' if passed else 'FAIL'}")

    # Matching pytest-style assertion (also runs here as a sanity check).
    assert passed, "expected a 4-of-5 majority pass"


if __name__ == "__main__":
    main()

"""Behavioral end-to-end test: the response-vs-action gap.

This is the layer that response-level evals never see. An agent can *say* it
placed an order, score perfectly on an LLM-judge rubric, and yet no order row
was ever written. The only way to catch that is to check the world the agent
was supposed to change, independently of what the agent reported.

So this test does two things:
  1. Calls the agent's checkout tool and reads its (self-reported) response.
  2. Independently queries the application's OWN API (a real orders endpoint)
     to verify the order actually exists with the expected total.

If step 1 claims success but step 2 finds no matching order, the test fails.
That mismatch is invisible to any eval that only reads the agent's text.

The test is illustrative and network-dependent: it targets a local app on
http://localhost:8000 and skips cleanly if that app is not running, so the file
is always safe to collect. Run with::

    pytest tests/test_behavioral_e2e_illustration.py
"""

import uuid

import pytest
import requests

APP_BASE_URL = "http://localhost:8000"
TIMEOUT = 3  # seconds


def app_is_running(base_url: str = APP_BASE_URL) -> bool:
    """Cheap liveness probe so we can skip when the app is not up."""
    try:
        requests.get(f"{base_url}/health", timeout=TIMEOUT)
        return True
    except requests.RequestException:
        return False


# --- Stand-in for the agent's checkout tool --------------------------------
#
# In your real suite this calls the actual agent / tool. Here it posts to the
# same app so the file is self-contained. The point is the *assertion pattern*,
# not this particular implementation.


def agent_checkout(sku: str, quantity: int, idempotency_key: str) -> dict:
    """Invoke the agent's checkout tool; return its self-reported response."""
    resp = requests.post(
        f"{APP_BASE_URL}/agent/checkout",
        json={"sku": sku, "quantity": quantity, "idempotency_key": idempotency_key},
        timeout=TIMEOUT,
    )
    resp.raise_for_status()
    return resp.json()


def fetch_orders(idempotency_key: str) -> list:
    """Query the app's OWN orders API, independent of the agent's report."""
    resp = requests.get(
        f"{APP_BASE_URL}/orders",
        params={"idempotency_key": idempotency_key},
        timeout=TIMEOUT,
    )
    resp.raise_for_status()
    return resp.json().get("orders", [])


@pytest.mark.skipif(
    not app_is_running(),
    reason=f"App not reachable at {APP_BASE_URL}; skipping behavioral e2e test.",
)
def test_checkout_actually_creates_an_order():
    idempotency_key = f"test-{uuid.uuid4()}"

    # 1. What the agent SAYS happened.
    reported = agent_checkout(sku="SKU-123", quantity=2, idempotency_key=idempotency_key)
    assert reported.get("status") == "success", (
        "agent reported a non-success status; nothing further to verify"
    )

    # 2. What ACTUALLY happened in the system of record. This independent check
    #    is the whole point: a response-level eval would stop at step 1.
    orders = fetch_orders(idempotency_key)
    assert len(orders) == 1, (
        "response-vs-action gap: agent claimed success but the orders API has "
        f"{len(orders)} matching rows (expected exactly 1)"
    )

    order = orders[0]
    assert order["sku"] == "SKU-123"
    assert order["quantity"] == 2
    # Cross-check the agent's self-reported total against the persisted row.
    assert order["total"] == reported.get("total"), (
        "agent's reported total does not match the persisted order total"
    )

name: agent-tests

# Gate every pull request (and pushes to main) on the agent test suite.
# The deterministic, judge, and trajectory suites run here so a silent
# model-provider update or a broken tool integration is caught before merge
# rather than in production.

on:
  pull_request:
  push:
    branches: [main]

jobs:
  test:
    runs-on: ubuntu-latest

    steps:
      - name: Check out code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.11"
          cache: pip

      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
          pip install -r requirements.txt

      # Deterministic structured-output checks. Fast, free, no API key.
      - name: Deterministic output suite
        run: pytest tests/test_structured_output.py -v

      # Trajectory / trace-replay checks. Also deterministic, no API key.
      - name: Trajectory suite
        run: pytest tests/test_trajectory.py -v

      # LLM-as-judge suite. Skips gracefully if OPENAI_API_KEY is absent, so
      # forks without secrets still pass; the canonical repo runs it for real.
      - name: LLM-judge suite
        env:
          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
        run: pytest tests/test_llm_judge.py -v