bat-cli 0.1.0__py3-none-any.whl

eval/engine/adapter.py

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+import json
+import time
+from typing import Any
+
+from a2a.client import A2ACardResolver, ClientConfig, ClientFactory
+from a2a.helpers import (
+    get_artifact_text,
+    get_message_text,
+    new_text_message,
+)
+from a2a.types import Role, SendMessageRequest, StreamResponse, TaskState
+from google.protobuf.json_format import MessageToDict
+from httpx import AsyncClient
+
+from .contracts import EpisodeResult, EpisodeTrace, TaskSpec, TraceEvent
+
+
+TERMINAL_STATUSES = {"completed", "error", "input-required"}
+
+
+_TASK_STATE_TO_STR = {
+    TaskState.TASK_STATE_SUBMITTED: "working",
+    TaskState.TASK_STATE_WORKING: "working",
+    TaskState.TASK_STATE_INPUT_REQUIRED: "input-required",
+    TaskState.TASK_STATE_COMPLETED: "completed",
+    TaskState.TASK_STATE_FAILED: "error",
+    TaskState.TASK_STATE_CANCELED: "error",
+    TaskState.TASK_STATE_REJECTED: "error",
+}
+
+
+def _to_dict(value: Any) -> dict[str, Any]:
+    if value is None:
+        return {}
+    if isinstance(value, dict):
+        return value
+    if hasattr(value, "model_dump"):
+        dumped = value.model_dump(by_alias=True)
+        return dumped if isinstance(dumped, dict) else {}
+    return {}
+
+
+def _struct_to_dict(metadata_struct: Any) -> dict[str, Any]:
+    if metadata_struct is None:
+        return {}
+    try:
+        return MessageToDict(metadata_struct) or {}
+    except Exception:
+        return {}
+
+
+def _chunk_key(chunk: StreamResponse) -> str:
+    if chunk.HasField("status_update"):
+        message_id = chunk.status_update.status.message.message_id
+        if message_id:
+            return f"status:{message_id}"
+    if chunk.HasField("artifact_update"):
+        artifact_id = chunk.artifact_update.artifact.artifact_id
+        if artifact_id:
+            return f"artifact:{artifact_id}"
+    if chunk.HasField("message"):
+        if chunk.message.message_id:
+            return f"message:{chunk.message.message_id}"
+    if chunk.HasField("task"):
+        if chunk.task.id:
+            return f"task:{chunk.task.id}"
+    return f"raw:{chunk.SerializeToString().hex()}"
+
+
+def _extract_status_and_content(chunk: StreamResponse) -> tuple[str | None, str]:
+    if chunk.HasField("message"):
+        return "completed", get_message_text(chunk.message)
+    if chunk.HasField("artifact_update"):
+        return "completed", get_artifact_text(chunk.artifact_update.artifact)
+    if chunk.HasField("status_update"):
+        state = chunk.status_update.status.state
+        status = _TASK_STATE_TO_STR.get(state)
+        content = ""
+        if chunk.status_update.status.HasField("message"):
+            content = get_message_text(chunk.status_update.status.message)
+        return status, content
+    if chunk.HasField("task"):
+        state = chunk.task.status.state
+        status = _TASK_STATE_TO_STR.get(state)
+        texts = [get_artifact_text(a) for a in chunk.task.artifacts]
+        return status, "\n".join(t for t in texts if t)
+    return None, ""
+
+
+def _extract_metadata(chunk: StreamResponse) -> dict[str, Any]:
+    metadata_struct: Any = None
+    if chunk.HasField("status_update"):
+        metadata_struct = chunk.status_update.metadata
+    elif chunk.HasField("artifact_update"):
+        metadata_struct = chunk.artifact_update.metadata
+    elif chunk.HasField("message"):
+        metadata_struct = chunk.message.metadata
+    elif chunk.HasField("task"):
+        metadata_struct = chunk.task.metadata
+
+    metadata = _struct_to_dict(metadata_struct)
+
+    if chunk.HasField("artifact_update"):
+        artifact_metadata = _struct_to_dict(chunk.artifact_update.artifact.metadata)
+        if artifact_metadata:
+            merged = dict(artifact_metadata)
+            merged.update(metadata)
+            metadata = merged
+
+    return metadata
+
+
+def _normalize_usage(metadata: dict[str, Any]) -> dict[str, Any]:
+    usage = _to_dict(metadata.get("usage"))
+    if not usage:
+        return {}
+
+    input_tokens = int(usage.get("input_tokens") or 0)
+    output_tokens = int(usage.get("output_tokens") or 0)
+    total_tokens = int(usage.get("total_tokens") or (input_tokens + output_tokens))
+    inference_time = float(usage.get("inference_time") or 0.0)
+
+    if input_tokens == 0 and output_tokens == 0 and total_tokens == 0 and inference_time == 0.0:
+        return {}
+
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": total_tokens,
+        "inference_time": inference_time,
+    }
+
+
+def _add_usage(total: dict[str, Any], incremental: dict[str, Any]) -> dict[str, Any]:
+    return {
+        "input_tokens": int(total.get("input_tokens") or 0) + int(incremental.get("input_tokens") or 0),
+        "output_tokens": int(total.get("output_tokens") or 0) + int(incremental.get("output_tokens") or 0),
+        "total_tokens": int(total.get("total_tokens") or 0) + int(incremental.get("total_tokens") or 0),
+        "inference_time": float(total.get("inference_time") or 0.0)
+        + float(incremental.get("inference_time") or 0.0),
+    }
+
+
+def _extract_tool_calls(metadata: dict[str, Any]) -> list[dict[str, Any]]:
+    trace = _to_dict(metadata.get("trace"))
+    tool_calls = trace.get("tool_calls")
+    if not isinstance(tool_calls, list):
+        return []
+    return [item for item in tool_calls if isinstance(item, dict)]
+
+
+def _tool_call_key(tool_call: dict[str, Any]) -> str:
+    call_id = tool_call.get("id")
+    if isinstance(call_id, str) and call_id:
+        return f"id:{call_id}"
+    return json.dumps(tool_call, sort_keys=True, ensure_ascii=True, default=str)
+
+
+class BatA2AAdapter:
+    def __init__(self, agent_url: str, request_timeout_s: float = 180.0, max_events: int = 200) -> None:
+        self.agent_url = agent_url
+        self.request_timeout_s = request_timeout_s
+        self.max_events = max_events
+
+    async def run_task(self, task: TaskSpec, *, thread_id: str) -> EpisodeResult:
+        t0_perf = time.perf_counter()
+        trace = EpisodeTrace()
+
+        usage_total: dict[str, Any] = {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "total_tokens": 0,
+            "inference_time": 0.0,
+        }
+        usage_seen: set[str] = set()
+        tool_calls_seen: set[str] = set()
+
+        last_status: str | None = None
+        last_content = ""
+
+        async with AsyncClient(timeout=self.request_timeout_s) as httpx_client:
+            resolver = A2ACardResolver(httpx_client=httpx_client, base_url=self.agent_url)
+            agent_card = await resolver.get_agent_card()
+
+            client = ClientFactory(ClientConfig(httpx_client=httpx_client, streaming=True)).create(card=agent_card)
+
+            try:
+                for turn in task.turns:
+                    turn_started = False
+                    message = new_text_message(
+                        text=turn,
+                        context_id=thread_id,
+                        role=Role.ROLE_USER,
+                    )
+                    stream = client.send_message(SendMessageRequest(message=message))
+
+                    async for chunk in stream:
+                        metadata = _extract_metadata(chunk)
+
+                        usage = _normalize_usage(metadata)
+                        if usage:
+                            usage_key = f"{_chunk_key(chunk)}::{json.dumps(usage, sort_keys=True, ensure_ascii=True)}"
+                            if usage_key not in usage_seen:
+                                usage_seen.add(usage_key)
+                                usage_total = _add_usage(usage_total, usage)
+
+                        for tool_call in _extract_tool_calls(metadata):
+                            key = _tool_call_key(tool_call)
+                            if key in tool_calls_seen:
+                                continue
+                            tool_calls_seen.add(key)
+                            trace.tool_calls.append(tool_call)
+
+                        status, content = _extract_status_and_content(chunk)
+                        if status is None:
+                            continue
+
+                        if len(trace.events) < self.max_events:
+                            trace.events.append(
+                                TraceEvent(
+                                    t_ms=(time.perf_counter() - t0_perf) * 1000.0,
+                                    task_status=status,
+                                    content_preview=content,
+                                    user_input=turn if not turn_started else None,
+                                )
+                            )
+                            turn_started = True
+
+                        if status in TERMINAL_STATUSES:
+                            last_status = status
+                            last_content = content or ""
+
+            except Exception as exc:
+                last_status = "error"
+                last_content = f"{type(exc).__name__}: {exc}"
+
+        trace.timings["wall_ms"] = (time.perf_counter() - t0_perf) * 1000.0
+        trace.usage = usage_total
+
+        final_status = last_status or "error"
+        final_output = last_content or ""
+
+        return EpisodeResult(
+            task_id=task.id,
+            final_status=final_status,
+            final_output=final_output,
+            trace=trace,
+            aux={"agent_url": self.agent_url},
+        )

