casebench 0.1.0__py3-none-any.whl

casebench/__init__.py

@@ -0,0 +1,5 @@
+"""Measure the legal accuracy of an AI system against court decisions and lawyers' references."""
+
+from importlib.metadata import version
+
+__version__ = version("casebench")

casebench/capture/__init__.py

@@ -0,0 +1 @@
+"""Capture: the step model and the parsers that read backend traffic."""

casebench/capture/assemble.py

@@ -0,0 +1,92 @@
+"""Turn one case's exchanges into steps, pair tool calls and show unaccounted time."""
+
+from collections.abc import Callable
+
+from casebench.capture.read import ids_in
+from casebench.capture.step import HttpExchange, Step, StepKind, ToolCall
+
+ParseExchange = Callable[[HttpExchange], list[Step]]
+
+
+def paired_tools(steps: list[Step]) -> list[Step]:
+    """A tool call asked for in one model response and answered in a later request is a tool
+    step: its name and arguments, the result's size, and the time between the two calls."""
+    pending: dict[str, tuple[ToolCall, float | None]] = {}
+    tools: list[Step] = []
+    for step in sorted((s for s in steps if s.kind is StepKind.LLM), key=_start):
+        for result in step.tool_results:
+            asked = pending.pop(result.id, None)
+            if asked is None:
+                continue
+            call, asked_at = asked
+            duration = _ms(asked_at, step.started)
+            tools.append(
+                Step(
+                    kind=StepKind.TOOL,
+                    name=call.name,
+                    arguments=call.arguments,
+                    result_size=result.size,
+                    started=asked_at,
+                    ended=step.started,
+                    duration_ms=duration,
+                )
+            )
+        for call in step.requested_tools:
+            pending[call.id] = (call, step.ended)
+    return tools
+
+
+def gaps(steps: list[Step]) -> list[Step]:
+    """Time between steps that no step covers, named after the steps on either side.
+    Gaps are reported to the millisecond, the resolution durations are shown in."""
+    timed = sorted((s for s in steps if s.started is not None and s.ended is not None), key=_start)
+    found: list[Step] = []
+    covered_until: float | None = None
+    before: Step | None = None
+    for step in timed:
+        assert step.started is not None and step.ended is not None
+        if covered_until is not None and before is not None:
+            missing = round((step.started - covered_until) * 1000)
+            if missing > 0:
+                found.append(
+                    Step(
+                        kind=StepKind.GAP,
+                        name=f"between {_label(before)} and {_label(step)}",
+                        started=covered_until,
+                        ended=step.started,
+                        duration_ms=float(missing),
+                    )
+                )
+        if covered_until is None or step.ended > covered_until:
+            covered_until, before = step.ended, step
+    return found
+
+
+def assemble(
+    exchanges: list[HttpExchange], parse: ParseExchange, in_process: list[Step]
+) -> list[Step]:
+    parsed = [(exchange, parse(exchange)) for exchange in exchanges]
+    known = {i for _, found in parsed for s in found for i in s.doc_ids_out}
+    known |= {i for s in in_process for i in s.doc_ids_out}
+    steps = [
+        step.model_copy(update={"doc_ids_in": ids_in(exchange.request_body, known)})
+        if step.kind is StepKind.LLM
+        else step
+        for exchange, found in parsed
+        for step in found
+    ] + list(in_process)
+    steps += paired_tools(steps)
+    steps += gaps(steps)
+    return sorted(steps, key=_start)
+
+
+def _start(step: Step) -> float:
+    return step.started if step.started is not None else float("inf")
+
+
+def _ms(start: float | None, end: float | None) -> float | None:
+    return (end - start) * 1000 if start is not None and end is not None else None
+
+
+def _label(step: Step) -> str:
+    return f"{step.kind.value} {step.name or step.model or ''}".rstrip()

casebench/capture/backend.py

