workflowval 0.1.0__py3-none-any.whl

workflowval/loop/stats.py

@@ -0,0 +1,178 @@
+"""Counts of model and tool calls per workflow, and a step limit suggested from past runs."""
+
+from __future__ import annotations
+
+import math
+import re
+from collections import Counter
+from collections.abc import Collection, Iterable, Mapping
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from pydantic import ValidationError
+
+from workflowval.models.base import ToolCall
+from workflowval.schemas import CallCounts, RunStats, StatsTotal
+
+MIN_RUNS_FOR_SUGGESTION = 5
+RECENT_RUNS = 20
+MIN_SUGGESTED_LIMIT = 10
+UNKNOWN_TOOL = "<unknown>"
+MAX_STEPS_REASON = "max_steps"
+# Runs that ended for other reasons stopped early, so they'd make the limit look too generous.
+USABLE_STOP_REASONS = frozenset({"completed", MAX_STEPS_REASON, "max_iterations", "no_progress"})
+
+_WORKFLOW_ID = re.compile(r"[a-z0-9_]+")
+_TEST_FILE = re.compile(r"(?:^|/)test_([a-z0-9_]+)\.py(?:::|$)")
+
+
+@dataclass
+class _Tally:
+    model_calls: int = 0
+    cost_usd: float = 0.0
+    tools: Counter[str] = field(default_factory=Counter)
+
+    def counts(self) -> CallCounts:
+        return CallCounts(
+            model_calls=self.model_calls,
+            tool_calls=dict(sorted(self.tools.items())),
+            cost_usd=self.cost_usd,
+        )
+
+
+def workflow_of(arguments: Mapping[str, Any]) -> str | None:
+    """Return the workflow a tool call's arguments name, if they name one validly."""
+    workflow = arguments.get("workflow")
+    candidates = [
+        arguments.get("workflow_id"),
+        workflow.get("id") if isinstance(workflow, Mapping) else None,
+    ]
+    for candidate in candidates:
+        if isinstance(candidate, str) and _WORKFLOW_ID.fullmatch(candidate):
+            return candidate
+    test_id = arguments.get("test_id")
+    if isinstance(test_id, str) and (match := _TEST_FILE.search(test_id)):
+        return match.group(1)
+    return None
+
+
+class StatsCollector:
+    """Collects one run's counts and costs, per workflow and per role."""
+
+    def __init__(self, limit: int, spec_hash: str = "") -> None:
+        self.limit = limit
+        self.spec_hash = spec_hash
+        self._owners: dict[str | None, _Tally] = {}
+        self._roles: dict[str, _Tally] = {}
+
+    def _tallies(self, owner: str | None, role: str | None) -> list[_Tally]:
+        tallies = [self._owners.setdefault(owner, _Tally())]
+        if role is not None:
+            tallies.append(self._roles.setdefault(role, _Tally()))
+        return tallies
+
+    def record_failed_call(self, scope: str | None, role: str | None = None) -> None:
+        """Count a model call that failed, which still used a step."""
+        for tally in self._tallies(scope, role):
+            tally.model_calls += 1
+
+    def record_reply(
+        self,
+        scope: str | None,
+        calls: Iterable[ToolCall],
+        known_tools: Collection[str],
+        role: str | None = None,
+        cost_usd: float | None = None,
+    ) -> None:
+        """Count one model call, its cost and its tool calls.
+
+        A model can invent tool names and they aren't safe to store, so those count as unknown.
+        """
+        owners: set[str | None] = set()
+        for call in calls:
+            owner = scope or workflow_of(call.arguments)
+            owners.add(owner)
+            name = call.name if call.name in known_tools else UNKNOWN_TOOL
+            for tally in self._tallies(owner, role):
+                tally.tools[name] += 1
+        model_owner = owners.pop() if scope is None and len(owners) == 1 else scope
+        for tally in self._tallies(model_owner, role):
+            tally.model_calls += 1
+            tally.cost_usd += cost_usd or 0.0
+
+    def snapshot(self, stop_reason: str | None = None) -> RunStats:
+        """Return the counts so far."""
+        tallies = list(self._owners.values())
+        return RunStats(
+            run=self._owners.get(None, _Tally()).counts(),
+            workflows={
+                owner: tally.counts()
+                for owner, tally in sorted(
+                    (item for item in self._owners.items() if item[0] is not None),
+                    key=lambda item: str(item[0]),
+                )
+                if owner is not None
+            },
+            roles={role: self._roles[role].counts() for role in sorted(self._roles)},
+            total=StatsTotal(
+                model_calls=sum(tally.model_calls for tally in tallies),
+                tool_calls=sum(sum(tally.tools.values()) for tally in tallies),
+                cost_usd=sum(tally.cost_usd for tally in tallies),
+                limit=self.limit,
+            ),
+            spec_hash=self.spec_hash,
+            stop_reason=stop_reason,
+            finished_at=datetime.now(UTC).isoformat() if stop_reason else None,
+        )
+
+
+def append_history(stats: RunStats, path: Path) -> None:
+    """Add a finished run to the history."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("a", encoding="utf-8") as handle:
+        handle.write(stats.model_dump_json(exclude_none=True) + "\n")
+
+
+def read_history(path: Path) -> list[RunStats]:
+    """Read past runs. Bad lines are skipped, since they must never stop a run from finishing."""
+    if not path.exists():
+        return []
+    runs = []
+    for line in path.read_text(encoding="utf-8").splitlines():
+        try:
+            runs.append(RunStats.model_validate_json(line))
+        except ValidationError:
+            continue
+    return runs
+
+
+def _one_and_a_half_times(value: int) -> int:
+    """Return 1.5 times a value, rounded up to the next 10."""
+    return math.ceil(value * 15 / 100) * 10
+
+
+def suggest_max_steps(history: Iterable[RunStats], limit: int) -> str | None:
+    """Suggest a `--max-steps` value from past runs. This only ever advises; it changes nothing."""
+    usable = [run for run in history if run.stop_reason in USABLE_STOP_REASONS][-RECENT_RUNS:]
+    if len(usable) < MIN_RUNS_FOR_SUGGESTION:
+        return None
+    # A run that hit a lower, older limit says nothing against the current limit.
+    hits = [
+        run for run in usable if run.stop_reason == MAX_STEPS_REASON and run.total.limit >= limit
+    ]
+    if hits:
+        return (
+            f"{len(hits)} of your last {len(usable)} runs hit the step limit of {limit}; "
+            f"consider --max-steps {_one_and_a_half_times(limit)}"
+        )
+    largest = max(run.total.model_calls for run in usable)
+    target = max(_one_and_a_half_times(largest), MIN_SUGGESTED_LIMIT)
+    if target == limit:
+        return None
+    steps = "step" if largest == 1 else "steps"
+    return (
+        f"the largest of your last {len(usable)} runs used {largest} {steps}; "
+        f"consider --max-steps {target}"
+    )

workflowval/loop/trace.py

@@ -0,0 +1,29 @@
+"""`.workflowval/trace.jsonl`: every model and tool call (NFR-6), with secrets removed (NFR-7)."""
+
+from __future__ import annotations
+
+import json
+import sys
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from workflowval.redact import redact
+
+
+class Tracer:
+    """Appends one JSON object per event, flushed at once so a killed run keeps its trace."""
+
+    def __init__(self, path: Path, *, echo: bool = False) -> None:
+        self.path = path
+        self.echo = echo
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+    def event(self, kind: str, **fields: Any) -> None:
+        """Record one event."""
+        record = {"ts": datetime.now(UTC).isoformat(), "kind": kind, **redact(fields)}
+        line = json.dumps(record, default=str)
+        with self.path.open("a", encoding="utf-8") as handle:
+            handle.write(line + "\n")
+        if self.echo:
+            print(line, file=sys.stderr)

workflowval/models/__init__.py

@@ -0,0 +1,19 @@
+"""Provider adapters behind one interface. LiteLLM is imported only when a call is made."""
+
+from workflowval.models.base import Message, Model, ModelError, Reply, ToolCall, ToolSpec, Usage
+from workflowval.models.fake import FakeModel, FakeScriptExhausted, call
+from workflowval.models.litellm_adapter import LiteLLMModel
+
+__all__ = [
+    "FakeModel",
+    "FakeScriptExhausted",
+    "LiteLLMModel",
+    "Message",
+    "Model",
+    "ModelError",
+    "Reply",
+    "ToolCall",
+    "ToolSpec",
+    "Usage",
+    "call",
+]

workflowval/models/base.py

@@ -0,0 +1,68 @@
+"""The one interface the rest of workflowval uses to talk to a model."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal, Protocol
+
+
+class ModelError(Exception):
+    """A provider call failed. The message is already redacted."""
+
+
+@dataclass(frozen=True)
+class ToolCall:
+    """A tool call from the model. `parse_error` is set when its arguments weren't valid JSON."""
+
+    id: str
+    name: str
+    arguments: dict[str, Any] = field(default_factory=dict)
+    parse_error: str | None = None
+
+
+@dataclass(frozen=True)
+class Message:
+    """One conversation message."""
+
+    role: Literal["system", "user", "assistant", "tool"]
+    content: str = ""
+    tool_calls: tuple[ToolCall, ...] = ()
+    tool_call_id: str | None = None
+
+
+@dataclass(frozen=True)
+class ToolSpec:
+    """What the model sees of a tool."""
+
+    name: str
+    description: str
+    parameters: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class Usage:
+    """Tokens and cost of one call. `cost_usd` is None when the model has no known price."""
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cost_usd: float | None = None
+
+
+@dataclass(frozen=True)
+class Reply:
+    """A model's answer."""
+
+    text: str = ""
+    tool_calls: tuple[ToolCall, ...] = ()
+    usage: Usage = field(default_factory=Usage)
+
+
+class Model(Protocol):
+    """Anything that can answer a conversation."""
+
+    name: str
+
+    def complete(self, messages: Sequence[Message], tools: Sequence[ToolSpec]) -> Reply:
+        """Return the model's reply. Raises ModelError if the provider call fails."""
+        ...

