python-yama 0.1.0__py3-none-any.whl

yama/mocks.py

@@ -0,0 +1,383 @@
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import inspect
+import json
+import time
+from copy import deepcopy
+from dataclasses import dataclass
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Callable, Sequence
+
+
+class ToolMockError(Exception):
+    def __init__(
+        self,
+        code: str,
+        message: str,
+        details: Any = None,
+        retryable: bool = False,
+    ) -> None:
+        super().__init__(message)
+        self.code = code
+        self.message = message
+        self.details = details
+        self.retryable = retryable
+
+    def as_result(self) -> dict[str, Any]:
+        return {
+            "error": {
+                "code": self.code,
+                "message": self.message,
+                "details": self.details,
+                "retryable": self.retryable,
+            }
+        }
+
+
+class MockConfigurationError(ValueError):
+    pass
+
+
+class MockInvocationError(RuntimeError):
+    pass
+
+
+class MockImplementationError(RuntimeError):
+    pass
+
+
+@dataclass
+class ToolMockContext:
+    plugin_root: Path
+    case_key: str
+    run_index: int
+    step_id: str
+    step_index: int
+    request_index: int
+    tool_call_id: str
+    call_index: int
+    transcript: Sequence[dict[str, Any]]
+    state: dict[str, Any]
+
+
+def load_python_callable(
+    reference: str,
+    plugin_root: Path,
+    *,
+    relative_root: Path | None = None,
+) -> Callable[..., Any]:
+    if ":" not in reference:
+        raise MockConfigurationError(f"Python callable must contain ':': {reference}")
+    location, attribute = reference.rsplit(":", 1)
+    if not attribute:
+        raise MockConfigurationError(f"Python callable name is missing: {reference}")
+    module: ModuleType
+    if location.endswith(".py") or "/" in location or "\\" in location:
+        path = ((relative_root or plugin_root) / location).resolve()
+        if not path.is_relative_to(plugin_root.resolve()):
+            raise MockConfigurationError(
+                f"Python mock must be inside plugin root: {path}"
+            )
+        if not path.is_file() or path.suffix != ".py":
+            raise MockConfigurationError(f"Python mock file does not exist: {path}")
+        module_name = f"yama_plugin_mock_{abs(hash(path))}"
+        spec = importlib.util.spec_from_file_location(module_name, path)
+        if spec is None or spec.loader is None:
+            raise MockConfigurationError(f"cannot load Python mock: {path}")
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+    else:
+        try:
+            module = importlib.import_module(location)
+        except ImportError as exc:
+            raise MockConfigurationError(
+                f"cannot import Python mock module: {location}"
+            ) from exc
+    target = getattr(module, attribute, None)
+    if not callable(target):
+        raise MockConfigurationError(f"Python mock target is not callable: {reference}")
+    return target
+
+
+def _validate_handler_signature(target: Callable[..., Any], reference: str) -> None:
+    signature = inspect.signature(target)
+    try:
+        signature.bind(arguments={}, context=None, config={})
+    except TypeError as exc:
+        raise MockConfigurationError(
+            f"Python mock {reference} must accept keyword arguments, context and config"
+        ) from exc
+
+
+def _json_clone(value: Any, label: str) -> Any:
+    try:
+        return json.loads(json.dumps(value, ensure_ascii=False))
+    except (TypeError, ValueError) as exc:
+        raise MockImplementationError(
+            f"{label} must be JSON serializable: {exc}"
+        ) from exc
+
+
+def _matches(actual: Any, expected: Any) -> bool:
+    if isinstance(expected, dict):
+        matcher_keys = {
+            "equals",
+            "contains",
+            "matches",
+            "is_instance",
+            "exists",
+            "one_of",
+        }
+        selected = matcher_keys.intersection(expected)
+        if selected:
+            if len(selected) != 1:
+                raise MockConfigurationError("a matcher must use exactly one operation")
+            operation = next(iter(selected))
+            wanted = expected[operation]
+            if operation == "equals":
+                return actual == wanted
+            if operation == "contains":
+                return (
+                    wanted in actual if isinstance(actual, (str, list, dict)) else False
+                )
+            if operation == "one_of":
+                return isinstance(wanted, list) and actual in wanted
+            if operation == "matches":
+                import re
+
+                return (
+                    isinstance(actual, str)
+                    and re.search(str(wanted), actual) is not None
+                )
+            if operation == "exists":
+                return (actual is not _MISSING) is bool(wanted)
+            types = {
+                "array": list,
+                "object": dict,
+                "string": str,
+                "number": (int, float),
+                "integer": int,
+                "boolean": bool,
+                "null": type(None),
+            }
+            return wanted in types and isinstance(actual, types[wanted])
+        if not isinstance(actual, dict):
+            return False
+        return all(
+            _matches(actual.get(key, _MISSING), value)
+            for key, value in expected.items()
+        )
+    if isinstance(expected, list):
+        return actual == expected
+    return actual == expected
+
+
+_MISSING = object()
+
+
+def _state_diff(before: Any, after: Any, path: str = "$") -> list[dict[str, Any]]:
+    if before == after:
+        return []
+    if isinstance(before, dict) and isinstance(after, dict):
+        changes: list[dict[str, Any]] = []
+        for key in sorted(set(before) | set(after)):
+            child_path = f"{path}.{key}"
+            if key not in before:
+                changes.append({"op": "add", "path": child_path, "value": after[key]})
+            elif key not in after:
+                changes.append({"op": "remove", "path": child_path, "old": before[key]})
+            else:
+                changes.extend(_state_diff(before[key], after[key], child_path))
+        return changes
+    return [{"op": "replace", "path": path, "old": before, "value": after}]
+
+
+class MockToolRegistry:
+    def __init__(
+        self,
+        config: dict[str, Any] | None,
+        *,
+        plugin_root: Path,
+        relative_root: Path | None = None,
+        exposed_tools: set[str],
+    ) -> None:
+        config = config or {}
+        if not isinstance(config, dict):
+            raise MockConfigurationError("mocks must be an object")
+        self.plugin_root = plugin_root
+        self.relative_root = relative_root or plugin_root
+        unknown = set(config) - {"state", "tools"}
+        if unknown:
+            hint = (
+                "; undeclared tools and unmatched arguments always fail"
+                if "strict" in unknown
+                else ""
+            )
+            raise MockConfigurationError(
+                f"unknown mock fields: {sorted(unknown)}{hint}"
+            )
+        initial_state = config.get("state", {})
+        if not isinstance(initial_state, dict):
+            raise MockConfigurationError("mocks.state must be an object")
+        self.state: dict[str, Any] = _json_clone(initial_state, "mocks.state")
+        tools = config.get("tools", {})
+        if not isinstance(tools, dict):
+            raise MockConfigurationError("mocks.tools must be an object")
+        extra = set(tools) - exposed_tools
+        missing = exposed_tools - set(tools)
+        if extra:
+            raise MockConfigurationError(
+                f"mock handlers are not exposed in context.tools: {sorted(extra)}"
+            )
+        if missing:
+            raise MockConfigurationError(
+                f"context tools have no mock handlers: {sorted(missing)}"
+            )
+        self.tools: dict[str, dict[str, Any]] = deepcopy(tools)
+        self.call_counts = {name: 0 for name in tools}
+        self.sequence_indexes = {name: 0 for name in tools}
+        self.events: list[dict[str, Any]] = []
+        self.handlers: dict[str, Callable[..., Any]] = {}
+        for name, tool_config in self.tools.items():
+            self._validate_tool_config(name, tool_config)
+
+    def _validate_tool_config(self, name: str, config: Any) -> None:
+        if not isinstance(config, dict):
+            raise MockConfigurationError(f"mock config must be an object: {name}")
+        unknown = set(config) - {"match", "respond", "handler"}
+        if unknown:
+            hint = (
+                "; tool-call expectations belong in assert.hard (tool_called)"
+                if "required" in unknown
+                else ""
+            )
+            raise MockConfigurationError(
+                f"unknown mock fields for {name}: {sorted(unknown)}{hint}"
+            )
+        if ("respond" in config) == ("handler" in config):
+            raise MockConfigurationError(
+                f"mock {name} must choose exactly one of respond or handler"
+            )
+        if "handler" in config:
+            handler = config["handler"]
+            if not isinstance(handler, dict) or not isinstance(
+                handler.get("python"), str
+            ):
+                raise MockConfigurationError(f"mock {name} handler.python is required")
+            unknown = set(handler) - {"python", "with"}
+            if unknown:
+                raise MockConfigurationError(
+                    f"unknown handler fields for {name}: {sorted(unknown)}"
+                )
+            target = load_python_callable(
+                handler["python"],
+                self.plugin_root,
+                relative_root=self.relative_root,
+            )
+            _validate_handler_signature(target, handler["python"])
+            self.handlers[name] = target
+        else:
+            respond = config["respond"]
+            if not isinstance(respond, dict):
+                raise MockConfigurationError(f"mock {name} respond must be an object")
+            modes = set(respond).intersection({"result", "error", "sequence"})
+            if len(modes) != 1:
+                raise MockConfigurationError(
+                    f"mock {name} respond must choose exactly one of result, error or sequence"
+                )
+            if "sequence" in respond and (
+                not isinstance(respond["sequence"], list) or not respond["sequence"]
+            ):
+                raise MockConfigurationError(f"mock {name} sequence must be non-empty")
+
+    def next_call_index(self, name: str) -> int:
+        return self.call_counts.get(name, 0) + 1
+
+    async def call(
+        self,
+        name: str,
+        arguments: dict[str, Any],
+        context: ToolMockContext,
+    ) -> Any:
+        if name not in self.tools:
+            raise MockInvocationError(f"tool has no mock handler: {name}")
+        config = self.tools[name]
+        expected = config.get("match", {}).get("arguments", {})
+        match_failed = bool(expected) and not _matches(arguments, expected)
+        self.call_counts[name] += 1
+        before = _json_clone(self.state, "mock state")
+        started = time.perf_counter()
+        is_error = False
+        try:
+            if match_failed:
+                result = {
+                    "error": {
+                        "code": "UNKNOWN_ERROR",
+                        "message": "unknown error",
+                        "details": None,
+                        "retryable": False,
+                    }
+                }
+                is_error = True
+            elif "handler" in config:
+                handler = config["handler"]
+                target = self.handlers[name]
+                value = target(
+                    arguments=deepcopy(arguments),
+                    context=context,
+                    config=deepcopy(handler.get("with", {})),
+                )
+                result = await value if inspect.isawaitable(value) else value
+            else:
+                result = self._declarative_result(name, config["respond"])
+        except ToolMockError as exc:
+            result = exc.as_result()
+            is_error = True
+        except (MockConfigurationError, MockInvocationError, MockImplementationError):
+            raise
+        except Exception as exc:
+            raise MockImplementationError(f"Python mock {name} failed: {exc}") from exc
+        result = _json_clone(result, f"mock result for {name}")
+        after = _json_clone(self.state, "mock state")
+        event = {
+            "tool": name,
+            "tool_call_id": context.tool_call_id,
+            "step_id": context.step_id,
+            "request_index": context.request_index,
+            "call_index": context.call_index,
+            "handler_kind": "python" if "handler" in config else "declarative",
+            "callable": config.get("handler", {}).get("python"),
+            "match_failed": match_failed,
+            "arguments": arguments,
+            "result": result,
+            "is_error": is_error or (isinstance(result, dict) and "error" in result),
+            "duration_ms": round((time.perf_counter() - started) * 1000, 3),
+            "state_diff": _state_diff(before, after),
+        }
+        self.events.append(event)
+        return result
+
+    def _declarative_result(self, name: str, respond: dict[str, Any]) -> Any:
+        item = respond
+        if "sequence" in respond:
+            index = self.sequence_indexes[name]
+            if index >= len(respond["sequence"]):
+                raise MockInvocationError(f"mock response sequence exhausted: {name}")
+            item = respond["sequence"][index]
+            self.sequence_indexes[name] += 1
+        if "result" in item:
+            return deepcopy(item["result"])
+        error = item.get("error")
+        if not isinstance(error, dict):
+            raise MockConfigurationError(f"mock error must be an object: {name}")
+        return {
+            "error": {
+                "code": error.get("code", "TOOL_ERROR"),
+                "message": error.get("message", "tool failed"),
+                "details": error.get("details"),
+                "retryable": bool(error.get("retryable", False)),
+            }
+        }