@@ -0,0 +1,47 @@
+"""Read a backend route's document ids by the JSONPath its config names, or from the path."""
+
+import re
+from typing import Any
+from urllib.parse import unquote, urlsplit
+
+from jsonpath_ng.ext.parser import parse as jsonpath
+
+from casebench.capture.proxy.routes import Route
+from casebench.capture.step import NO_DIAGNOSTIC, HttpExchange, Step, StepKind
+
+
+def _ids(expression: str | None, document: Any) -> list[str]:
+    if expression is None or document is None:
+        return []
+    return [str(match.value) for match in jsonpath(expression).find(document)]
+
+
+def fetched(route: Route, url: str) -> str | None:
+    """The id a request reads, when its path ends as the route's `fetch_path` does."""
+    if route.fetch_path is None or "{id}" not in route.fetch_path:
+        return None
+    before, _, after = route.fetch_path.partition("{id}")
+    pattern = re.escape(before) + r"([^/]+)" + re.escape(after) + "$"
+    found = re.search(pattern, urlsplit(url).path)
+    return unquote(found.group(1)) if found else None
+
+
+def backend_step(exchange: HttpExchange, route: Route) -> Step:
+    """One retrieval, rerank or document fetch; without ids found, it says so."""
+    document_id = fetched(route, exchange.url)
+    if document_id is not None:
+        kind, ids_in, ids_out = StepKind.DOCUMENT, [], [document_id]
+    else:
+        kind = StepKind(route.step)
+        ids_out = _ids(route.doc_ids, exchange.response_body)
+        ids_in = _ids(route.doc_ids_in, exchange.request_body)
+    return Step(
+        kind=kind,
+        name=route.name,
+        doc_ids_in=ids_in,
+        doc_ids_out=ids_out,
+        started=exchange.started,
+        ended=exchange.ended,
+        duration_ms=exchange.duration_ms,
+        note=None if ids_out else NO_DIAGNOSTIC,
+    )

casebench/capture/cli.py

@@ -0,0 +1,36 @@
+"""`casebench proxy env`: the variables that point an already-running agent at the proxy."""
+
+from pathlib import Path
+
+import click
+
+from casebench.capture.config import routes_for
+from casebench.capture.proxy.routes import environment
+from casebench.capture.session import TRACES_ENDPOINT, traces_endpoint
+from casebench.config import DEFAULT_CONFIG, load_config
+
+
+@click.group()
+def proxy() -> None:
+    """The recording proxy that captures model and backend calls."""
+
+
+@proxy.command()
+@click.option(
+    "--config",
+    "config_path",
+    type=click.Path(dir_okay=False, path_type=Path),
+    default=DEFAULT_CONFIG,
+    show_default=True,
+)
+def env(config_path: Path) -> None:
+    """Print the lines to set in the agent's environment."""
+    config = load_config(config_path)
+    if config.capture is None or config.capture.port == 0:
+        raise click.ClickException(
+            "set capture.port in the config so an agent that is already running can reach "
+            "the proxy; when Casebench starts the agent it sets these itself"
+        )
+    for name, value in environment(routes_for(config.capture), config.capture.port).items():
+        click.echo(f"{name}={value}")
+    click.echo(f"{TRACES_ENDPOINT}={traces_endpoint(config.capture.port)}")

casebench/capture/config.py

@@ -0,0 +1,28 @@
+"""The `capture:` section of casebench.yaml: the proxy's port and its routes."""
+
+from typing import Any
+
+from pydantic import Field
+
+from casebench.capture.functions import StepFunction
+from casebench.capture.proxy.routes import Route, default_routes
+from casebench.core.model import Record
+
+
+class CaptureConfig(Record):
+    """`port: 0` lets the operating system choose; name a port for an agent already running.
+    `traces: true` points the OpenTelemetry exporter of agents Casebench starts at it."""
+
+    port: int = 0
+    routes: list[dict[str, Any]] = Field(default_factory=list)
+    functions: list[StepFunction] = Field(default_factory=list)
+    traces: bool = False
+
+
+def routes_for(capture: CaptureConfig | None) -> list[Route]:
+    """The parsers' default routes, with the config's fields laid over a route of the same
+    name, and the config's new routes added."""
+    merged = {route.name: route.model_dump() for route in default_routes()}
+    for route in capture.routes if capture else []:
+        merged[route["name"]] = {**merged.get(route["name"], {}), **route}
+    return [Route.model_validate(fields) for fields in merged.values()]

casebench/capture/conformance.py

