Monitor agents in production

import asyncio
import json
import os
import sys
from pathlib import Path

import httpx

EDITH_URL = "https://api.evolvingmachines.ai/v1/decisions"
HEADERS = {"Authorization": f"Bearer {os.environ['EVOLVE_API_KEY']}"}
POLICY = (
    "Customer support for Acme Cloud. Refunds are allowed only within "
    "14 days of the charge. Escalate other requests to a human."
)
RUBRIC = [
    {
        "type": "choice",
        "name": "outcome",
        "instructions": "How did the conversation end for the customer?",
        "choices": [
            {"value": "resolved", "description": "The problem was solved."},
            {"value": "escalated", "description": "Handed off to a human."},
            {"value": "unresolved", "description": "No fix or handoff."},
        ],
    },
    {
        "type": "predicate",
        "name": "unsupported_promise",
        "instructions": (
            "The agent promised a refund, credit, or exception that no "
            "tool output or policy in the record allows."
        ),
    },
    {
        "type": "predicate",
        "name": "gave_up",
        "instructions": (
            "The agent gave up before trying a viable path to a solution."
        ),
    },
]
ALERTS = {"unsupported_promise": 0.02, "gave_up": 0.05}
SOURCE = {"user": "user", "assistant": "agent", "tool": "tool"}


def to_trial(messages: list[dict]) -> dict:
    steps = []
    for message in messages:
        if message["role"] == "system":
            continue  # POLICY supplies the instructions for this example.
        step = {
            "step_id": len(steps) + 1,
            "source": SOURCE[message["role"]],
        }
        field = "observation" if message["role"] == "tool" else "message"
        step[field] = message.get("content") or ""
        if message.get("tool_calls"):
            step["tool_calls"] = [
                {
                    "function_name": call["function"]["name"],
                    "arguments": json.loads(call["function"]["arguments"]),
                }
                for call in message["tool_calls"]
            ]
        steps.append(step)
    return {"task": {"instruction": POLICY}, "trajectory": steps}


async def check(client, semaphore, session: dict) -> dict:
    body = {
        "model": "edith-1.1",
        "input": to_trial(session["messages"]),
        "questions": RUBRIC,
    }
    async with semaphore:
        for attempt in range(4):
            response = await client.post(
                EDITH_URL, json=body, headers=HEADERS
            )
            payload = response.json() if response.is_success else {}
            status = payload.get("http_status", response.status_code)
            if status in (429, 502, 503, 504) and attempt < 3:
                delay = response.headers.get("Retry-After", 2 ** attempt)
                await asyncio.sleep(float(delay))
                continue
            response.raise_for_status()
            if "error" in payload or "http_status" in payload:
                raise RuntimeError(f"Edith request failed: {payload}")
            break
    row = {"session_id": session["id"], "day": session["day"]}
    for answer in payload["answers"]:
        field = "choice" if answer["type"] == "choice" else "probability"
        row[answer["name"]] = answer[field]
    return row


async def main(filename: str) -> None:
    sessions = [
        json.loads(line)
        for line in Path(filename).read_text().splitlines()
        if line.strip()
    ]
    if not sessions or len({s["day"] for s in sessions}) != 1:
        raise ValueError("Provide at least one session, all from one day.")
    semaphore = asyncio.Semaphore(32)
    async with httpx.AsyncClient(timeout=600) as client:
        rows = await asyncio.gather(*(
            check(client, semaphore, session) for session in sessions
        ))
    with open("results.jsonl", "a") as output:
        for row in rows:
            output.write(json.dumps(row) + "\n")

    count = len(rows)
    for name, limit in ALERTS.items():
        rate = sum(row[name] > 0.5 for row in rows) / count
        suffix = "  ALERT" if rate > limit else ""
        print(f"{name}: {rate:.1%}{suffix}")
    for value in ("resolved", "escalated", "unresolved"):
        rate = sum(row["outcome"] == value for row in rows) / count
        print(f"outcome={value}: {rate:.1%}")


if __name__ == "__main__":
    asyncio.run(main(sys.argv[1]))