spanlight 0.1.0__py3-none-any.whl

spanlight/_context.py

@@ -0,0 +1,79 @@
+"""Propagation of the active span through :mod:`contextvars`.
+
+Using a :class:`~contextvars.ContextVar` means the current span follows the
+natural flow of the program: nested function calls see their caller's span,
+and :mod:`asyncio` tasks inherit the span that was active when they were
+created (each task runs in a copy of its creator's context). Threads do not
+inherit context automatically; use :func:`contextvars.copy_context` and
+``ctx.run(...)`` when handing work to a thread pool.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Iterator
+from contextvars import ContextVar, Token
+
+from ._span import Span
+
+_current_span: ContextVar[Span | None] = ContextVar("spanlight_current_span", default=None)
+_current_prompt_version_id: ContextVar[str | None] = ContextVar(
+    "spanlight_current_prompt_version_id", default=None
+)
+
+
+def get_current_span() -> Span | None:
+    """Return the span active in the current context, if any."""
+    return _current_span.get()
+
+
+def get_current_trace_id() -> str | None:
+    """Return the trace id of the active span, if any."""
+    span = _current_span.get()
+    return span.trace_id if span is not None else None
+
+
+@contextlib.contextmanager
+def activate(span: Span) -> Iterator[Span]:
+    """Make ``span`` the current span for the duration of the ``with`` block."""
+    token = _current_span.set(span)
+    try:
+        yield span
+    finally:
+        _current_span.reset(token)
+
+
+def set_current_span(span: Span | None) -> Token[Span | None]:
+    """Set the current span and return a token for :func:`restore_current_span`."""
+    return _current_span.set(span)
+
+
+def restore_current_span(token: Token[Span | None], previous: Span | None) -> None:
+    """Undo :func:`set_current_span`.
+
+    ``ContextVar.reset`` refuses tokens created in a different context (for
+    example a ``with`` block entered in one task and exited in another). In
+    that case the previous span is restored by value instead of raising.
+    """
+    try:
+        _current_span.reset(token)
+    except ValueError:
+        _current_span.set(previous)
+
+
+def get_current_prompt_version_id() -> str | None:
+    """Return the prompt version id spans started now are linked to, if any."""
+    return _current_prompt_version_id.get()
+
+
+def set_current_prompt_version_id(version_id: str | None) -> Token[str | None]:
+    """Link spans started from now on to ``version_id``; returns a token to undo it."""
+    return _current_prompt_version_id.set(version_id)
+
+
+def restore_current_prompt_version_id(token: Token[str | None], previous: str | None) -> None:
+    """Undo :func:`set_current_prompt_version_id`, by value when the token is foreign."""
+    try:
+        _current_prompt_version_id.reset(token)
+    except ValueError:
+        _current_prompt_version_id.set(previous)

spanlight/_decorator.py

@@ -0,0 +1,345 @@
+"""Implementation of the ``@observe`` decorator.
+
+The decorator supports four kinds of callables, each needing slightly
+different treatment so that the span covers the *real* duration of the work:
+
+* plain functions: the span covers the call;
+* coroutine functions: the span covers the awaited coroutine;
+* generator functions: the span stays open until the generator is exhausted,
+  closed or fails, and the current-span context is active only while the
+  generator body runs (never leaking into the consumer between items);
+* async generator functions: same as generators, for ``async for``.
+
+In every case user exceptions propagate unchanged, and failures inside the
+SDK itself are swallowed so they cannot break the decorated function.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from collections.abc import AsyncGenerator, Callable, Generator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast
+
+from ._context import get_current_span, restore_current_span, set_current_span
+from ._serialize import MAX_ITEMS
+from ._span import Span, SpanKind
+
+if TYPE_CHECKING:
+    from ._client import Spanlight
+
+logger = logging.getLogger("spanlight")
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+_IMPLICIT_RECEIVERS = frozenset({"self", "cls"})
+
+
+@dataclass(frozen=True)
+class _ObserveOptions:
+    name: str
+    kind: SpanKind
+    capture_input: bool
+    capture_output: bool
+
+
+def build_observe(
+    get_client: Callable[[], Spanlight],
+    name: Callable[P, R] | str | None,
+    *,
+    kind: SpanKind,
+    capture_input: bool,
+    capture_output: bool,
+) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
+    """Shared implementation behind both ``observe`` entry points.
+
+    ``name`` is either the span name (``@observe(name="answer")`` or
+    ``@observe("answer")``), ``None`` (``@observe()``), or, for the bare
+    ``@observe`` form, the decorated function itself.
+    """
+    span_name = name if isinstance(name, str) else None
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        options = _ObserveOptions(
+            name=span_name or str(getattr(func, "__name__", "function")),
+            kind=kind,
+            capture_input=capture_input,
+            capture_output=capture_output,
+        )
+        return _decorate(func, get_client, options)
+
+    if callable(name):
+        return decorator(name)
+    return decorator
+
+
+def _decorate(
+    func: Callable[P, R],
+    get_client: Callable[[], Spanlight],
+    options: _ObserveOptions,
+) -> Callable[P, R]:
+    signature = _safe_signature(func)
+
+    def start(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Span | None:
+        """Start the span for one call, or return ``None`` to run untraced."""
+        try:
+            client = get_client()
+            if not client.enabled:
+                return None
+            span = client.start_span(options.name, kind=options.kind)
+            if options.capture_input:
+                span.set_input(_bind_arguments(signature, args, kwargs))
+            return span
+        except Exception:
+            logger.debug("observe: failed to start span %r", options.name, exc_info=True)
+            return None
+
+    if inspect.isasyncgenfunction(func):
+        return cast("Callable[P, R]", _wrap_async_generator(func, start, options))
+    if inspect.isgeneratorfunction(func):
+        return cast("Callable[P, R]", _wrap_generator(func, start, options))
+    if inspect.iscoroutinefunction(func):
+        return cast("Callable[P, R]", _wrap_coroutine(func, start, options))
+    return _wrap_function(func, start, options)
+
+
+_Starter = Callable[[tuple[Any, ...], dict[str, Any]], Span | None]
+
+
+def _wrap_function(
+    func: Callable[P, R],
+    start: _Starter,
+    options: _ObserveOptions,
+) -> Callable[P, R]:
+    @functools.wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        span = start(args, kwargs)
+        if span is None:
+            return func(*args, **kwargs)
+
+        previous = get_current_span()
+        token = set_current_span(span)
+        try:
+            result = func(*args, **kwargs)
+        except BaseException as exc:
+            _finish(span, error=exc)
+            raise
+        finally:
+            restore_current_span(token, previous)
+        _finish(span, output=result, capture_output=options.capture_output)
+        return result
+
+    return wrapper
+
+
+def _wrap_coroutine(
+    func: Callable[..., Any],
+    start: _Starter,
+    options: _ObserveOptions,
+) -> Callable[..., Any]:
+    @functools.wraps(func)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        span = start(args, kwargs)
+        if span is None:
+            return await func(*args, **kwargs)
+
+        previous = get_current_span()
+        token = set_current_span(span)
+        try:
+            result = await func(*args, **kwargs)
+        except BaseException as exc:
+            _finish(span, error=exc)
+            raise
+        finally:
+            restore_current_span(token, previous)
+        _finish(span, output=result, capture_output=options.capture_output)
+        return result
+
+    return wrapper
+
+
+def _wrap_generator(
+    func: Callable[..., Generator[Any, Any, Any]],
+    start: _Starter,
+    options: _ObserveOptions,
+) -> Callable[..., Generator[Any, Any, Any]]:
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Generator[Any, Any, Any]:
+        span = start(args, kwargs)
+        generator = func(*args, **kwargs)
+        if span is None:
+            return generator
+        return _traced_generator(generator, span, options)
+
+    return wrapper
+
+
+def _traced_generator(
+    generator: Generator[Any, Any, Any],
+    span: Span,
+    options: _ObserveOptions,
+) -> Generator[Any, Any, Any]:
+    """Drive ``generator`` step by step with ``span`` active only during each step.
+
+    ``send()``, ``throw()`` and ``close()`` from the consumer are forwarded
+    to the inner generator, so the wrapper is transparent.
+    """
+    collected = _OutputCollector(enabled=options.capture_output)
+    step: Callable[[], Any] = functools.partial(generator.send, None)
+    while True:
+        previous = get_current_span()
+        token = set_current_span(span)
+        try:
+            item = step()
+        except StopIteration as stop:
+            _finish(
+                span, output=collected.result(stop.value), capture_output=options.capture_output
+            )
+            return stop.value
+        except BaseException as exc:
+            _finish(span, error=exc)
+            raise
+        finally:
+            restore_current_span(token, previous)
+
+        collected.add(item)
+        try:
+            sent = yield item
+        except GeneratorExit:
+            try:
+                generator.close()
+            finally:
+                _finish(span, output=collected.result(None), capture_output=options.capture_output)
+            raise
+        except BaseException as exc:
+            step = functools.partial(generator.throw, exc)
+        else:
+            step = functools.partial(generator.send, sent)
+
+
+def _wrap_async_generator(
+    func: Callable[..., AsyncGenerator[Any, Any]],
+    start: _Starter,
+    options: _ObserveOptions,
+) -> Callable[..., AsyncGenerator[Any, Any]]:
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> AsyncGenerator[Any, Any]:
+        span = start(args, kwargs)
+        generator = func(*args, **kwargs)
+        if span is None:
+            return generator
+        return _traced_async_generator(generator, span, options)
+
+    return wrapper
+
+
+async def _traced_async_generator(
+    generator: AsyncGenerator[Any, Any],
+    span: Span,
+    options: _ObserveOptions,
+) -> AsyncGenerator[Any, Any]:
+    """Async counterpart of :func:`_traced_generator`."""
+    collected = _OutputCollector(enabled=options.capture_output)
+    step: Callable[[], Any] = functools.partial(generator.asend, None)
+    while True:
+        previous = get_current_span()
+        token = set_current_span(span)
+        try:
+            item = await step()
+        except StopAsyncIteration:
+            _finish(span, output=collected.result(None), capture_output=options.capture_output)
+            return
+        except BaseException as exc:
+            _finish(span, error=exc)
+            raise
+        finally:
+            restore_current_span(token, previous)
+
+        collected.add(item)
+        try:
+            sent = yield item
+        except GeneratorExit:
+            try:
+                await generator.aclose()
+            finally:
+                _finish(span, output=collected.result(None), capture_output=options.capture_output)
+            raise
+        except BaseException as exc:
+            step = functools.partial(generator.athrow, exc)
+        else:
+            step = functools.partial(generator.asend, sent)
+
+
+class _OutputCollector:
+    """Accumulates yielded items (bounded) to use as a generator span's output."""
+
+    def __init__(self, *, enabled: bool) -> None:
+        self._enabled = enabled
+        self._items: list[Any] = []
+        self.truncated = False
+
+    def add(self, item: Any) -> None:
+        if not self._enabled:
+            return
+        if len(self._items) < MAX_ITEMS:
+            self._items.append(item)
+        else:
+            self.truncated = True
+
+    def result(self, return_value: Any) -> Any:
+        """Return the items, joined into one string when they are all strings."""
+        if not self._enabled:
+            return None
+        if self._items and all(isinstance(item, str) for item in self._items):
+            return "".join(self._items)
+        if not self._items and return_value is not None:
+            return return_value
+        return list(self._items)
+
+
+def _finish(
+    span: Span,
+    *,
+    output: Any = None,
+    capture_output: bool = False,
+    error: BaseException | None = None,
+) -> None:
+    """End ``span`` recording either its output or the error. Never raises."""
+    try:
+        if error is not None:
+            if not isinstance(error, GeneratorExit):
+                span.record_error(error)
+        elif capture_output:
+            span.set_output(output)
+        span.end()
+    except Exception:
+        logger.debug("observe: failed to finish span", exc_info=True)
+
+
+def _safe_signature(func: Callable[..., Any]) -> inspect.Signature | None:
+    try:
+        return inspect.signature(func)
+    except (TypeError, ValueError):
+        return None
+
+
+def _bind_arguments(
+    signature: inspect.Signature | None,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> dict[str, Any]:
+    """Map call arguments to parameter names, omitting ``self``/``cls``."""
+    if signature is None:
+        return {"args": list(args), "kwargs": kwargs}
+    try:
+        bound = signature.bind_partial(*args, **kwargs)
+    except TypeError:
+        return {"args": list(args), "kwargs": kwargs}
+    arguments = dict(bound.arguments)
+    parameters = list(signature.parameters)
+    if parameters and parameters[0] in _IMPLICIT_RECEIVERS:
+        arguments.pop(parameters[0], None)
+    return arguments

spanlight/_experiments.py

@@ -0,0 +1,256 @@
+"""Run a prompt over a dataset on the server and gate on the scores.
+
+Use it as ``python -m spanlight._experiments``.
+
+The runner talks to the dashboard API (``/api/v1``) with a personal access token. It finds the
+dataset, prompt version and evaluators by name, reuses or creates the experiment for that
+combination, starts a run and polls it until it ends. Exit codes: 0 every threshold held, 1 a
+score fell below its ``--fail-below`` threshold, 2 the run could not be started, failed or timed
+out. The ``eval-gate`` GitHub Action wraps this module.
+
+The token is read from ``SPANLIGHT_TOKEN`` and never printed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+from collections.abc import Callable, Mapping, Sequence
+from decimal import Decimal, InvalidOperation
+from typing import Any
+
+import httpx
+
+from ._experiments_api import (
+    Api,
+    ExperimentError,
+    ResolvedPrompt,
+    resolve_dataset,
+    resolve_evaluators,
+    resolve_prompt,
+)
+from ._experiments_parse import PromptRef, parse_fail_below, parse_prompt_ref
+from ._experiments_report import render_markdown
+from ._experiments_runs import (
+    ExperimentOutcome,
+    org_of,
+    outcome_of,
+    start_run,
+    wait,
+    wait_idle,
+)
+from ._experiments_setup import experiment_for, normalize_credential_id
+
+__all__ = [
+    "ExperimentError",
+    "ExperimentOutcome",
+    "PromptRef",
+    "main",
+    "parse_fail_below",
+    "parse_prompt_ref",
+    "run_experiment",
+]
+
+EXIT_PASSED = 0
+EXIT_THRESHOLD = 1
+EXIT_FAILED = 2
+_MAX_COST_CEILING = Decimal(1000)
+
+
+def _prepare(
+    api: Api,
+    dataset: str,
+    prompt: str,
+    evaluators: list[str],
+    credential_id: str,
+    model: str,
+    max_cost_usd: Decimal | None,
+) -> tuple[dict[str, Any], ResolvedPrompt]:
+    """Resolve the names to ids and find or create the experiment."""
+    dataset_row = resolve_dataset(api, dataset)
+    version = resolve_prompt(api, parse_prompt_ref(prompt))
+    experiment = experiment_for(
+        api,
+        dataset_row,
+        version,
+        evaluators,
+        resolve_evaluators(api, evaluators),
+        credential_id=credential_id,
+        model=model,
+        max_cost_usd=max_cost_usd,
+    )
+    return experiment, version
+
+
+def run_experiment(
+    *,
+    host: str,
+    token: str,
+    project_id: str,
+    dataset: str,
+    prompt: str,
+    evaluators: list[str],
+    fail_below: dict[str, Decimal],
+    credential_id: str,
+    model: str,
+    max_cost_usd: Decimal | None = None,
+    timeout_s: int = 1800,
+    poll_s: float = 5,
+    transport: httpx.BaseTransport | None = None,
+    environ: Mapping[str, str] | None = None,
+    sleep: Callable[[float], None] = time.sleep,
+    clock: Callable[[], float] = time.monotonic,
+) -> ExperimentOutcome:
+    """Run ``prompt`` (``name``, ``name@label`` or ``name:version``) over ``dataset`` and score it.
+
+    ``credential_id`` and ``model`` choose the provider the prompt runs on and, with the
+    evaluators, which experiment is used (a different combination is a different experiment).
+    ``max_cost_usd`` caps a run and applies only when the experiment is created. The whole call,
+    including waiting for an earlier run of the same experiment, is bounded by ``timeout_s``.
+    Raises :class:`ExperimentError` when the run cannot be set up; a run that fails or times out
+    is returned with that ``status``.
+    """
+    env = os.environ if environ is None else environ
+    credential_id = normalize_credential_id(credential_id)
+    api = Api(host=host, token=token, project_id=project_id, transport=transport)
+    try:
+        deadline = clock() + timeout_s
+        experiment, version = _prepare(
+            api, dataset, prompt, evaluators, credential_id, model, max_cost_usd
+        )
+        url = f"{host.rstrip('/')}/{org_of(api)}/{project_id}/experiments/{experiment['id']}"
+        wait_idle(api, experiment["id"], deadline, poll_s, sleep, clock)
+        run = start_run(api, experiment["id"], version.version_id, env)
+        try:
+            final = wait(api, experiment["id"], run["id"], deadline, poll_s, sleep, clock)
+        except ExperimentError as error:
+            # The run exists: report it with its id rather than lose it.
+            final = {"status": "failed", "error": str(error), "totals": None}
+    finally:
+        api.close()
+    title = f"{dataset} / {version.name} v{version.version}"
+    return outcome_of(final, run["id"], url, title, timeout_s, fail_below)
+
+
+def _poll_seconds(raw: str) -> float:
+    """``--poll``: at least one second, so the gate cannot hammer the server."""
+    value = float(raw)
+    if value < 1:
+        raise argparse.ArgumentTypeError("must be at least 1 second")
+    return value
+
+
+def _max_cost(raw: str) -> Decimal:
+    """``--max-cost-usd``: a positive amount of at most the server's ceiling of 1000."""
+    try:
+        value = Decimal(raw)
+    except InvalidOperation:
+        raise argparse.ArgumentTypeError("must be a number") from None
+    if not value.is_finite() or not 0 < value <= _MAX_COST_CEILING:
+        raise argparse.ArgumentTypeError("must be above 0 and at most 1000")
+    return value
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m spanlight._experiments",
+        description="Run a prompt over a dataset and fail when a score is below its threshold. "
+        "The access token is read from SPANLIGHT_TOKEN.",
+    )
+    env = os.environ.get
+    parser.add_argument("--host", default=env("SPANLIGHT_HOST"), help="Spanlight URL (https)")
+    parser.add_argument("--project-id", default=env("SPANLIGHT_PROJECT_ID"))
+    parser.add_argument("--dataset", required=True, help="dataset name")
+    parser.add_argument("--prompt", required=True, help="name, name@label or name:version")
+    parser.add_argument("--evaluators", default="", help="comma-separated evaluator names")
+    parser.add_argument("--fail-below", default="", help="e.g. relevance=0.8,format=1")
+    parser.add_argument("--credential-id", default=None, help="provider credential to run on")
+    parser.add_argument("--model", default=None, help="model to run the prompt on")
+    parser.add_argument(
+        "--max-cost-usd", type=_max_cost, default=None, help="run cost cap (new experiments)"
+    )
+    parser.add_argument("--timeout", type=int, default=1800, help="seconds to wait in total")
+    parser.add_argument("--poll", type=_poll_seconds, default=5.0, help="seconds between polls")
+    return parser
+
+
+def _publish(outcome: ExperimentOutcome, markdown: str) -> None:
+    """Print the report; inside GitHub Actions also fill the step summary and step outputs."""
+    print(markdown)
+    summary = os.environ.get("GITHUB_STEP_SUMMARY")
+    if summary:
+        with open(summary, "a", encoding="utf-8") as handle:
+            handle.write(markdown + "\n")
+    outputs = os.environ.get("GITHUB_OUTPUT")
+    if outputs:
+        with open(outputs, "a", encoding="utf-8") as handle:
+            handle.write(f"run-id={outcome.run_id}\nrun-url={outcome.url}\n")
+
+
+def _missing_settings(args: argparse.Namespace, token: str | None) -> list[str]:
+    """What the caller left out, as messages."""
+    checks = (
+        ("--host or SPANLIGHT_HOST is required", args.host),
+        ("--project-id or SPANLIGHT_PROJECT_ID is required", args.project_id),
+        ("--credential-id is required", args.credential_id),
+        ("--model is required", args.model),
+        ("SPANLIGHT_TOKEN is required", token),
+    )
+    return [message for message, value in checks if not value]
+
+
+def _gate(args: argparse.Namespace, token: str) -> int:
+    """Run the experiment, publish the report and return the exit code."""
+    fail_below = parse_fail_below(args.fail_below)
+    outcome = run_experiment(
+        host=args.host,
+        token=token,
+        project_id=args.project_id,
+        dataset=args.dataset,
+        prompt=args.prompt,
+        evaluators=[name.strip() for name in args.evaluators.split(",") if name.strip()],
+        fail_below=fail_below,
+        credential_id=args.credential_id,
+        model=args.model,
+        max_cost_usd=args.max_cost_usd,
+        timeout_s=args.timeout,
+        poll_s=args.poll,
+    )
+    report = render_markdown(
+        title=outcome.title,
+        status=outcome.status,
+        error=outcome.error,
+        url=outcome.url,
+        totals=outcome.totals,
+        fail_below=fail_below,
+        failures=outcome.failures,
+    )
+    _publish(outcome, report)
+    if outcome.status != "done":
+        return EXIT_FAILED
+    return EXIT_PASSED if outcome.passed else EXIT_THRESHOLD
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Run the gate. Returns the exit code: 0 passed, 1 a threshold missed, 2 anything else."""
+    args = _parser().parse_args(argv)
+    token = os.environ.get("SPANLIGHT_TOKEN")
+    problems = _missing_settings(args, token)
+    if problems:
+        print("error: " + "; ".join(problems), file=sys.stderr)
+        return EXIT_FAILED
+    try:
+        return _gate(args, token or "")
+    except (ExperimentError, ValueError) as error:
+        print(f"error: {error}", file=sys.stderr)
+    except Exception as error:
+        # Exit 1 means "a score fell below its threshold", so a bug or an answer of an
+        # unexpected shape must not look like one. The type only: a repr could carry a header.
+        print(f"error: unexpected {type(error).__name__}", file=sys.stderr)
+    return EXIT_FAILED
+
+
+if __name__ == "__main__":
+    sys.exit(main())