workflowval/models/fake.py

@@ -0,0 +1,55 @@
+"""A scripted model, which makes the loop fully deterministic in tests and CI (NFR-8)."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from workflowval.models.base import Message, Reply, ToolCall, ToolSpec, Usage
+
+
+class FakeScriptExhausted(RuntimeError):
+    """The code under test asked for more replies than were scripted."""
+
+
+def call(name: str, /, **arguments: Any) -> Reply:
+    """Build a reply holding one tool call."""
+    return Reply(tool_calls=(ToolCall(id=f"call_{name}", name=name, arguments=arguments),))
+
+
+@dataclass
+class FakeModel:
+    """Replays replies in order and records every request it received."""
+
+    script: Iterable[Reply] = ()
+    name: str = "fake/scripted"
+    requests: list[tuple[list[Message], list[ToolSpec]]] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        self._replies = iter(self.script)
+
+    def complete(self, messages: Sequence[Message], tools: Sequence[ToolSpec]) -> Reply:
+        """Return the next scripted reply."""
+        self.requests.append((list(messages), list(tools)))
+        try:
+            return next(self._replies)
+        except StopIteration:
+            raise FakeScriptExhausted(
+                f"{self.name}: script ran out after {len(self.requests) - 1} replies"
+            ) from None
+
+    @classmethod
+    def from_file(cls, path: Path) -> FakeModel:
+        """Load a script: a JSON list of {text, tool_calls: [{id, name, arguments}], usage}."""
+        replies = [
+            Reply(
+                text=item.get("text", ""),
+                tool_calls=tuple(ToolCall(**tc) for tc in item.get("tool_calls", [])),
+                usage=Usage(**item.get("usage", {})),
+            )
+            for item in json.loads(path.read_text(encoding="utf-8"))
+        ]
+        return cls(script=replies)

workflowval/models/litellm_adapter.py

@@ -0,0 +1,153 @@
+"""The default adapter. This is the only module that imports a provider library (NFR-5)."""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+from workflowval.models.base import Message, ModelError, Reply, ToolCall, ToolSpec, Usage
+from workflowval.redact import redact_text
+
+
+def _load_litellm() -> Any:
+    # Without this, importing litellm fetches its price list from the network.
+    os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+    import litellm
+
+    return litellm
+
+
+# Seconds to wait before each new try. A temporary provider error shouldn't end a whole run.
+RETRY_DELAYS = (2.0, 4.0, 8.0, 16.0, 32.0)
+_TEMPORARY_ERRORS = frozenset(
+    {
+        "APIConnectionError",
+        "APITimeoutError",
+        "InternalServerError",
+        "RateLimitError",
+        "ServiceUnavailableError",
+        "Timeout",
+    }
+)
+
+
+def _is_temporary(exc: Exception) -> bool:
+    """True for rate limits, provider outages and dropped connections, which may clear."""
+    status = getattr(exc, "status_code", None)
+    if isinstance(status, int):
+        return status in (408, 429) or status >= 500
+    return type(exc).__name__ in _TEMPORARY_ERRORS
+
+
+def _to_wire(message: Message) -> dict[str, Any]:
+    wire: dict[str, Any] = {"role": message.role, "content": message.content}
+    if message.tool_calls:
+        wire["tool_calls"] = [
+            {
+                "id": call.id,
+                "type": "function",
+                "function": {"name": call.name, "arguments": json.dumps(call.arguments)},
+            }
+            for call in message.tool_calls
+        ]
+    if message.tool_call_id is not None:
+        wire["tool_call_id"] = message.tool_call_id
+    return wire
+
+
+def _from_wire(raw: Any) -> ToolCall:
+    name = raw.function.name
+    try:
+        parsed = json.loads(raw.function.arguments or "{}")
+    except json.JSONDecodeError as exc:
+        return ToolCall(id=raw.id, name=name, parse_error=f"arguments are not valid JSON: {exc}")
+    if not isinstance(parsed, dict):
+        return ToolCall(id=raw.id, name=name, parse_error="arguments must be a JSON object")
+    return ToolCall(id=raw.id, name=name, arguments=parsed)
+
+
+@dataclass
+class LiteLLMModel:
+    """Any provider LiteLLM supports, named `provider/model`. Keys come from env vars."""
+
+    name: str
+    temperature: float | None = None
+    seed: int | None = None
+    extra: dict[str, Any] = field(default_factory=dict)
+    track_cost: bool = True
+    sleep: Callable[[float], None] = time.sleep
+
+    def complete(self, messages: Sequence[Message], tools: Sequence[ToolSpec]) -> Reply:
+        """Call the provider and translate its reply. Failures raise a redacted ModelError."""
+        litellm = _load_litellm()
+        params: dict[str, Any] = {
+            "model": self.name,
+            "messages": [_to_wire(m) for m in messages],
+            # Drop params a provider doesn't support (e.g. seed) instead of failing the call.
+            "drop_params": True,
+            **self.extra,
+        }
+        if tools:
+            params["tools"] = [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": t.name,
+                        "description": t.description,
+                        "parameters": t.parameters,
+                    },
+                }
+                for t in tools
+            ]
+        if self.temperature is not None:
+            params["temperature"] = self.temperature
+        if self.seed is not None:
+            params["seed"] = self.seed
+        response = self._call_with_retries(litellm, params)
+        message = response.choices[0].message
+        usage = getattr(response, "usage", None)
+        input_tokens = getattr(usage, "prompt_tokens", 0) or 0
+        output_tokens = getattr(usage, "completion_tokens", 0) or 0
+        return Reply(
+            text=message.content or "",
+            tool_calls=tuple(_from_wire(c) for c in message.tool_calls or ()),
+            usage=Usage(
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                cost_usd=self._cost(litellm, response, input_tokens + output_tokens),
+            ),
+        )
+
+    def _call_with_retries(self, litellm: Any, params: dict[str, Any]) -> Any:
+        delays = iter(RETRY_DELAYS)
+        while True:
+            try:
+                return litellm.completion(**params)
+            except Exception as exc:
+                delay = next(delays, None) if _is_temporary(exc) else None
+                if delay is None:
+                    # "from None": a chained provider exception could carry request headers.
+                    detail = redact_text(str(exc))
+                    raise ModelError(f"{self.name}: {type(exc).__name__}: {detail}") from None
+            self.sleep(delay)
+
+    def _cost(self, litellm: Any, response: Any, tokens: int) -> float | None:
+        """Prefer the price the provider reported, since LiteLLM's estimate can be wrongly 0."""
+        if not self.track_cost:
+            return None
+        hidden = getattr(response, "_hidden_params", None) or {}
+        reported = hidden.get("response_cost")
+        if reported is None:
+            reported = getattr(getattr(response, "usage", None), "cost", None)
+        if reported is None:
+            try:
+                reported = litellm.completion_cost(completion_response=response)
+            except Exception:
+                return None  # An unpriced model (e.g. a brand-new or local one) is not an error.
+        cost = float(reported)
+        # A call that used tokens wasn't free: the price is unknown, so don't report it as $0.
+        return None if cost == 0.0 and tokens else cost