yama/models.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, NotRequired, TypedDict
+
+
+Json = dict[str, Any]
+
+
+class CheckerResult(TypedDict):
+    """Structured return value of a `python` hard-assertion checker.
+
+    Only `passed` is required. The optional fields refine the diagnostic
+    shown in reports; omitted ones fall back as implemented in
+    `assertions.py` (`detail` doubles as the failure reason by default).
+    Checkers may also return a bare bool or a `(passed, detail)` tuple.
+    """
+
+    passed: bool
+    detail: NotRequired[str]
+    label: NotRequired[str]
+    expected: NotRequired[str]
+    actual: NotRequired[str]
+    reason: NotRequired[str]
+
+
+def sanitize_case_key(case_key: str) -> str:
+    key = case_key.replace("/", "-").replace("\\", "-")
+    if key.endswith(".yaml"):
+        key = key[:-5]
+    elif key.endswith(".yml"):
+        key = key[:-4]
+    return key
+
+
+@dataclass(frozen=True)
+class ToolSchema:
+    name: str
+    description: str
+    parameters: Json
+    source_path: Path | None = None
+
+    def for_provider(self) -> Json:
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters,
+            },
+        }
+
+@dataclass(frozen=True)
+class ToolCall:
+    id: str
+    name: str
+    arguments: Json
+
+
+@dataclass
+class LLMResponse:
+    content: str = ""
+    tool_calls: list[ToolCall] = field(default_factory=list)
+    raw: Any = None
+
+    def assistant_message(self) -> Json:
+        message: Json = {"role": "assistant", "content": self.content or ""}
+        if self.tool_calls:
+            message["tool_calls"] = [
+                {
+                    "id": call.id,
+                    "type": "function",
+                    "function": {
+                        "name": call.name,
+                        "arguments": call.arguments,
+                    },
+                }
+                for call in self.tool_calls
+            ]
+        return message
+
+
+@dataclass
+class ToolCallRecord:
+    step_id: str
+    request_index: int
+    call_index: int
+    call: ToolCall
+    result: Any
+    is_error: bool = False
+
+    def to_dict(self) -> Json:
+        return {
+            "step_id": self.step_id,
+            "request_index": self.request_index,
+            "call_index": self.call_index,
+            "call": asdict(self.call),
+            "result": self.result,
+            "is_error": self.is_error,
+        }
+
+
+@dataclass
+class StepResult:
+    step_id: str
+    user_message: Json
+    assistant_message: Json
+    tool_calls: list[ToolCallRecord]
+    messages: list[Json]
+    hard_checks: list[Json] = field(default_factory=list)
+    hard_pass: bool = True
+    judge_result: Json | None = None
+
+    def to_dict(self) -> Json:
+        return {
+            "step_id": self.step_id,
+            "user_message": self.user_message,
+            "assistant_message": self.assistant_message,
+            "tool_calls": [record.to_dict() for record in self.tool_calls],
+            "messages": self.messages,
+            "hard_checks": self.hard_checks,
+            "hard_pass": self.hard_pass,
+            "judge_result": self.judge_result,
+        }
+
+
+@dataclass
+class LoadedSkill:
+    requested_name: str
+    name: str
+    description: str
+    body: str
+    path: Path
+    content_hash: str
+
+    def to_dict(self) -> Json:
+        data = asdict(self)
+        data["path"] = str(self.path)
+        return data
+
+
+@dataclass
+class ResolvedCase:
+    path: Path
+    case_key: str
+    plugin_root: Path
+    evals_root: Path
+    model: Json
+    context: Json
+    mocks: Json
+    execution: Json
+    steps: list[Json]
+    outcome: Json
+    system_prompt: str
+    skills: list[LoadedSkill]
+    tools: list[ToolSchema]
+    initial_messages: list[Json]
+    fixture_manifest: list[Json] = field(default_factory=list)