proofstep 0.1.0__py3-none-any.whl

proofstep/__init__.py

@@ -0,0 +1,192 @@
+"""Proofstep Python SDK — tracing for AI applications and tool-using agents.
+
+    import proofstep
+
+    proofstep.init(project="my-app")
+
+    @proofstep.trace("generate_outreach")
+    async def generate_outreach(prospect_id: str) -> Email:
+        ...
+
+    @proofstep.tool("gmail.send")
+    async def send_email(to: str, subject: str, body: str) -> str:
+        ...
+
+Two guarantees hold everywhere in this package:
+
+- **It never raises into your application.** Every public entry point is wrapped;
+  internal failures are logged once per window and return a no-op.
+- **It never blocks your application.** Export is a non-blocking enqueue onto a
+  bounded buffer drained by a background thread. When the buffer is full it drops
+  the oldest trace and counts it, because visible loss beats an invisible stall.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Iterator
+from importlib import metadata as _metadata
+from typing import Any
+
+from proofstep import redaction
+from proofstep.client import Captured, Client, sampled
+from proofstep.config import Config
+from proofstep.context import current_span, current_trace, propagate
+from proofstep.decorators import make_span, make_tool, make_trace
+from proofstep.propagation import extract, inject
+from proofstep.recorder import SpanRecorder, TraceRecorder
+from proofstep.safety import NOOP
+from proofstep_types import CaptureMode, SpanType, Status, Trace
+
+# Read from the installed distribution rather than written here twice. A hand-maintained
+# copy drifts the first time a release bumps one and not the other — which it already did,
+# reporting 0.1.0.dev0 from a 0.1.0 wheel.
+__version__ = _metadata.version("proofstep")
+
+_client: Client | None = None
+
+
+def init(**settings: Any) -> Client:
+    """Configure the SDK. Safe to call more than once; the last call wins."""
+    global _client  # noqa: PLW0603 — one process-wide client is the intended shape
+    _client = Client(Config.from_env(**settings))
+    return _client
+
+
+def get_client() -> Client:
+    """The active client, created from the environment on first use.
+
+    Implicit initialization is deliberate: an unconfigured import must still work,
+    so that adding a decorator never breaks a script that has not called `init`.
+    """
+    global _client  # noqa: PLW0603
+    if _client is None:
+        _client = Client(Config.from_env())
+    return _client
+
+
+def configure(**settings: Any) -> None:
+    """Update settings on the existing client without replacing it."""
+    client = get_client()
+    for key, value in settings.items():
+        if not hasattr(client.config, key):
+            msg = f"unknown Proofstep setting {key!r}"
+            raise TypeError(msg)
+        setattr(client.config, key, value)
+
+
+def reset() -> None:
+    """Drop the active client. For tests."""
+    global _client  # noqa: PLW0603
+    if _client is not None:
+        _client.shutdown(0.1)
+    _client = None
+
+
+trace = make_trace(get_client)
+span = make_span(get_client)
+tool = make_tool(span)
+
+
+@contextlib.contextmanager
+def start_trace(name: str, **kwargs: Any) -> Iterator[Any]:
+    """Context-manager form of `@trace`."""
+    with get_client().trace(name, **kwargs) as recorder:
+        yield recorder
+
+
+@contextlib.contextmanager
+def start_span(name: str, **kwargs: Any) -> Iterator[Any]:
+    """Context-manager form of `@span`."""
+    with get_client().span(name, **kwargs) as recorder:
+        yield recorder
+
+
+@contextlib.contextmanager
+def capture(name: str = "task", **kwargs: Any) -> Iterator[list[Trace]]:
+    """Record a trace and hand it back rather than only exporting it.
+
+    This is how an instrumented task feeds the local evaluation engine: the captured
+    `Trace` is what trajectory policies are evaluated against.
+
+        with proofstep.capture("classify") as captured:
+            result = await classify(example.input)
+        return proofstep.Captured(output=result, trace=captured[0])
+    """
+    sink: list[Trace] = []
+    with get_client().trace(name, **kwargs) as recorder:
+        yield sink
+        if recorder is not NOOP:
+            sink.append(recorder.snapshot())
+
+
+def set_metadata(**values: Any) -> None:
+    if (active := current_trace()) is not None:
+        active.set_metadata(**values)
+
+
+def set_tags(**values: str) -> None:
+    if (active := current_trace()) is not None:
+        active.set_tags(**values)
+
+
+def set_state(**values: Any) -> None:
+    """Record explicit workflow state for `final_state` policy rules."""
+    if (active := current_trace()) is not None:
+        active.set_state(**values)
+
+
+def record_event(name: str, **attributes: Any) -> None:
+    if (active := current_span()) is not None:
+        active.record_event(name, **attributes)
+
+
+def set_attributes(**values: Any) -> None:
+    if (active := current_span()) is not None:
+        active.set_attributes(**values)
+
+
+def flush(timeout: float | None = None) -> bool:
+    return bool(get_client().flush(timeout))
+
+
+def shutdown(timeout: float | None = None) -> None:
+    get_client().shutdown(timeout)
+
+
+__all__ = [
+    "NOOP",
+    "CaptureMode",
+    "Captured",
+    "Client",
+    "Config",
+    "SpanRecorder",
+    "SpanType",
+    "Status",
+    "Trace",
+    "TraceRecorder",
+    "capture",
+    "configure",
+    "current_span",
+    "current_trace",
+    "extract",
+    "flush",
+    "get_client",
+    "init",
+    "inject",
+    "propagate",
+    "record_event",
+    "redaction",
+    "reset",
+    "sampled",
+    "set_attributes",
+    "set_metadata",
+    "set_state",
+    "set_tags",
+    "shutdown",
+    "span",
+    "start_span",
+    "start_trace",
+    "tool",
+    "trace",
+]

proofstep/client.py

@@ -0,0 +1,209 @@
+"""The client: owns configuration, sampling, the exporter, and span creation."""
+
+from __future__ import annotations
+
+import atexit
+import contextlib
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Any
+
+from proofstep import context as ctx
+from proofstep.config import Config
+from proofstep.exporter import Exporter
+from proofstep.recorder import SpanRecorder, TraceRecorder, new_trace_id
+from proofstep.safety import NOOP, log_once, never_raises
+from proofstep_types import SpanType, Status, Trace
+
+
+@dataclass(frozen=True, slots=True)
+class Captured:
+    """A task's return value paired with the trace it produced.
+
+    Shaped to match what the evaluation runner already looks for (`.output` and
+    `.trace`), so an instrumented task drops into a suite with no adapter and no
+    dependency from the engine back to the SDK.
+    """
+
+    output: Any
+    trace: Trace
+
+
+def sampled(trace_id: str, rate: float) -> bool:
+    """Deterministic head sampling on the trace id.
+
+    Hash-mod rather than a coin flip per span, so a sampled trace is captured
+    *whole*. A half-recorded trajectory is worse than none: the policy engine would
+    read the gaps as evidence.
+    """
+    if rate >= 1.0:
+        return True
+    if rate <= 0.0:
+        return False
+    return (int(trace_id[:8], 16) / 0xFFFFFFFF) < rate
+
+
+class Client:
+    def __init__(
+        self, config: Config | None = None, *, transport: Any = None, **overrides: object
+    ) -> None:
+        self.config = config or Config.from_env(**overrides)
+        self.exporter = Exporter(self.config, transport=transport)
+        self._atexit_registered = False
+
+    # ------------------------------------------------------------------- tracing
+
+    @never_raises(default=NOOP)
+    def start_trace(
+        self,
+        name: str,
+        *,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> TraceRecorder | Any:
+        if not self.config.records:
+            return NOOP
+        identifier = trace_id or new_trace_id()
+        recorder = TraceRecorder(
+            name,
+            config=self.config,
+            trace_id=identifier,
+            parent_span_id=parent_span_id,
+            sampled=sampled(identifier, self.config.sample_rate),
+        )
+        if metadata:
+            recorder.set_metadata(**metadata)
+        self._register_atexit()
+        return recorder
+
+    @never_raises(default=NOOP)
+    def start_span(
+        self,
+        name: str,
+        *,
+        span_type: SpanType | str = SpanType.CUSTOM,
+        tool_name: str | None = None,
+        trace: TraceRecorder | None = None,
+        parent: SpanRecorder | None = None,
+    ) -> SpanRecorder | Any:
+        if not self.config.records:
+            return NOOP
+
+        active_trace = trace or ctx.current_trace()
+        if active_trace is None:
+            # An orphan span is more confusing than a synthetic root, and losing it
+            # entirely is worse than both. Create a trace so the span has somewhere
+            # to live, then emit it when the span closes.
+            #
+            # The common cause is a raw ThreadPoolExecutor: Python does not copy
+            # contextvars across threads, so the worker sees no active trace. Say so,
+            # because the fix (`proofstep.propagate`) is not discoverable otherwise.
+            log_once(
+                "client.orphan_span",
+                f"span {name!r} was created with no active trace and has been recorded "
+                "as its own trace. If this is a thread, wrap the callable with "
+                "proofstep.propagate() to keep it attached to its parent.",
+            )
+            active_trace = self.start_trace(name)
+            if active_trace is NOOP:
+                return NOOP
+            ctx.set_trace(active_trace)
+
+        active_parent = parent or ctx.current_span()
+        recorder = SpanRecorder(
+            name,
+            trace=active_trace,
+            span_type=SpanType(span_type),
+            parent_span_id=active_parent.span_id if active_parent else None,
+            tool_name=tool_name,
+            depth=(active_parent.depth + 1) if active_parent else 0,
+        )
+        if not active_trace.register(recorder):
+            return NOOP
+        return recorder
+
+    @contextlib.contextmanager
+    def trace(self, name: str, **kwargs: Any) -> Iterator[Any]:
+        recorder = self.start_trace(name, **kwargs)
+        if recorder is NOOP:
+            yield NOOP
+            return
+
+        token = ctx.set_trace(recorder)
+        try:
+            yield recorder
+        except BaseException as exc:
+            recorder.status = Status.ERROR
+            recorder.set_metadata(error=f"{type(exc).__name__}: {exc}"[:500])
+            raise
+        finally:
+            recorder.end()
+            ctx.reset_trace(token)
+            self.emit(recorder)
+
+    @contextlib.contextmanager
+    def span(self, name: str, **kwargs: Any) -> Iterator[Any]:
+        had_trace = ctx.current_trace() is not None
+        recorder = self.start_span(name, **kwargs)
+        if recorder is NOOP:
+            yield NOOP
+            return
+
+        token = ctx.set_span(recorder)
+        try:
+            yield recorder
+        except BaseException as exc:
+            # Record and re-raise, untouched. The traceback the user sees must be
+            # exactly the one their code produced.
+            if isinstance(exc, Exception):
+                recorder.set_error(exc)
+            else:
+                recorder.status = Status.ERROR
+            raise
+        finally:
+            recorder.end()
+            ctx.reset_span(token)
+            if not had_trace:
+                self._close_implicit_trace()
+
+    def _close_implicit_trace(self) -> None:
+        """Emit a trace `start_span` created on the caller's behalf.
+
+        Without this the span is recorded into a trace nobody ever ends, and the
+        data vanishes silently — the exact failure mode the SDK promises to avoid.
+        """
+        implicit = ctx.current_trace()
+        if implicit is None:
+            return
+        implicit.end()
+        self.emit(implicit)
+        ctx.set_trace(None)
+
+    # -------------------------------------------------------------------- export
+
+    @never_raises()
+    def emit(self, recorder: TraceRecorder) -> None:
+        """Finish a trace: snapshot it and hand it to the exporter."""
+        if not self.config.records:
+            return
+        keep = recorder.sampled or (
+            self.config.always_sample_on_error and recorder.status is Status.ERROR
+        )
+        if not keep:
+            return
+        self.exporter.submit(recorder.snapshot())
+
+    @never_raises(default=False)
+    def flush(self, timeout: float | None = None) -> bool:
+        return bool(self.exporter.flush(timeout))
+
+    @never_raises()
+    def shutdown(self, timeout: float | None = None) -> None:
+        self.exporter.shutdown(timeout)
+
+    def _register_atexit(self) -> None:
+        if self._atexit_registered:
+            return
+        self._atexit_registered = True
+        atexit.register(self.shutdown, self.config.shutdown_timeout_s)

proofstep/config.py

@@ -0,0 +1,152 @@
+"""SDK configuration.
+
+Precedence: explicit arguments, then environment variables, then defaults. Env vars
+matter more than they look — they are how a deployment turns capture down, or off,
+without a code change.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from proofstep_types import CaptureMode
+
+DEFAULT_ENDPOINT = "http://localhost:8000"
+
+# Deliberately conservative. A telemetry library that quietly buffers unbounded
+# work, or ships megabyte payloads by default, becomes the outage.
+DEFAULT_MAX_BUFFERED_SPANS = 10_000
+DEFAULT_BATCH_SIZE = 512
+DEFAULT_FLUSH_INTERVAL_S = 2.0
+DEFAULT_MAX_FIELD_BYTES = 256 * 1024
+DEFAULT_MAX_SPAN_BYTES = 1024 * 1024
+DEFAULT_MAX_BATCH_BYTES = 5 * 1024 * 1024
+
+
+def _env_bool(name: str, default: bool) -> bool:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() not in ("0", "false", "no", "off", "")
+
+
+def _env_float(name: str, default: float) -> float:
+    raw = os.environ.get(name)
+    try:
+        return float(raw) if raw is not None else default
+    except ValueError:
+        return default
+
+
+def _env_int(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    try:
+        return int(raw) if raw is not None else default
+    except ValueError:
+        return default
+
+
+@dataclass
+class Config:
+    """Resolved SDK settings."""
+
+    api_key: str | None = None
+    endpoint: str = DEFAULT_ENDPOINT
+    project: str | None = None
+    environment: str = "development"
+    capture_mode: CaptureMode = CaptureMode.REDACTED
+    sample_rate: float = 1.0
+    always_sample_on_error: bool = True
+
+    enabled: bool = True
+    export: bool = True
+    """When False the SDK records in-process but never sends. This is what `--local`
+    uses, and what makes the whole engine usable before anyone has an account."""
+
+    max_buffered_spans: int = DEFAULT_MAX_BUFFERED_SPANS
+    batch_size: int = DEFAULT_BATCH_SIZE
+    flush_interval_s: float = DEFAULT_FLUSH_INTERVAL_S
+    max_field_bytes: int = DEFAULT_MAX_FIELD_BYTES
+    max_span_bytes: int = DEFAULT_MAX_SPAN_BYTES
+    max_batch_bytes: int = DEFAULT_MAX_BATCH_BYTES
+    max_spans_per_trace: int = 10_000
+
+    export_timeout_s: float = 10.0
+    max_retries: int = 5
+    shutdown_timeout_s: float = 5.0
+
+    spool_dir: Path | None = None
+    """Where to persist batches the API refused. Off by default; on in CI, where a
+    lost run is a lost signal rather than a monitoring gap."""
+
+    git_commit: str | None = None
+    service_name: str | None = None
+    redact_keys: list[str] = field(default_factory=list)
+    debug: bool = False
+
+    @classmethod
+    def from_env(cls, **overrides: object) -> Config:
+        """Build a config from the environment, then apply explicit overrides."""
+        capture_raw = os.environ.get("PROOFSTEP_CAPTURE_MODE", "").strip().lower()
+        try:
+            capture = CaptureMode(capture_raw) if capture_raw else CaptureMode.REDACTED
+        except ValueError:
+            capture = CaptureMode.REDACTED
+
+        spool = os.environ.get("PROOFSTEP_SPOOL_DIR")
+
+        config = cls(
+            api_key=os.environ.get("PROOFSTEP_API_KEY"),
+            endpoint=os.environ.get("PROOFSTEP_ENDPOINT", DEFAULT_ENDPOINT).rstrip("/"),
+            project=os.environ.get("PROOFSTEP_PROJECT"),
+            environment=os.environ.get("PROOFSTEP_ENVIRONMENT", "development"),
+            capture_mode=capture,
+            sample_rate=_env_float("PROOFSTEP_SAMPLE_RATE", 1.0),
+            enabled=_env_bool("PROOFSTEP_ENABLED", default=True),
+            export=_env_bool("PROOFSTEP_EXPORT", default=True),
+            max_buffered_spans=_env_int("PROOFSTEP_MAX_BUFFERED_SPANS", DEFAULT_MAX_BUFFERED_SPANS),
+            batch_size=_env_int("PROOFSTEP_BATCH_SIZE", DEFAULT_BATCH_SIZE),
+            flush_interval_s=_env_float("PROOFSTEP_FLUSH_INTERVAL", DEFAULT_FLUSH_INTERVAL_S),
+            spool_dir=Path(spool) if spool else None,
+            git_commit=os.environ.get("PROOFSTEP_GIT_COMMIT") or _git_sha(),
+            service_name=os.environ.get("PROOFSTEP_SERVICE_NAME"),
+            debug=_env_bool("PROOFSTEP_DEBUG", default=False),
+        )
+
+        for key, value in overrides.items():
+            if value is None:
+                continue
+            if not hasattr(config, key):
+                msg = f"unknown Proofstep setting {key!r}"
+                raise TypeError(msg)
+            setattr(config, key, value)
+
+        config.capture_mode = CaptureMode(config.capture_mode)
+        config.sample_rate = min(1.0, max(0.0, config.sample_rate))
+        return config
+
+    @property
+    def records(self) -> bool:
+        return self.enabled and self.capture_mode is not CaptureMode.DISABLED
+
+    @property
+    def sends(self) -> bool:
+        return self.records and self.export and bool(self.api_key)
+
+    @property
+    def stores_payloads(self) -> bool:
+        return self.capture_mode.stores_payloads
+
+
+def _git_sha() -> str | None:
+    """Read the commit from CI environment variables only.
+
+    Deliberately does not shell out to git: a telemetry import must not fork a
+    subprocess, and in a container the repo usually is not there anyway.
+    """
+    for name in ("GITHUB_SHA", "GIT_COMMIT", "CI_COMMIT_SHA", "VERCEL_GIT_COMMIT_SHA"):
+        if value := os.environ.get(name):
+            return value
+    return None

proofstep/context.py

@@ -0,0 +1,74 @@
+"""Context propagation.
+
+`contextvars`, not thread-locals. `asyncio.create_task` and `TaskGroup` copy the
+current context automatically, so spans created inside concurrent children attach to
+the right parent with no user action — which is the whole reason to use contextvars
+here, since an agent framework spawns tasks constantly.
+
+Threads are the exception: Python does not copy contextvars across them, so
+`propagate()` is provided and its necessity is documented rather than hidden.
+"""
+
+from __future__ import annotations
+
+import contextvars
+import functools
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TypeVar
+
+if TYPE_CHECKING:
+    from proofstep.recorder import SpanRecorder, TraceRecorder
+
+T = TypeVar("T")
+
+_current_trace: contextvars.ContextVar[TraceRecorder | None] = contextvars.ContextVar(
+    "proofstep_trace", default=None
+)
+_current_span: contextvars.ContextVar[SpanRecorder | None] = contextvars.ContextVar(
+    "proofstep_span", default=None
+)
+
+
+def current_trace() -> TraceRecorder | None:
+    return _current_trace.get()
+
+
+def current_span() -> SpanRecorder | None:
+    return _current_span.get()
+
+
+def set_trace(trace: TraceRecorder | None) -> contextvars.Token[TraceRecorder | None]:
+    return _current_trace.set(trace)
+
+
+def set_span(span: SpanRecorder | None) -> contextvars.Token[SpanRecorder | None]:
+    return _current_span.set(span)
+
+
+def reset_trace(token: contextvars.Token[TraceRecorder | None]) -> None:
+    _current_trace.reset(token)
+
+
+def reset_span(token: contextvars.Token[SpanRecorder | None]) -> None:
+    _current_span.reset(token)
+
+
+def propagate(fn: Callable[..., T]) -> Callable[..., T]:  # noqa: UP047 — SDK targets py3.10
+    """Carry the current context into a callable that will run on another thread.
+
+    `ThreadPoolExecutor` does not copy contextvars, so a span created inside a
+    worker would otherwise attach to nothing and appear as an orphan.
+
+        pool.submit(proofstep.propagate(do_work), arg)
+    """
+    context = contextvars.copy_context()
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> T:
+        return context.run(fn, *args, **kwargs)
+
+    return wrapper
+
+
+def snapshot() -> tuple[TraceRecorder | None, SpanRecorder | None]:
+    return _current_trace.get(), _current_span.get()