python-yama 0.1.0__py3-none-any.whl

yama/assertions.py

@@ -0,0 +1,577 @@
+from __future__ import annotations
+
+import inspect
+import json
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Any
+
+from .llm import LLMClient
+from .mocks import load_python_callable
+from .models import StepResult, ToolCallRecord
+
+
+class AssertionConfigurationError(ValueError):
+    pass
+
+
+_MISSING = object()
+
+
+def _select_jsonpath(value: Any, path: str) -> Any:
+    if path == "$":
+        return value
+    if not path.startswith("$."):
+        raise AssertionConfigurationError(f"unsupported JSONPath: {path}")
+    current = value
+    for name, raw_index in re.findall(r"([^.\[\]]+)(?:\[(\d+)\])?", path[2:]):
+        if not isinstance(current, dict) or name not in current:
+            return _MISSING
+        current = current[name]
+        if raw_index:
+            index = int(raw_index)
+            if not isinstance(current, list) or index >= len(current):
+                return _MISSING
+            current = current[index]
+    return current
+
+
+def _match(value: Any, config: dict[str, Any]) -> bool:
+    operations = [
+        operation
+        for operation in ("equals", "contains", "matches", "is_instance", "exists")
+        if operation in config
+    ]
+    if len(operations) != 1:
+        raise AssertionConfigurationError("matcher requires exactly one operation")
+    operation = operations[0]
+    expected = config[operation]
+    if operation == "equals":
+        return value == expected
+    if operation == "contains":
+        return expected in value if isinstance(value, (str, list, dict)) else False
+    if operation == "matches":
+        return isinstance(value, str) and re.search(str(expected), value) is not None
+    if operation == "exists":
+        return (value is not _MISSING) is bool(expected)
+    types = {
+        "array": list,
+        "object": dict,
+        "string": str,
+        "number": (int, float),
+        "integer": int,
+        "boolean": bool,
+        "null": type(None),
+    }
+    return expected in types and isinstance(value, types[expected])
+
+
+def _records(result: StepResult, name: str | None = None) -> list[ToolCallRecord]:
+    return [
+        record
+        for record in result.tool_calls
+        if name is None or record.call.name == name
+    ]
+
+
+def _tool_count_pass(count: int, config: dict[str, Any]) -> bool:
+    if "times" in config:
+        if "min_times" in config or "max_times" in config:
+            raise AssertionConfigurationError(
+                "times is mutually exclusive with min_times/max_times"
+            )
+        return count == config["times"]
+    minimum = config.get("min_times", 1)
+    maximum = config.get("max_times")
+    return count >= minimum and (maximum is None or count <= maximum)
+
+
+def _diagnostic(
+    *,
+    passed: bool,
+    label: str,
+    expected: str,
+    actual: str,
+    reason: str,
+    detail: str | None = None,
+) -> dict[str, Any]:
+    return {
+        "passed": passed,
+        "label": label,
+        "expected": expected,
+        "actual": actual,
+        "reason": reason,
+        "detail": detail or (actual if passed else reason),
+    }
+
+
+def _render_value(value: Any, *, limit: int = 240) -> str:
+    if value is _MISSING:
+        return "<missing>"
+    rendered = repr(value)
+    if len(rendered) <= limit:
+        return rendered
+    return f"{rendered[: limit - 3]}..."
+
+
+def _count_expectation(config: dict[str, Any]) -> str:
+    if "times" in config:
+        return f"exactly {config['times']} time(s)"
+    minimum = config.get("min_times", 1)
+    maximum = config.get("max_times")
+    if maximum is None:
+        return f"at least {minimum} time(s)"
+    return f"between {minimum} and {maximum} time(s)"
+
+
+def _matcher(config: dict[str, Any]) -> tuple[str, Any, str]:
+    operations = [
+        operation
+        for operation in ("equals", "contains", "matches", "is_instance", "exists")
+        if operation in config
+    ]
+    if len(operations) != 1:
+        raise AssertionConfigurationError("matcher requires exactly one operation")
+    operation = operations[0]
+    expected = config[operation]
+    descriptions = {
+        "equals": f"equal {_render_value(expected)}",
+        "contains": f"contain {_render_value(expected)}",
+        "matches": f"match regex {_render_value(str(expected))}",
+        "is_instance": f"be an instance of {_render_value(expected)}",
+        "exists": "exist" if expected else "be absent",
+    }
+    return operation, expected, descriptions[operation]
+
+
+def _selected_values(values: list[tuple[int, Any]]) -> str:
+    if not values:
+        return "no values selected"
+    return "; ".join(
+        f"call #{call_number}: {_render_value(value)}"
+        for call_number, value in values
+    )
+
+
+async def run_hard_assertions(
+    assertions: list[Any],
+    *,
+    result: StepResult,
+    transcript: list[dict[str, Any]],
+    state: dict[str, Any],
+    plugin_root: Path,
+    relative_root: Path | None = None,
+) -> list[dict[str, Any]]:
+    checks: list[dict[str, Any]] = []
+    for assertion in assertions:
+        if not isinstance(assertion, dict) or len(assertion) != 1:
+            raise AssertionConfigurationError(
+                "each hard assertion must contain exactly one type"
+            )
+        kind, config = next(iter(assertion.items()))
+        diagnostic = await _run_one(
+            kind,
+            config,
+            result=result,
+            transcript=transcript,
+            state=state,
+            plugin_root=plugin_root,
+            relative_root=relative_root,
+        )
+        checks.append({"type": kind, **diagnostic})
+    return checks
+
+
+async def _run_one(
+    kind: str,
+    config: Any,
+    *,
+    result: StepResult,
+    transcript: list[dict[str, Any]],
+    state: dict[str, Any],
+    plugin_root: Path,
+    relative_root: Path | None,
+) -> dict[str, Any]:
+    assistant = str(result.assistant_message.get("content", ""))
+    if kind == "tool_called":
+        if isinstance(config, str):
+            config = {"name": config}
+        records = _records(result, config["name"])
+        passed = _tool_count_pass(len(records), config)
+        expected_count = _count_expectation(config)
+        return _diagnostic(
+            passed=passed,
+            label=f"tool {config['name']!r} call count",
+            expected=f"tool {config['name']!r} called {expected_count}",
+            actual=f"tool {config['name']!r} called {len(records)} time(s)",
+            reason=(
+                "call count satisfied the configured expectation"
+                if passed
+                else f"expected {expected_count}, but observed {len(records)} time(s)"
+            ),
+            detail=f"{config['name']} called {len(records)} time(s)",
+        )
+    if kind == "tool_not_called":
+        name = config if isinstance(config, str) else config["name"]
+        count = len(_records(result, name))
+        passed = count == 0
+        return _diagnostic(
+            passed=passed,
+            label=f"tool {name!r} is not called",
+            expected=f"tool {name!r} called 0 time(s)",
+            actual=f"tool {name!r} called {count} time(s)",
+            reason=(
+                "tool was not called"
+                if passed
+                else f"tool {name!r} was called {count} unexpected time(s)"
+            ),
+            detail=f"{name} called {count} time(s)",
+        )
+    if kind == "tool_arguments":
+        name = config["name"]
+        records = _records(result, name)
+        call_number = config.get("call")
+        if call_number is not None:
+            selected_records = (
+                [(call_number, records[call_number - 1])]
+                if 0 < call_number <= len(records)
+                else []
+            )
+        else:
+            selected_records = list(enumerate(records, 1))
+        values = [
+            (number, _select_jsonpath(record.call.arguments, config["jsonpath"]))
+            for number, record in selected_records
+        ]
+        _, _, matcher_description = _matcher(config)
+        passed = any(_match(value, config) for _, value in values)
+        scope = f"call #{call_number}" if call_number is not None else "any call"
+        expected = (
+            f"{scope} to tool {name!r} has {config['jsonpath']} that "
+            f"{matcher_description}"
+        )
+        actual = _selected_values(values)
+        if passed:
+            reason = "at least one selected argument value satisfied the matcher"
+        elif not records:
+            reason = f"tool {name!r} was never called, so no arguments could be checked"
+        elif call_number is not None and not selected_records:
+            reason = (
+                f"tool {name!r} had {len(records)} call(s), so call #{call_number} "
+                "does not exist"
+            )
+        elif all(value is _MISSING for _, value in values):
+            reason = (
+                f"JSONPath {config['jsonpath']!r} was missing from every selected "
+                f"call to tool {name!r}"
+            )
+        else:
+            reason = (
+                f"none of the {len(values)} selected value(s) "
+                f"{matcher_description}"
+            )
+        return _diagnostic(
+            passed=passed,
+            label=(
+                f"tool {name!r} argument {config['jsonpath']} must "
+                f"{matcher_description}"
+            ),
+            expected=expected,
+            actual=actual,
+            reason=reason,
+            detail=f"selected values: {actual}",
+        )
+    if kind == "tool_call_order":
+        if ("names" in config) == ("groups" in config):
+            raise AssertionConfigurationError(
+                "tool_call_order must choose names or groups"
+            )
+        actual_calls = [record.call.name for record in result.tool_calls]
+        if "names" in config:
+            expected = config["names"]
+            if config.get("exact", False):
+                passed = actual_calls == expected
+                expectation = f"exact tool sequence {expected!r}"
+            else:
+                cursor = iter(actual_calls)
+                passed = all(
+                    any(item == wanted for item in cursor) for wanted in expected
+                )
+                expectation = f"ordered tool subsequence {expected!r}"
+            return _diagnostic(
+                passed=passed,
+                label="tool call order",
+                expected=expectation,
+                actual=f"tool sequence {actual_calls!r}",
+                reason=(
+                    "tool sequence satisfied the configured order"
+                    if passed
+                    else f"expected {expectation}, but it was not found in the actual sequence"
+                ),
+                detail=f"actual order: {actual_calls}",
+            )
+        groups = config["groups"]
+        expected_groups = " -> ".join(repr(group) for group in groups)
+        # Existential matching: the assertion passes when one call per group
+        # member can be picked such that every pick of a later group lands in
+        # a strictly later request round than the previous group's completion.
+        # Calls outside the picked ones (repeats, interleaved extras) never
+        # fail the order on their own. Greedily picking the earliest eligible
+        # round per member minimizes each group's completion round, so if the
+        # greedy walk fails no other assignment can succeed.
+        completions: list[int] = []
+        boundary = -1
+        failure_reason = ""
+        failure_detail = ""
+        for index, group in enumerate(groups, 1):
+            rounds_by_name = {
+                name: [
+                    record.request_index
+                    for record in result.tool_calls
+                    if record.call.name == name and record.request_index > boundary
+                ]
+                for name in group
+            }
+            missing = [
+                name
+                for name, rounds in rounds_by_name.items()
+                if not rounds and not _records(result, name)
+            ]
+            if missing:
+                failure_reason = (
+                    f"group #{index} was incomplete; missing tool(s): {missing!r}"
+                )
+                failure_detail = f"group not fully called: {group}"
+                break
+            late = [name for name, rounds in rounds_by_name.items() if not rounds]
+            if late:
+                failure_reason = (
+                    f"group #{index} could not be matched: {late!r} had no call "
+                    f"after request {boundary}, where group #{index - 1} completed"
+                )
+                failure_detail = f"group order not matchable for: {late}"
+                break
+            boundary = max(min(rounds) for rounds in rounds_by_name.values())
+            completions.append(boundary)
+        passed = not failure_reason
+        return _diagnostic(
+            passed=passed,
+            label="tool call order by groups",
+            expected=f"groups matchable in request order: {expected_groups}",
+            actual=(
+                f"matched group request indexes {completions!r}; "
+                f"tool sequence {actual_calls!r}"
+            ),
+            reason=(
+                "every group was matched by calls in strictly later requests "
+                "than the previous group"
+                if passed
+                else failure_reason
+            ),
+            detail=(
+                f"matched group request indexes: {completions}"
+                if passed
+                else failure_detail
+            ),
+        )
+    if kind == "assistant_contains":
+        target = str(config)
+        position = assistant.find(target)
+        passed = position >= 0
+        return _diagnostic(
+            passed=passed,
+            label=f"assistant contains {_render_value(target)}",
+            expected=f"assistant response contains literal text {_render_value(target)}",
+            actual=(
+                f"text found at character {position}"
+                if passed
+                else f"text not found in assistant response ({len(assistant)} characters)"
+            ),
+            reason=(
+                "required text was present"
+                if passed
+                else f"required literal text {_render_value(target)} was absent"
+            ),
+            detail=f"assistant contains {config!r}: {passed}",
+        )
+    if kind == "assistant_not_contains":
+        target = str(config)
+        position = assistant.find(target)
+        passed = position < 0
+        return _diagnostic(
+            passed=passed,
+            label=f"assistant omits {_render_value(target)}",
+            expected=f"assistant response omits literal text {_render_value(target)}",
+            actual=(
+                f"text not found in assistant response ({len(assistant)} characters)"
+                if passed
+                else f"text found at character {position}"
+            ),
+            reason=(
+                "forbidden text was absent"
+                if passed
+                else f"forbidden literal text {_render_value(target)} was present"
+            ),
+            detail=f"assistant omits {config!r}: {passed}",
+        )
+    if kind == "assistant_matches":
+        pattern = str(config)
+        match = re.search(pattern, assistant)
+        passed = match is not None
+        return _diagnostic(
+            passed=passed,
+            label="assistant matches regex",
+            expected=f"assistant response matches regex {_render_value(pattern)}",
+            actual=(
+                f"matched {_render_value(match.group(0))} at characters "
+                f"{match.start()}..{match.end()}"
+                if match
+                else f"no match in assistant response ({len(assistant)} characters)"
+            ),
+            reason=(
+                "regular expression matched the assistant response"
+                if passed
+                else f"regular expression {_render_value(pattern)} did not match"
+            ),
+            detail=f"assistant matches regex: {passed}",
+        )
+    if kind == "assistant_semantic_count":
+        target = config["target"]
+        count = len(re.findall(re.escape(target), assistant, flags=re.IGNORECASE))
+        if "equals" in config:
+            passed = count == config["equals"]
+            expected_count = f"exactly {config['equals']} occurrence(s)"
+        else:
+            passed = count >= config.get("min", 0) and count <= config.get("max", count)
+            minimum = config.get("min", 0)
+            maximum = config.get("max")
+            expected_count = (
+                f"at least {minimum} occurrence(s)"
+                if maximum is None
+                else f"between {minimum} and {maximum} occurrence(s)"
+            )
+        return _diagnostic(
+            passed=passed,
+            label=f"assistant semantic count for {target!r}",
+            expected=f"literal marker {target!r} occurs {expected_count}",
+            actual=f"literal marker {target!r} occurred {count} time(s)",
+            reason=(
+                "semantic marker count satisfied the configured expectation"
+                if passed
+                else f"expected {expected_count}, but observed {count} occurrence(s)"
+            ),
+            detail=f"literal semantic marker {target!r} occurred {count} time(s)",
+        )
+    if kind == "python":
+        if not isinstance(config, dict) or not isinstance(config.get("checker"), str):
+            raise AssertionConfigurationError("python assertion requires checker")
+        checker = load_python_callable(
+            config["checker"], plugin_root, relative_root=relative_root
+        )
+        value = checker(
+            step_result=result,
+            transcript=tuple(deepcopy(transcript)),
+            state=state,
+            config=deepcopy(config.get("with", {})),
+        )
+        value = await value if inspect.isawaitable(value) else value
+        if isinstance(value, dict):
+            passed = bool(value.get("passed"))
+            detail = str(value.get("detail", value))
+            return _diagnostic(
+                passed=passed,
+                label=str(value.get("label", f"python checker {config['checker']}")),
+                expected=str(value.get("expected", "custom checker returns passed=true")),
+                actual=str(value.get("actual", detail)),
+                reason=str(
+                    value.get(
+                        "reason",
+                        "custom checker passed" if passed else detail,
+                    )
+                ),
+                detail=detail,
+            )
+        if isinstance(value, tuple) and len(value) == 2:
+            passed, detail = bool(value[0]), str(value[1])
+        else:
+            passed, detail = bool(value), f"checker returned {value!r}"
+        return _diagnostic(
+            passed=passed,
+            label=f"python checker {config['checker']}",
+            expected="custom checker returns a passing result",
+            actual=detail,
+            reason="custom checker passed" if passed else detail,
+            detail=detail,
+        )
+    raise AssertionConfigurationError(f"unknown hard assertion: {kind}")
+
+
+async def run_judge(
+    config: dict[str, Any],
+    *,
+    client: LLMClient,
+    step_result: StepResult,
+    transcript_before_step: list[dict[str, Any]],
+    default_model: dict[str, Any],
+) -> dict[str, Any]:
+    dimensions = config.get("dimensions", [])
+    if not isinstance(dimensions, list) or not dimensions:
+        raise AssertionConfigurationError("judge.dimensions must be non-empty")
+    model = deepcopy(default_model)
+    if "model" in config:
+        model["name"] = config["model"]
+    model["temperature"] = config.get("temperature", 0)
+    prompt = {
+        "dimensions": dimensions,
+        "transcript_before_step": transcript_before_step,
+        "user_message": step_result.user_message,
+        "assistant_message": step_result.assistant_message,
+        "tool_calls": [record.to_dict() for record in step_result.tool_calls],
+    }
+    response = await client.complete(
+        messages=[
+            {
+                "role": "system",
+                "content": (
+                    "You are an evaluation judge. Score every requested dimension from 1 to 5. "
+                    'Return only JSON: {"dimensions":[{"name":str,"score":int,'
+                    '"reason":str,"evidence":str}]}'
+                ),
+            },
+            {"role": "user", "content": json.dumps(prompt, ensure_ascii=False)},
+        ],
+        tools=[],
+        model=model,
+    )
+    text = response.content.strip()
+    if text.startswith("```"):
+        text = re.sub(r"^```(?:json)?\s*|\s*```$", "", text, flags=re.DOTALL)
+    try:
+        judged = json.loads(text)
+    except json.JSONDecodeError as exc:
+        raise AssertionConfigurationError(
+            f"judge returned invalid JSON: {text}"
+        ) from exc
+    scores = {item["name"]: item for item in judged.get("dimensions", [])}
+    weighted = 0.0
+    total_weight = 0.0
+    passed = True
+    for dimension in dimensions:
+        item = scores.get(dimension["name"])
+        if (
+            not item
+            or not isinstance(item.get("score"), int)
+            or not 1 <= item["score"] <= 5
+        ):
+            raise AssertionConfigurationError(
+                f"judge omitted dimension: {dimension['name']}"
+            )
+        weight = float(dimension.get("weight", 1.0))
+        weighted += ((item["score"] - 1) / 4) * weight
+        total_weight += weight
+        if item["score"] < dimension.get("min_score", 1):
+            passed = False
+    judged["overall"] = weighted / total_weight
+    judged["passed"] = passed
+    return judged