eval/engine/bench_runner.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+from .contracts import EpisodeResult, TaskSpec
+from .evaluator import EpisodeEvaluator
+
+
+_QUALITATIVE_FIELDS = (
+    "response_relevance",
+    "task_completion_quality",
+    "hallucination_score",
+    "tool_call_appropriateness",
+)
+
+
+def _episode_passed(ep: EpisodeResult) -> bool:
+    return ep.verdict.passed if ep.verdict is not None else False
+
+
+def _average_qualitative_scores(results: list[EpisodeResult]) -> dict[str, float]:
+    out: dict[str, float] = {}
+    for field in _QUALITATIVE_FIELDS:
+        values = [
+            getattr(r.qualitative_scores, field)
+            for r in results
+            if r.qualitative_scores is not None
+            and getattr(r.qualitative_scores, field) is not None
+        ]
+        if values:
+            out[field] = sum(values) / len(values)
+    return out
+
+
+def _safe_task_id(task_id: str) -> str:
+    return "".join(ch if ch.isalnum() or ch in ("-", "_", ".") else "_" for ch in task_id)
+
+
+def _write_json(path: Path, obj: object) -> None:
+    path.write_text(json.dumps(obj, indent=2, ensure_ascii=False), encoding="utf-8")
+
+
+@dataclass
+class RunConfig:
+    run_name: str
+    out_dir: str = "output"
+    k: int = 1
+    model: str = "default"
+    task_id: str = ""
+
+
+class BenchRunner:
+    def __init__(
+        self,
+        adapter: object,
+        config: RunConfig,
+        evaluator: EpisodeEvaluator | None = None,
+    ):
+        self.adapter = adapter
+        self.config = config
+        self.evaluator = evaluator or EpisodeEvaluator()
+        self.task_dir: Path | None = None
+        self.run_dir: Path | None = None
+
+    def _episodes_dir(self) -> Path:
+        if self.run_dir is None:
+            raise ValueError("run_dir is not initialized")
+        return self.run_dir / "episodes"
+
+    def persist_results(self, results: list[EpisodeResult]) -> None:
+        episodes_dir = self._episodes_dir()
+        episodes_dir.mkdir(parents=True, exist_ok=True)
+
+        for episode in results:
+            task_file_id = _safe_task_id(episode.task_id)
+            attempt_index = int(episode.aux.get("attempt_index", 0))
+            (episodes_dir / f"{task_file_id}__try{attempt_index}.json").write_text(
+                episode.model_dump_json(indent=2),
+                encoding="utf-8",
+            )
+
+    async def run(self, tasks: list[TaskSpec]) -> list[EpisodeResult]:
+        stamp = time.strftime("%Y%m%d_%H%M%S", time.gmtime())
+        task_id = self.config.task_id or stamp
+        safe_model_name = self.config.model.replace(":", "-")
+
+        self.task_dir = Path(self.config.out_dir) / task_id
+        self.run_dir = self.task_dir / f"{self.config.run_name}_{safe_model_name}"
+        self._episodes_dir().mkdir(parents=True, exist_ok=True)
+        self._run_timestamp = stamp
+
+        all_attempts: list[EpisodeResult] = []
+
+        for task in tasks:
+            for i in range(max(1, int(self.config.k))):
+                thread_id = f"{task.id}__try{i}"
+                episode = await self.adapter.run_task(task=task, thread_id=thread_id)
+                episode.verdict = self.evaluator.evaluate(
+                    episode.final_status, episode.final_output, episode.trace.tool_calls, task.expected
+                )
+                episode.expected_outcome = task.expected.expected_outcome
+                episode.model_name = self.config.model
+                episode.aux["attempt_index"] = i
+                all_attempts.append(episode)
+
+        self.persist_results(all_attempts)
+        return all_attempts
+
+    def write_summary(self, results: list[EpisodeResult]) -> None:
+        if self.run_dir is None:
+            raise ValueError("run_dir is not initialized; call run() first")
+
+        display_model_name = (
+            self.config.model.split(":")[-1] if ":" in self.config.model else self.config.model
+        )
+
+        attempts_by_task: dict[str, list[EpisodeResult]] = {}
+        for attempt in results:
+            attempts_by_task.setdefault(attempt.task_id, []).append(attempt)
+
+        attempts = []
+        for task_id, task_attempts in attempts_by_task.items():
+            total = len(task_attempts)
+            passed = sum(1 for ep in task_attempts if _episode_passed(ep))
+            attempts.append(
+                {
+                    "task_id": task_id,
+                    "attempts": total,
+                    "passed": passed,
+                    "failed": total - passed,
+                    "success_percentage": (passed / total) * 100.0 if total else 0.0,
+                }
+            )
+
+        _write_json(
+            self.run_dir / "summary.json",
+            {
+                "run_name": self.config.run_name,
+                "timestamp_utc": getattr(self, "_run_timestamp", ""),
+                "k": self.config.k,
+                "model_name": display_model_name,
+                "attempts": attempts,
+                "qualitative_scores": _average_qualitative_scores(results),
+                "passed": sum(1 for ep in results if _episode_passed(ep)),
+            },
+        )

