diff --git a/README.md b/README.md index 7587be1..38e773b 100644 --- a/README.md +++ b/README.md @@ -174,6 +174,7 @@ cat agentfailbench/failures/tool_drift/tool-semantic-drift-001.yaml ```text parallax/ ├── agentfailbench/ # benchmark cases + injectors +│ └── environments/ # TaskEnvironment protocol · customer_api · database (sqlite) ├── runtime/ # semantic monitor + detectors ├── recovery/ # RecoverAI policies ├── baselines/ # comparison methods diff --git a/agentfailbench/environments/__init__.py b/agentfailbench/environments/__init__.py index e69de29..f8bdd4d 100644 --- a/agentfailbench/environments/__init__.py +++ b/agentfailbench/environments/__init__.py @@ -0,0 +1,7 @@ +"""AgentFailBench environment packages and shared protocol.""" + +from agentfailbench.environments.base import TaskEnvironment +from agentfailbench.environments.customer_api import CustomerApiEnv +from agentfailbench.environments.database import DatabaseEnv + +__all__ = ["CustomerApiEnv", "DatabaseEnv", "TaskEnvironment"] diff --git a/agentfailbench/environments/base.py b/agentfailbench/environments/base.py new file mode 100644 index 0000000..f5bc60d --- /dev/null +++ b/agentfailbench/environments/base.py @@ -0,0 +1,28 @@ +"""Shared task-environment protocol for AgentFailBench. + +Every environment (customer API, database investigation, …) drives an episode +through the same three hooks so runners and injectors stay env-agnostic. +""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from runtime.schemas.episode import Action, EnvObservation + + +@runtime_checkable +class TaskEnvironment(Protocol): + """Lifecycle contract every AgentFailBench environment must implement.""" + + def reset(self) -> None: + """Return the environment to its initial deterministic state.""" + ... + + def step(self, action: Action) -> EnvObservation: + """Apply one agent action and return the transport-level observation.""" + ... + + def validate_success(self) -> bool: + """Return True when the task objective has been achieved.""" + ... diff --git a/agentfailbench/environments/customer_api/env.py b/agentfailbench/environments/customer_api/env.py index 2593db9..4003b2e 100644 --- a/agentfailbench/environments/customer_api/env.py +++ b/agentfailbench/environments/customer_api/env.py @@ -56,7 +56,11 @@ def plan_id_meaning(version: ContractVersion) -> str: @dataclass class CustomerApiEnv: - """Deterministic in-memory customer subscription API.""" + """Deterministic in-memory customer subscription API. + + Implements :class:`~agentfailbench.environments.base.TaskEnvironment` + (``reset``, ``step``, ``validate_success``). + """ task: TaskSpec contract_version: ContractVersion = "v1" diff --git a/agentfailbench/environments/database/__init__.py b/agentfailbench/environments/database/__init__.py new file mode 100644 index 0000000..4edba8b --- /dev/null +++ b/agentfailbench/environments/database/__init__.py @@ -0,0 +1,5 @@ +"""Deterministic in-memory SQLite investigation environment.""" + +from agentfailbench.environments.database.env import DatabaseEnv, seed_schema + +__all__ = ["DatabaseEnv", "seed_schema"] diff --git a/agentfailbench/environments/database/env.py b/agentfailbench/environments/database/env.py new file mode 100644 index 0000000..a480476 --- /dev/null +++ b/agentfailbench/environments/database/env.py @@ -0,0 +1,150 @@ +"""In-memory SQLite investigation environment for AgentFailBench Milestone 1.""" + +from __future__ import annotations + +import sqlite3 +from dataclasses import dataclass, field +from typing import Any + +from runtime.schemas.episode import Action, EnvObservation, TaskSpec + +# Default seed: a tiny CRM suitable for lookup / aggregation tasks. +SEED_SQL = """ +CREATE TABLE customers ( + customer_id TEXT PRIMARY KEY, + name TEXT NOT NULL, + segment TEXT NOT NULL +); +CREATE TABLE orders ( + order_id TEXT PRIMARY KEY, + customer_id TEXT NOT NULL REFERENCES customers(customer_id), + status TEXT NOT NULL, + amount_cents INTEGER NOT NULL +); +INSERT INTO customers VALUES + ('cust_001', 'Ada Lovelace', 'enterprise'), + ('cust_002', 'Grace Hopper', 'smb'), + ('cust_003', 'Alan Turing', 'enterprise'); +INSERT INTO orders VALUES + ('ord_100', 'cust_001', 'pending', 4500), + ('ord_101', 'cust_001', 'shipped', 12000), + ('ord_102', 'cust_002', 'pending', 800), + ('ord_103', 'cust_002', 'cancelled', 2000), + ('ord_104', 'cust_003', 'pending', 9900), + ('ord_105', 'cust_003', 'shipped', 1500); +""" + + +def seed_schema(conn: sqlite3.Connection) -> None: + """Create tables and insert the canonical fixture rows.""" + conn.executescript(SEED_SQL) + conn.commit() + + +def _rows_to_dicts(cursor: sqlite3.Cursor) -> list[dict[str, Any]]: + columns = [d[0] for d in cursor.description] if cursor.description else [] + return [dict(zip(columns, row, strict=True)) for row in cursor.fetchall()] + + +@dataclass +class DatabaseEnv: + """Deterministic in-memory SQLite investigation environment. + + Actions + ------- + ``list_tables`` + Return table names in the schema. + ``describe_table`` + Return column metadata for ``table``. + ``execute_sql`` + Run a single read-only ``SELECT`` (writes are rejected). + ``submit_answer`` + Record the agent's final answer for :meth:`validate_success`. + """ + + task: TaskSpec + conn: sqlite3.Connection = field(default_factory=lambda: sqlite3.connect(":memory:")) + call_count: int = 0 + _submitted: Any = None + _seeded: bool = False + + def __post_init__(self) -> None: + if not self._seeded: + seed_schema(self.conn) + self._seeded = True + + def reset(self) -> None: + self.conn.close() + self.conn = sqlite3.connect(":memory:") + seed_schema(self.conn) + self._seeded = True + self.call_count = 0 + self._submitted = None + + def step(self, action: Action) -> EnvObservation: + self.call_count += 1 + name = action.name + args = action.arguments + + if name == "list_tables": + cur = self.conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name" + ) + tables = [row[0] for row in cur.fetchall()] + return EnvObservation(success=True, data={"tables": tables}) + + if name == "describe_table": + table = str(args.get("table", "")) + if not table or not table.replace("_", "").isalnum(): + return EnvObservation(success=False, status_code=400, error="invalid_table") + cur = self.conn.execute(f"PRAGMA table_info({table})") + columns = [ + {"cid": r[0], "name": r[1], "type": r[2], "notnull": r[3], "pk": r[5]} + for r in cur.fetchall() + ] + if not columns: + return EnvObservation(success=False, status_code=404, error="table_not_found") + return EnvObservation(success=True, data={"table": table, "columns": columns}) + + if name == "execute_sql": + sql = str(args.get("sql", "")).strip() + if not sql: + return EnvObservation(success=False, status_code=400, error="empty_sql") + # Read-only: reject anything that is not a single SELECT. + lowered = sql.lstrip().lower() + if not lowered.startswith("select"): + return EnvObservation(success=False, status_code=400, error="write_not_allowed") + if ";" in sql.rstrip(";"): + return EnvObservation(success=False, status_code=400, error="multiple_statements") + try: + cur = self.conn.execute(sql) + except sqlite3.Error as exc: + return EnvObservation( + success=False, status_code=400, error=f"sql_error:{exc}", data={} + ) + rows = _rows_to_dicts(cur) + return EnvObservation(success=True, data={"rows": rows, "row_count": len(rows)}) + + if name == "submit_answer": + self._submitted = args.get("value") + return EnvObservation( + success=True, + data={"submitted": self._submitted}, + ) + + return EnvObservation(success=False, status_code=400, error=f"unknown_action:{name}") + + def validate_success(self) -> bool: + """True when the agent submitted the task's expected answer.""" + expected = self.task.expected_answer + if expected is None: + return False + if self._submitted is None: + return False + # Normalize numeric strings ("15300" vs 15300). + if isinstance(expected, (int, float)) and not isinstance(expected, bool): + try: + return float(self._submitted) == float(expected) + except (TypeError, ValueError): + return False + return str(self._submitted) == str(expected) diff --git a/agentfailbench/runners/episode.py b/agentfailbench/runners/episode.py index 9d78233..5800263 100644 --- a/agentfailbench/runners/episode.py +++ b/agentfailbench/runners/episode.py @@ -4,9 +4,10 @@ from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, cast from agentfailbench.agents.scripted import ScriptedApiAgent +from agentfailbench.environments.base import TaskEnvironment from agentfailbench.environments.customer_api.env import CustomerApiEnv from agentfailbench.failures.base import FailureInjector, FailureInjectorRegistry from agentfailbench.failures.memory.injector import StaleMemoryInjector @@ -99,11 +100,12 @@ def _run_traced( case: BenchmarkCaseModel, *, inject_failure: bool, -) -> tuple[CustomerApiEnv, ScriptedApiAgent, TraceCollector, FailureInjector]: +) -> tuple[TaskEnvironment, ScriptedApiAgent, TraceCollector, FailureInjector]: task = build_task(case) - env = CustomerApiEnv(task=task) + api_env = CustomerApiEnv(task=task) + env: TaskEnvironment = api_env memory = AgentMemory() if _is_memory_case(case) else None - agent = ScriptedApiAgent(env=env, memory=memory) + agent = ScriptedApiAgent(env=api_env, memory=memory) injector, injectors = _build_injector(case, memory=memory) collector = TraceCollector(task_id=case.case_id) @@ -191,7 +193,7 @@ def run_episode( "injector_triggered": injector.triggered, "ground_truth_first_detectable": case.ground_truth.first_detectable_step, "ground_truth_root_cause": case.ground_truth.root_cause.value, - "contract_final": env.contract_version, + "contract_final": cast(CustomerApiEnv, env).contract_version, "suite": "memory" if _is_memory_case(case) else "tool_drift", }, ) diff --git a/docs/benchmark-specification.md b/docs/benchmark-specification.md index 6a8fc46..714c8ad 100644 --- a/docs/benchmark-specification.md +++ b/docs/benchmark-specification.md @@ -66,12 +66,36 @@ risk: - ground-truth root-cause and recovery labels - reproducible configurations +## Shared task-environment interface + +Every environment implements the ``TaskEnvironment`` protocol +(``agentfailbench.environments.base``): + +| Method | Role | +| --- | --- | +| `reset()` | Restore deterministic initial state | +| `step(action) -> EnvObservation` | Apply one agent action | +| `validate_success() -> bool` | Task-success validator | + +Shipped environments: + +| Package | Environment id | Notes | +| --- | --- | --- | +| `environments/customer_api` | `customer_service_api` | Versioned subscription tool contracts | +| `environments/database` | `sqlite_investigation` | In-memory SQLite; `submit_answer` closes the task | + +Action / observation schemas live in `runtime.schemas.episode` +(`Action`, `EnvObservation`, `TaskSpec`). + ## Package layout ```text agentfailbench/ tasks/ environments/ + base.py # TaskEnvironment protocol + customer_api/ + database/ # in-memory SQLite investigation failures/ labels/ runners/ diff --git a/runtime/schemas/episode.py b/runtime/schemas/episode.py index bf638d2..46ef2f3 100644 --- a/runtime/schemas/episode.py +++ b/runtime/schemas/episode.py @@ -8,7 +8,12 @@ class TaskSpec(BaseModel): - """Specification for a benchmark task instance.""" + """Specification for a benchmark task instance. + + Customer-API tasks use ``customer_id`` / ``target_plan_code``. + Database investigation tasks set ``expected_answer`` for + :meth:`~agentfailbench.environments.database.env.DatabaseEnv.validate_success`. + """ task_id: str objective: str @@ -16,6 +21,7 @@ class TaskSpec(BaseModel): customer_id: str = "cust_001" target_plan_code: str = "GOLD_ANNUAL" expected_steps: int = 8 + expected_answer: Any | None = None class Action(BaseModel): diff --git a/tests/unit/test_task_environment.py b/tests/unit/test_task_environment.py new file mode 100644 index 0000000..dbeaa32 --- /dev/null +++ b/tests/unit/test_task_environment.py @@ -0,0 +1,87 @@ +"""Unit tests for the shared TaskEnvironment protocol and DatabaseEnv.""" + +from __future__ import annotations + +from agentfailbench.environments import CustomerApiEnv, DatabaseEnv, TaskEnvironment +from agentfailbench.environments.customer_api.env import CustomerApiEnv as CustomerApiEnvDirect +from agentfailbench.environments.database.env import DatabaseEnv as DatabaseEnvDirect +from runtime.schemas.episode import Action, TaskSpec + + +def _api_task() -> TaskSpec: + return TaskSpec( + task_id="t-api", + objective="update_customer_subscription", + environment="customer_service_api", + ) + + +def _db_task(*, expected_answer: int | str = 15300) -> TaskSpec: + return TaskSpec( + task_id="t-db", + objective="sum_pending_order_amounts", + environment="sqlite_investigation", + expected_answer=expected_answer, + expected_steps=4, + ) + + +def test_customer_api_satisfies_protocol() -> None: + env = CustomerApiEnv(task=_api_task()) + assert isinstance(env, TaskEnvironment) + + +def test_database_env_satisfies_protocol() -> None: + env = DatabaseEnv(task=_db_task()) + assert isinstance(env, TaskEnvironment) + + +def test_database_clean_episode_succeeds() -> None: + """Normal (non-failure) path: explore schema, aggregate, submit answer.""" + env = DatabaseEnvDirect(task=_db_task(expected_answer=15300)) + # pending amounts: 4500 + 800 + 9900 = 15300 + steps = [ + Action(name="list_tables", arguments={}), + Action(name="describe_table", arguments={"table": "orders"}), + Action( + name="execute_sql", + arguments={ + "sql": "SELECT SUM(amount_cents) AS total FROM orders WHERE status = 'pending'" + }, + ), + Action(name="submit_answer", arguments={"value": 15300}), + ] + for action in steps: + obs = env.step(action) + assert obs.success, obs.error + assert env.validate_success() + + +def test_database_wrong_answer_fails_validator() -> None: + env = DatabaseEnv(task=_db_task(expected_answer=15300)) + env.step(Action(name="submit_answer", arguments={"value": 0})) + assert not env.validate_success() + + +def test_database_rejects_writes() -> None: + env = DatabaseEnv(task=_db_task()) + obs = env.step(Action(name="execute_sql", arguments={"sql": "DELETE FROM orders"})) + assert not obs.success + assert obs.error == "write_not_allowed" + + +def test_database_reset_restores_seed() -> None: + env = DatabaseEnv(task=_db_task(expected_answer=3)) + env.step(Action(name="submit_answer", arguments={"value": 3})) + assert env.validate_success() + env.reset() + assert not env.validate_success() + obs = env.step( + Action(name="execute_sql", arguments={"sql": "SELECT COUNT(*) AS n FROM customers"}) + ) + assert obs.success + assert obs.data["rows"][0]["n"] == 3 + + +def test_customer_api_still_exports() -> None: + assert CustomerApiEnv is CustomerApiEnvDirect