python-yama 0.1.0__py3-none-any.whl

yama/runner.py

@@ -0,0 +1,531 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import shutil
+import traceback
+from copy import deepcopy
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from .assertions import run_hard_assertions, run_judge
+from .tools.bash.fs import cleanup_session
+from .llm import LLMClient
+from .mocks import MockToolRegistry, ToolMockContext, _matches, load_python_callable
+from .models import ResolvedCase, StepResult, ToolCallRecord, sanitize_case_key
+
+
+def _tool_result_content(result: Any) -> str:
+    """Render a mock tool result as the `content` of a `role: tool` message.
+
+    The `bash` tool's handlers (run_bash_fs/run_bash_cli) return the raw
+    `{stdout, stderr, exit_code}` envelope so the framework can inspect it
+    for records/assertions -- but the model should see the command's actual
+    output text, the same way a real shell tool would return it, not our
+    internal JSON envelope. The `skill` tool returns the skill file's text
+    verbatim (see `tools/skill/skill.py`), which should also reach the
+    model as-is rather than as a JSON-quoted string. Every other tool
+    result (including bash error envelopes) keeps the previous JSON
+    encoding.
+    """
+    if (
+        isinstance(result, dict)
+        and isinstance(result.get("stdout"), str)
+        and isinstance(result.get("stderr"), str)
+        and "exit_code" in result
+    ):
+        return result["stdout"] + result["stderr"]
+    if isinstance(result, str):
+        return result
+    return json.dumps(result, ensure_ascii=False)
+
+
+class MaxToolRoundsExceeded(RuntimeError):
+    def __init__(self, step_id: str, tool_calls: list[ToolCallRecord]) -> None:
+        super().__init__(f"max tool rounds exceeded for step {step_id}")
+        self.step_id = step_id
+        self.tool_calls = tool_calls
+
+
+@dataclass
+class RunResult:
+    run_index: int
+    passed: bool
+    steps: list[StepResult] = field(default_factory=list)
+    error: str | None = None
+    traceback: str | None = None
+    mock_events: list[dict[str, Any]] = field(default_factory=list)
+    mock_state: dict[str, Any] = field(default_factory=dict)
+    requests: list[dict[str, Any]] = field(default_factory=list)
+    artifact_dir: Path | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "run_index": self.run_index,
+            "passed": self.passed,
+            "steps": [step.to_dict() for step in self.steps],
+            "error": self.error,
+            "mock_events": self.mock_events,
+            "mock_state": self.mock_state,
+            "requests": self.requests,
+            "artifact_dir": str(self.artifact_dir) if self.artifact_dir else None,
+        }
+
+
+@dataclass
+class CaseResult:
+    case_key: str
+    passed: bool
+    runs: list[RunResult]
+    pass_rate: float
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "case_key": self.case_key,
+            "passed": self.passed,
+            "pass_rate": self.pass_rate,
+            "runs": [run.to_dict() for run in self.runs],
+        }
+
+
+class YamaRunner:
+    def __init__(
+        self,
+        llm: LLMClient,
+        *,
+        result_root: Path | None = None,
+        write_artifacts: bool = True,
+    ) -> None:
+        self.llm = llm
+        self.result_root = result_root
+        self.write_artifacts = write_artifacts
+
+    async def run_case(self, case: ResolvedCase) -> CaseResult:
+        repeat = int(case.execution.get("repeat", 1))
+        if repeat < 1:
+            raise ValueError("execution.repeat must be at least 1")
+        runs = [
+            await self._run_once(case, run_index) for run_index in range(1, repeat + 1)
+        ]
+        passed_count = sum(run.passed for run in runs)
+        result = CaseResult(
+            case_key=case.case_key,
+            passed=passed_count == repeat,
+            runs=runs,
+            pass_rate=passed_count / repeat,
+        )
+        return result
+
+    async def _run_once(self, case: ResolvedCase, run_index: int) -> RunResult:
+        artifact_dir = self._artifact_dir(case, run_index)
+        registry: MockToolRegistry | None = None
+        steps: list[StepResult] = []
+        requests: list[dict[str, Any]] = []
+        try:
+            registry = MockToolRegistry(
+                case.mocks,
+                plugin_root=case.plugin_root,
+                relative_root=case.evals_root,
+                exposed_tools={tool.name for tool in case.tools},
+            )
+            messages = await self._initial_messages(case, registry, run_index)
+            for step in case.steps:
+                transcript_before = deepcopy(messages)
+                timeout = float(case.execution.get("timeout_seconds_per_step", 120))
+                try:
+                    step_result = await asyncio.wait_for(
+                        self._run_step(
+                            case,
+                            step,
+                            messages,
+                            registry,
+                            run_index,
+                            requests,
+                        ),
+                        timeout=timeout,
+                    )
+                except MaxToolRoundsExceeded as exc:
+                    messages[:] = transcript_before
+                    steps.append(
+                        self._max_tool_rounds_step_result(case, step, exc)
+                    )
+                    continue
+                assertion_config = step.get("assert", {})
+                hard = assertion_config.get("hard", [])
+                step_result.hard_checks = await run_hard_assertions(
+                    hard,
+                    result=step_result,
+                    transcript=messages,
+                    state=registry.state,
+                    plugin_root=case.plugin_root,
+                    relative_root=case.evals_root,
+                )
+                step_result.hard_pass = all(
+                    check["passed"] for check in step_result.hard_checks
+                )
+                judge_config = assertion_config.get("judge")
+                if judge_config and (
+                    step_result.hard_pass
+                    or judge_config.get("run_if_hard_failed", False)
+                ):
+                    step_result.judge_result = await run_judge(
+                        judge_config,
+                        client=self.llm,
+                        step_result=step_result,
+                        transcript_before_step=transcript_before,
+                        default_model=case.model,
+                    )
+                steps.append(step_result)
+            passed = self._case_passed(case, steps)
+            run = RunResult(
+                run_index=run_index,
+                passed=passed,
+                steps=steps,
+                mock_events=deepcopy(registry.events),
+                mock_state=deepcopy(registry.state),
+                requests=requests,
+                artifact_dir=artifact_dir,
+            )
+        except Exception as exc:
+            run = RunResult(
+                run_index=run_index,
+                passed=False,
+                steps=steps,
+                error=f"{type(exc).__name__}: {exc}",
+                traceback=traceback.format_exc(),
+                mock_events=deepcopy(registry.events) if registry else [],
+                mock_state=deepcopy(registry.state) if registry else {},
+                requests=requests,
+                artifact_dir=artifact_dir,
+            )
+        finally:
+            if registry is not None:
+                await cleanup_session(registry.state)
+        if self.write_artifacts:
+            self._write_artifacts(case, run)
+        return run
+
+    def _max_tool_rounds_step_result(
+        self,
+        case: ResolvedCase,
+        step: dict[str, Any],
+        exc: MaxToolRoundsExceeded,
+    ) -> StepResult:
+        maximum = int(case.execution.get("max_tool_rounds_per_step", 10))
+        user_message = {"role": "user", "content": deepcopy(step["user"])}
+        return StepResult(
+            step_id=step["id"],
+            user_message=user_message,
+            assistant_message={"role": "assistant", "content": ""},
+            tool_calls=exc.tool_calls,
+            messages=[],
+            hard_checks=[
+                {
+                    "type": "max_tool_rounds_exceeded",
+                    "passed": False,
+                    "label": "tool call loop limit",
+                    "expected": f"step completes within {maximum} tool round(s)",
+                    "actual": f"{len(exc.tool_calls)} tool call(s) made before limit",
+                    "reason": str(exc),
+                    "detail": "step was reset; transcript rolled back and run continued",
+                }
+            ],
+            hard_pass=False,
+        )
+
+    async def _initial_messages(
+        self,
+        case: ResolvedCase,
+        registry: MockToolRegistry,
+        run_index: int,
+    ) -> list[dict[str, Any]]:
+        messages = deepcopy(case.initial_messages)
+        builders = case.context.get("builders", [])
+        case_messages: dict[str, list[dict[str, Any]]] = {
+            "after_system": [],
+            "before_history": [],
+            "after_history": [],
+        }
+        for position in case_messages:
+            case_messages[position] = await self._run_builders(
+                builders,
+                position=position,
+                case=case,
+                transcript=messages,
+                state=registry.state,
+                run_index=run_index,
+            )
+        history = messages[1:]
+        return [
+            messages[0],
+            *case_messages["after_system"],
+            *case_messages["before_history"],
+            *history,
+            *case_messages["after_history"],
+        ]
+
+    async def _run_step(
+        self,
+        case: ResolvedCase,
+        step: dict[str, Any],
+        messages: list[dict[str, Any]],
+        registry: MockToolRegistry,
+        run_index: int,
+        requests: list[dict[str, Any]],
+    ) -> StepResult:
+        injections = step.get("message_inject", [])
+        messages.extend(
+            deepcopy(item["message"])
+            for item in injections
+            if item["position"] == "before_user_message"
+        )
+        user_message = {"role": "user", "content": deepcopy(step["user"])}
+        messages.append(user_message)
+        messages.extend(
+            deepcopy(item["message"])
+            for item in injections
+            if item["position"] == "after_user_message"
+        )
+        records: list[ToolCallRecord] = []
+        step_call_counts: dict[str, int] = {}
+        maximum = int(case.execution.get("max_tool_rounds_per_step", 10))
+        for request_index in range(maximum + 1):
+            request_messages = deepcopy(messages)
+            request_record = {
+                "step_id": step["id"],
+                "request_index": request_index,
+            }
+            requests.append(request_record)
+            try:
+                response = await self.llm.complete(
+                    messages=request_messages,
+                    tools=case.tools,
+                    model=case.model,
+                )
+            except Exception as exc:
+                request_record["error"] = f"{type(exc).__name__}: {exc}"
+                raise
+            finally:
+                invocation = getattr(self.llm, "last_request", None)
+                if invocation is not None:
+                    request_record["invocation"] = deepcopy(invocation)
+                resolved_provider = getattr(self.llm, "last_provider", None)
+                if resolved_provider is not None:
+                    request_record["resolved_provider"] = deepcopy(resolved_provider)
+            assistant = response.assistant_message()
+            request_record["response"] = deepcopy(assistant)
+            messages.append(assistant)
+            if not response.tool_calls:
+                return StepResult(
+                    step_id=step["id"],
+                    user_message=user_message,
+                    assistant_message=assistant,
+                    tool_calls=records,
+                    messages=deepcopy(messages),
+                )
+            for call in response.tool_calls:
+                call_index = registry.next_call_index(call.name)
+                step_call_counts[call.name] = step_call_counts.get(call.name, 0) + 1
+                step_call_index = step_call_counts[call.name]
+                context = ToolMockContext(
+                    plugin_root=case.plugin_root,
+                    case_key=case.case_key,
+                    run_index=run_index,
+                    step_id=step["id"],
+                    step_index=step["index"],
+                    request_index=request_index,
+                    tool_call_id=call.id,
+                    call_index=call_index,
+                    transcript=tuple(deepcopy(messages)),
+                    state=registry.state,
+                )
+                result = await registry.call(call.name, call.arguments, context)
+                record = ToolCallRecord(
+                    step_id=step["id"],
+                    request_index=request_index,
+                    call_index=call_index,
+                    call=call,
+                    result=result,
+                    is_error=isinstance(result, dict) and "error" in result,
+                )
+                records.append(record)
+                messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": call.id,
+                        "content": _tool_result_content(result),
+                    }
+                )
+                for item in injections:
+                    if item["position"] != "after_tool_call":
+                        continue
+                    match = item["match"]
+                    if match["name"] != call.name:
+                        continue
+                    if "call" in match and match["call"] != step_call_index:
+                        continue
+                    if "arguments" in match and not _matches(
+                        call.arguments, match["arguments"]
+                    ):
+                        continue
+                    messages.append(deepcopy(item["message"]))
+            messages.extend(
+                deepcopy(item["message"])
+                for item in injections
+                if item["position"] == "after_all_tool_calls"
+            )
+        raise MaxToolRoundsExceeded(step["id"], records)
+
+    async def _run_builders(
+        self,
+        builders: Any,
+        *,
+        position: str,
+        case: ResolvedCase,
+        transcript: list[dict[str, Any]],
+        state: dict[str, Any],
+        run_index: int,
+    ) -> list[dict[str, Any]]:
+        output: list[dict[str, Any]] = []
+        for builder in builders or []:
+            if builder.get("position") != position:
+                continue
+            use = builder.get("use")
+            if use == "message":
+                produced = [deepcopy(builder["message"])]
+            elif use == "python":
+                target = load_python_callable(
+                    builder["callable"],
+                    case.plugin_root,
+                    relative_root=case.evals_root,
+                )
+                context = {
+                    "plugin_root": case.plugin_root,
+                    "case_key": case.case_key,
+                    "run_index": run_index,
+                    "step_id": None,
+                    "step_index": None,
+                    "request_index": 0,
+                    "transcript": tuple(deepcopy(transcript)),
+                    "tool_state": state,
+                }
+                value = target(
+                    context=context, config=deepcopy(builder.get("with", {}))
+                )
+                produced = await value if inspect.isawaitable(value) else value
+                if not isinstance(produced, list):
+                    raise TypeError("Python builder must return a message list")
+            else:
+                raise ValueError(f"unsupported builder type: {use}")
+            output.extend(produced)
+        return output
+
+    def _case_passed(self, case: ResolvedCase, steps: list[StepResult]) -> bool:
+        if not all(step.hard_pass for step in steps):
+            return False
+        judge_results = [
+            step.judge_result for step in steps if step.judge_result is not None
+        ]
+        if any(not result.get("passed", False) for result in judge_results):
+            return False
+        requirement = case.outcome.get("require", {})
+        judge_requirement = requirement.get("judge", {})
+        if judge_requirement and not judge_results:
+            return False
+        if "overall_gte" in judge_requirement:
+            overall = sum(result["overall"] for result in judge_results) / len(
+                judge_results
+            )
+            if overall < judge_requirement["overall_gte"]:
+                return False
+        if "each_dimension_gte" in judge_requirement:
+            threshold = judge_requirement["each_dimension_gte"]
+            for result in judge_results:
+                if any(
+                    (item["score"] - 1) / 4 < threshold for item in result["dimensions"]
+                ):
+                    return False
+        return True
+
+    def _artifact_dir(self, case: ResolvedCase, run_index: int) -> Path | None:
+        if not self.write_artifacts:
+            return None
+        root = self.result_root or case.plugin_root / ".yama" / "runs"
+        key = sanitize_case_key(case.case_key)
+        return root / key / f"run-{run_index:03d}"
+
+    def _write_artifacts(self, case: ResolvedCase, run: RunResult) -> None:
+        assert run.artifact_dir is not None
+        directory = run.artifact_dir
+        if directory.exists():
+            shutil.rmtree(directory)
+        directory.mkdir(parents=True, exist_ok=True)
+        fixture_entries = deepcopy(case.fixture_manifest)
+        fixture_dir = directory / "fixtures"
+        for entry in fixture_entries:
+            source = Path(entry["fixture"])
+            fixture_dir.mkdir(exist_ok=True)
+            target = fixture_dir / f"{entry['sha256'][:12]}-{entry['filename']}"
+            shutil.copy2(source, target)
+            entry["materialized_path"] = str(target)
+        resolved = {
+            "case_key": case.case_key,
+            "path": str(case.path),
+            "plugin_root": str(case.plugin_root),
+            "evals_root": str(case.evals_root),
+            "model": case.model,
+            "execution": case.execution,
+            "context": case.context,
+            "mocks": case.mocks,
+            "steps": case.steps,
+            "outcome": case.outcome,
+        }
+        (directory / "case.resolved.yaml").write_text(
+            yaml.safe_dump(resolved, allow_unicode=True, sort_keys=False),
+            encoding="utf-8",
+        )
+        (directory / "rendered-system-prompt.txt").write_text(
+            case.system_prompt, encoding="utf-8"
+        )
+        self._write_json(
+            directory / "loaded-skills.json", [s.to_dict() for s in case.skills]
+        )
+        self._write_json(directory / "fixtures.json", fixture_entries)
+        request_dir = directory / "requests"
+        request_dir.mkdir(exist_ok=True)
+        for index, request in enumerate(run.requests, 1):
+            self._write_json(request_dir / f"request-{index:03d}.json", request)
+        self._write_json(
+            directory / "transcript.json", [step.to_dict() for step in run.steps]
+        )
+        self._write_json(directory / "mock-events.json", run.mock_events)
+        self._write_json(directory / "mock-state.json", run.mock_state)
+        self._write_json(
+            directory / "hard-checks.json",
+            {step.step_id: step.hard_checks for step in run.steps},
+        )
+        self._write_json(
+            directory / "judge-result.json",
+            {
+                step.step_id: step.judge_result
+                for step in run.steps
+                if step.judge_result
+            },
+        )
+        result = run.to_dict()
+        result["generated_at"] = datetime.now(timezone.utc).isoformat()
+        if run.traceback:
+            result["traceback"] = run.traceback
+        self._write_json(directory / "result.json", result)
+
+    @staticmethod
+    def _write_json(path: Path, value: Any) -> None:
+        path.write_text(
+            json.dumps(value, ensure_ascii=False, indent=2, default=str) + "\n",
+            encoding="utf-8",
+        )