@@ -0,0 +1,33 @@
+"""Conformance for capture parsers, measured as the share of recorded calls parsed.
+
+A sample is `{"exchange": <HTTP exchange>, "expected": [<step>, ...]}`.
+"""
+
+from pathlib import Path
+
+from casebench.capture.interface import CaptureParser
+from casebench.capture.step import HttpExchange, Step
+from casebench.plugins.samples import SuiteResult, ratio, read_json, sample_files
+
+
+def check(parser: CaptureParser, samples: Path) -> SuiteResult:
+    result = SuiteResult()
+    parsed = 0
+    if not parser.hosts:
+        result.problems.append("hosts is empty, so no traffic would reach this parser")
+    for path in sample_files(samples):
+        result.samples += 1
+        sample = read_json(path)
+        exchange = HttpExchange.model_validate(sample["exchange"])
+        expected = [Step.model_validate(step) for step in sample["expected"]]
+        try:
+            steps = parser.parse(exchange)
+        except Exception as error:
+            result.problems.append(f"{path.name}: {type(error).__name__}: {error}")
+            continue
+        if not all(isinstance(step, Step) for step in steps):
+            result.problems.append(f"{path.name}: parse() must return Step records")
+            continue
+        parsed += steps == expected
+    result.figure = ratio(parsed, result.samples)
+    return result

casebench/capture/context.py

@@ -0,0 +1,53 @@
+"""Tie every captured step to the case that caused it."""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+
+from casebench.capture.step import Step
+
+
+@dataclass
+class CaseCapture:
+    case_id: str
+    traceparent: str | None = None
+    worker: str | None = None
+    steps: list[Step] = field(default_factory=list)
+
+
+_current: ContextVar[CaseCapture | None] = ContextVar("casebench_case", default=None)
+
+
+@contextmanager
+def capturing(case_id: str, traceparent: str | None = None) -> Iterator[CaseCapture]:
+    """Steps recorded inside this block belong to this case, even with other cases
+    running at the same time, because each asyncio task and thread gets its own copy."""
+    capture = CaseCapture(case_id, traceparent)
+    token = _current.set(capture)
+    try:
+        yield capture
+    finally:
+        _current.reset(token)
+
+
+def record(step: Step) -> bool:
+    """Attach a step to the running case; False when no case is running."""
+    capture = _current.get()
+    if capture is None:
+        return False
+    capture.steps.append(step)
+    return True
+
+
+def current_traceparent() -> str | None:
+    """The running case's trace context, to pass on to an agent outside this process."""
+    capture = _current.get()
+    return capture.traceparent if capture is not None else None
+
+
+def set_worker(worker: str) -> None:
+    """Note which worker process answers the running case, so its calls are found by path."""
+    capture = _current.get()
+    if capture is not None:
+        capture.worker = worker

casebench/capture/corpus_routes.py

@@ -0,0 +1,50 @@
+"""Capture routes for the corpora the config names, so an agent's own searches are recorded.
+
+An agent that searches the same document store Casebench checks citations against calls it
+at the corpus's address. A Python agent's calls arrive at the proxy by where they go
+(decision 0044), so a route at that host reads them: a search as a retrieval with the ids
+the corpus's search results hold, a fetch as the document its path names.
+"""
+
+from typing import Any
+from urllib.parse import urlsplit
+
+from casebench.capture.proxy.routes import Route
+from casebench.corpus.binding import CorpusConfig
+from casebench.targets.placeholders import expand_env
+
+
+def ids_path(options: dict[str, Any]) -> str | None:
+    """Where a corpus's search response holds its ids, as one JSONPath."""
+    results, fields = options.get("search_results"), options.get("result_fields") or {}
+    if results and fields.get("id"):
+        return f"{results}{str(fields['id']).removeprefix('$')}"
+    return options.get("search_ids")
+
+
+def corpus_routes(corpora: dict[str, CorpusConfig], taken: list[Route]) -> list[Route]:
+    """A route at each HTTP corpus's host that no route covers yet. A base URL that names a
+    variable is read as the corpus reads it; one that cannot be read here gets no route."""
+    covered = {urlsplit(route.upstream).netloc for route in taken}
+    made = []
+    for name, corpus in corpora.items():
+        if corpus.connector != "http" or not corpus.options.get("base_url"):
+            continue
+        try:
+            base = expand_env({"base_url": str(corpus.options["base_url"])})["base_url"]
+        except KeyError:
+            continue
+        parts = urlsplit(base)
+        if parts.scheme not in ("http", "https") or not parts.netloc or parts.netloc in covered:
+            continue
+        covered.add(parts.netloc)
+        made.append(
+            Route(
+                name=f"corpus-{name}",
+                upstream=f"{parts.scheme}://{parts.netloc}",
+                step="retrieval",
+                doc_ids=ids_path(corpus.options),
+                fetch_path=corpus.options.get("fetch_path"),
+            )
+        )
+    return made