eval/engine/contracts.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+from pathlib import Path
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+AgentTaskStatus = Literal["working", "input-required", "completed", "error"]
+
+
+class ExpectedToolCall(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    name: str
+    args_subset: dict[str, Any] = Field(default_factory=dict)
+    times: int = 1
+
+
+class TaskExpected(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    # Expected final status from the agent. None = skip status check entirely.
+    status: AgentTaskStatus | None = "completed"
+    # Free-text description of the desired outcome, evaluated semantically by the LLM judges.
+    expected_outcome: str | None = None
+    # All phrases must appear in the final output text. None/empty = skip substring check.
+    output_must_contain: list[str] | None = None
+    tool_calls: list[ExpectedToolCall] = Field(default_factory=list)
+
+
+class TaskSpec(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    id: str
+    turns: list[str]
+    expected: TaskExpected = Field(default_factory=TaskExpected)
+    meta: dict[str, Any] = Field(default_factory=dict)
+
+
+class TraceEvent(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    t_ms: float
+    task_status: AgentTaskStatus
+    content_preview: str
+    user_input: str | None = None
+
+
+class EpisodeTrace(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    events: list[TraceEvent] = Field(default_factory=list)
+    usage: dict[str, Any] = Field(default_factory=dict)
+    timings: dict[str, float] = Field(default_factory=dict)
+    tool_calls: list[dict[str, Any]] = Field(default_factory=list)
+
+
+class QualitativeScores(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    response_relevance: float | None = None
+    task_completion_quality: float | None = None
+    hallucination_score: float | None = None
+    tool_call_appropriateness: float | None = None
+    judge_reasoning: dict[str, str] = Field(default_factory=dict)
+
+
+class EpisodeVerdict(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    passed: bool
+    reason: str = ""
+
+
+class EpisodeResult(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    model_name: str | None = None
+    task_id: str
+    expected_outcome: str | None = None
+    final_status: AgentTaskStatus
+    final_output: str
+    verdict: EpisodeVerdict | None = None
+    qualitative_scores: QualitativeScores | None = None
+    aux: dict[str, Any] = Field(default_factory=dict)
+    trace: EpisodeTrace = Field(default_factory=EpisodeTrace)
+
+
+class ModelSpec(BaseModel):
+    provider: str
+    model: str
+    base_url: str | None = None
+    env: dict[str, str] = Field(default_factory=dict)
+
+
+class JudgeSpec(BaseModel):
+    provider: str
+    model: str
+    base_url: str | None = None
+    api_key_env: str | None = None
+    env: dict[str, str] = Field(default_factory=dict)
+    prompts: dict[str, str] = Field(default_factory=dict)
+
+
+class EvalConfig(BaseModel):
+    dataset: Path
+    output_dir: Path
+    agent_url: str
+    agent_startup_timeout_s: int = 45
+    agent_shutdown_timeout_s: int = 10
+    k: int
+    qualitative: bool
+    run_name: str
+    models: list[ModelSpec]
+    judge: JudgeSpec | None