yama/tools/base.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Mapping
+
+from ..mocks import ToolMockContext
+
+
+class Tool(ABC):
+    """Common interface implemented by every built-in Tool Mock handler."""
+
+    @abstractmethod
+    def run(
+        self,
+        arguments: Mapping[str, Any],
+        context: ToolMockContext | None,
+        config: Mapping[str, Any],
+    ) -> Any:
+        raise NotImplementedError

yama/tools/bash/_cli_shim.py

@@ -0,0 +1,71 @@
+"""Standalone resolver invoked by compiled mocks.cli shell scripts.
+
+Not imported by the framework at runtime: fs.compile_cli_config writes
+a tiny shell wrapper into the sandbox's shadow PATH for every mocked command,
+and that wrapper runs `python -m yama.tools.bash._cli_shim <rules.json> "$@"`. This
+module reuses the exact rule-tree walk and _matches() semantics the old
+in-process bash interpreter used, so mocks.cli's matching behavior (including
+regex fidelity) is unchanged even though a real shell now owns command-line
+parsing instead of a hand-rolled tokenizer.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from typing import Any
+
+from ...mocks import _matches
+
+
+def _resolve(node: dict[str, Any], argv: list[str]) -> dict[str, Any] | None:
+    remaining = list(argv)
+    fallback = node.get("default")
+    while True:
+        if node.get("default") is not None:
+            fallback = node["default"]
+        commands = node.get("commands") or {}
+        if remaining and remaining[0] in commands:
+            node = commands[remaining[0]]
+            remaining = remaining[1:]
+            continue
+        break
+    view = {"argv": remaining, "argline": " ".join(remaining)}
+    for rule in node.get("rules") or []:
+        match = rule.get("match") or {}
+        if all(_matches(view.get(key), expected) for key, expected in match.items()):
+            return rule
+    return fallback
+
+
+def main(argv: list[str]) -> int:
+    if not argv:
+        sys.stderr.write("yama-cli-shim: missing rules file argument\n")
+        return 70
+    rules_path, *rest = argv
+    with open(rules_path, "r", encoding="utf-8") as handle:
+        spec = json.load(handle)
+    name = spec["name"]
+    rule = _resolve(spec["node"], rest)
+    if rule is None:
+        sys.stderr.write(f"{name}: unknown error")
+        return 1
+    if rule.get("passthrough"):
+        env = dict(os.environ)
+        env["PATH"] = os.environ.get("YAMA_REAL_PATH", env.get("PATH", ""))
+        try:
+            os.execvpe(name, [name, *rest], env)
+        except FileNotFoundError:
+            sys.stderr.write(f"bash: {name}: command not found\n")
+            return 127
+        except PermissionError:
+            sys.stderr.write(f"bash: {name}: permission denied\n")
+            return 126
+    sys.stdout.write(rule.get("stdout", ""))
+    sys.stderr.write(rule.get("stderr", ""))
+    return int(rule.get("exit_code", 0))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))