casebench/capture/functions.py

@@ -0,0 +1,115 @@
+"""Wrap functions the config names so each call records a step with its document ids."""
+
+import dataclasses
+import functools
+import inspect
+import json
+import time
+from collections.abc import Callable
+from typing import Any
+
+from jsonpath_ng.ext.parser import parse as jsonpath
+
+from casebench.capture.context import record
+from casebench.capture.step import NO_DIAGNOSTIC, Step, StepKind
+from casebench.core.model import Record
+from casebench.loading import import_location, resolve, split_entry
+
+# Set by the capture session for worker processes (decision 0022).
+STEP_FUNCTIONS = "CASEBENCH_STEP_FUNCTIONS"
+
+
+class StepFunction(Record):
+    """`capture.functions` in casebench.yaml."""
+
+    function: str
+    step: StepKind = StepKind.RERANK
+    name: str | None = None
+    doc_ids_in: str | None = None
+    doc_ids_out: str | None = None
+
+
+def _plain(value: Any) -> Any:
+    """Dataclasses, pydantic models and plain objects as their fields, for JSONPath."""
+    if dataclasses.is_dataclass(value) and not isinstance(value, type):
+        return {k: _plain(v) for k, v in dataclasses.asdict(value).items()}
+    dump = getattr(value, "model_dump", None)
+    if callable(dump) and not isinstance(value, type):
+        return _plain(dump())
+    if isinstance(value, dict):
+        return {str(k): _plain(v) for k, v in value.items()}
+    if isinstance(value, list | tuple | set):
+        return [_plain(v) for v in value]
+    if hasattr(value, "__dict__") and not callable(value):
+        return {k: _plain(v) for k, v in vars(value).items() if not k.startswith("_")}
+    return value
+
+
+def _ids(expression: str | None, document: Any) -> list[str]:
+    if expression is None:
+        return []
+    return [str(match.value) for match in jsonpath(expression).find(document)]
+
+
+def _step(
+    spec: StepFunction, arguments: dict[str, Any], result: Any, started: float, begun: float
+) -> Step:
+    """`started` orders the step by the wall clock; `begun` times it by the monotonic one."""
+    ended = time.time()
+    document = {"arguments": _plain(arguments), "result": _plain(result)}
+    ids_out = _ids(spec.doc_ids_out, document)
+    return Step(
+        kind=spec.step,
+        name=spec.name or spec.function.rsplit(":", 1)[-1],
+        doc_ids_in=_ids(spec.doc_ids_in, document),
+        doc_ids_out=ids_out,
+        started=started,
+        ended=ended,
+        duration_ms=(time.perf_counter() - begun) * 1000,
+        note=None if ids_out else NO_DIAGNOSTIC,
+    )
+
+
+def wrap(spec: StepFunction, original: Callable[..., Any]) -> Callable[..., Any]:
+    signature = inspect.signature(original)
+
+    def bound(args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+        try:
+            return dict(signature.bind(*args, **kwargs).arguments)
+        except TypeError:
+            return {"args": list(args), **kwargs}
+
+    if inspect.iscoroutinefunction(original):
+
+        @functools.wraps(original)
+        async def wrapped_async(*args: Any, **kwargs: Any) -> Any:
+            started, begun = time.time(), time.perf_counter()
+            result = await original(*args, **kwargs)
+            record(_step(spec, bound(args, kwargs), result, started, begun))
+            return result
+
+        return wrapped_async
+
+    @functools.wraps(original)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+        started, begun = time.time(), time.perf_counter()
+        result = original(*args, **kwargs)
+        record(_step(spec, bound(args, kwargs), result, started, begun))
+        return result
+
+    return wrapped
+
+
+def install(specs: list[StepFunction]) -> None:
+    """Replace each named function on the object that holds it; a missing name stops here."""
+    for spec in specs:
+        location, attribute, is_file = split_entry(spec.function)
+        owner, name = resolve(import_location(location, is_file), attribute)
+        original = getattr(owner, name, None)
+        if not callable(original):
+            raise LookupError(f"capture.functions names {spec.function}, which is not a function")
+        setattr(owner, name, wrap(spec, original))
+
+
+def specs_from_environment(value: str | None) -> list[StepFunction]:
+    return [StepFunction.model_validate(item) for item in json.loads(value)] if value else []

casebench/capture/interface.py

@@ -0,0 +1,23 @@
+"""The capture parser contract: how to read one backend's traffic."""
+
+from typing import Protocol, runtime_checkable
+
+from casebench.capture.step import HttpExchange, Step
+
+
+@runtime_checkable
+class CaptureParser(Protocol):
+    """Turns one backend's HTTP exchanges into normalised steps.
+
+    A parser may also list `packages`: the SDKs whose calls it reads, as their package
+    registries spell them, so `casebench init` can name the SDKs it finds in a project.
+    """
+
+    @property
+    def hosts(self) -> list[str]:
+        """The hosts whose traffic this parser reads, such as `api.example.com`."""
+        ...
+
+    def parse(self, exchange: HttpExchange) -> list[Step]:
+        """The steps one exchange shows; an empty list if it shows none."""
+        ...

casebench/capture/otlp/__init__.py

@@ -0,0 +1 @@
+"""Receive OpenTelemetry spans a team already emits and turn them into steps."""

casebench/capture/otlp/decode.py

@@ -0,0 +1,82 @@
+"""Read an OTLP trace export, protobuf or JSON, into plain spans."""
+
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
+from opentelemetry.proto.common.v1.common_pb2 import AnyValue
+
+NANOS = 1_000_000_000
+
+
+@dataclass
+class RawSpan:
+    trace_id: str
+    name: str
+    started: float
+    ended: float
+    attributes: dict[str, Any] = field(default_factory=dict)
+
+
+def _proto_value(value: AnyValue) -> Any:
+    kind = value.WhichOneof("value")
+    if kind == "array_value":
+        return [_proto_value(item) for item in value.array_value.values]
+    if kind == "kvlist_value":
+        return {item.key: _proto_value(item.value) for item in value.kvlist_value.values}
+    if kind == "bytes_value":
+        return value.bytes_value.hex()
+    return getattr(value, kind) if kind else None
+
+
+def from_protobuf(body: bytes) -> list[RawSpan]:
+    request = ExportTraceServiceRequest.FromString(body)
+    return [
+        RawSpan(
+            trace_id=span.trace_id.hex(),
+            name=span.name,
+            started=span.start_time_unix_nano / NANOS,
+            ended=span.end_time_unix_nano / NANOS,
+            attributes={a.key: _proto_value(a.value) for a in span.attributes},
+        )
+        for resource in request.resource_spans
+        for scope in resource.scope_spans
+        for span in scope.spans
+    ]
+
+
+def _json_value(value: dict[str, Any]) -> Any:
+    if "arrayValue" in value:
+        return [_json_value(item) for item in value["arrayValue"].get("values", [])]
+    if "kvlistValue" in value:
+        return {i["key"]: _json_value(i["value"]) for i in value["kvlistValue"].get("values", [])}
+    if "intValue" in value:
+        return int(value["intValue"])  # OTLP/JSON writes 64-bit integers as strings
+    for key in ("stringValue", "doubleValue", "boolValue", "bytesValue"):
+        if key in value:
+            return value[key]
+    return None
+
+
+def from_json(body: bytes) -> list[RawSpan]:
+    """OTLP/JSON: ids are hex strings, times are nanoseconds written as strings."""
+    request = json.loads(body)
+    return [
+        RawSpan(
+            trace_id=span.get("traceId", "").lower(),
+            name=span.get("name", ""),
+            started=int(span.get("startTimeUnixNano", 0)) / NANOS,
+            ended=int(span.get("endTimeUnixNano", 0)) / NANOS,
+            attributes={
+                a["key"]: _json_value(a.get("value", {})) for a in span.get("attributes", [])
+            },
+        )
+        for resource in request.get("resourceSpans", [])
+        for scope in resource.get("scopeSpans", [])
+        for span in scope.get("spans", [])
+    ]
+
+
+def decode(body: bytes, content_type: str) -> list[RawSpan]:
+    return from_json(body) if "json" in content_type else from_protobuf(body)