nullrun 0.4.0__py3-none-any.whl

nullrun/toolbox/__init__.py

@@ -0,0 +1,20 @@
+"""
+NullRun Toolbox.
+
+A curated set of higher-level, ready-to-use integration helpers for
+specific AI SDKs and frameworks. The `instrumentation/` package ships
+the low-level patches (httpx, OpenAI v1+ attribute path, auto mode);
+the `toolbox/` package ships opinionated wrappers that combine
+instrumentation + cost enforcement + workflow scoping for the most
+common agent runtimes (LangGraph, LlamaIndex, etc.).
+
+The split keeps the curated public surface (`nullrun.init`,
+`nullrun.protect`, `nullrun.track_*`) discoverable in `dir(nullrun)`
+while the framework-specific glue lives one import away at
+`nullrun.toolbox.<framework>`.
+"""
+from __future__ import annotations
+
+__all__ = [
+    "langgraph",
+]

nullrun/toolbox/langgraph.py

@@ -0,0 +1,94 @@
+"""
+LangGraph toolbox helpers for NullRun.
+
+This module is the user-facing entry point for LangGraph
+integrations. It is a thin convenience layer that wires the
+`NullRunCallback` from `nullrun.instrumentation.langgraph` onto a
+LangGraph compiled app so that every `app.invoke(...)` and
+`app.stream(...)` call fires the LangChain callback hooks. The
+callback extracts `input_tokens` / `output_tokens` from the LLM
+response and forwards them to the runtime's `track()` method —
+cost is then recomputed by the backend from the org's pricing
+policy.
+
+Why this lives in `toolbox/`, not `instrumentation/`:
+  - `instrumentation/` ships the generic, low-level patches
+    (httpx, OpenAI v1+ attribute path, LangChain callback class).
+    These are reusable building blocks.
+  - `toolbox/langgraph.py` ships a ready-to-use `wrapper(app)`
+    that is a single function call for the most common
+    LangGraph case. It is the entry point the user is pointed
+    to from the LangGraph integration docs.
+
+The previous location `nullrun.instrumentation.langgraph.instrument`
+is removed as of Phase 1 Commit 6. Users who imported it should
+switch to `nullrun.toolbox.langgraph.wrapper`.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from nullrun.instrumentation.langgraph import NullRunCallback
+from nullrun.runtime import NullRunRuntime, get_runtime
+
+logger = logging.getLogger(__name__)
+
+
+def wrapper(app: Any, runtime: Any | None = None) -> Any:
+    """
+    Wrap a compiled LangGraph app with NullRun tracking.
+
+    Every `app.invoke(...)` and `app.stream(...)` call gets a
+    `NullRunCallback` attached so the runtime sees the LLM
+    usage for cost accounting and policy enforcement.
+
+    Usage:
+        from nullrun import init
+        from nullrun.toolbox.langgraph import wrapper
+
+        runtime = init()
+        graph = build_my_graph()
+        graph = wrapper(graph, runtime=runtime)
+
+        result = graph.invoke({"messages": [("user", "hi")]})
+
+    Args:
+        app: A compiled LangGraph `StateGraph` (anything with
+             `.invoke` and `.stream`).
+        runtime: Optional `NullRunRuntime`. Defaults to the
+             module-level singleton from `get_runtime()`.
+
+    Returns:
+        The same `app` object, with `.invoke` and `.stream`
+        wrapped in place. The callback is added to LangChain's
+        `config["callbacks"]` list per call, so multiple
+        wrappers compose without colliding.
+    """
+    rt: NullRunRuntime = runtime or get_runtime()
+    callback = NullRunCallback(runtime=rt)
+    original_invoke = getattr(app, "invoke", None)
+    original_stream = getattr(app, "stream", None)
+
+    if original_invoke is not None:
+        def wrapped_invoke(input: Any, config: Any | None = None, **kwargs: Any) -> Any:
+            if config is None:
+                config = {}
+            if "callbacks" not in config:
+                config["callbacks"] = []
+            config["callbacks"].append(callback)
+            return original_invoke(input, config, **kwargs)
+        app.invoke = wrapped_invoke
+
+    if original_stream is not None:
+        def wrapped_stream(input: Any, config: Any | None = None, **kwargs: Any) -> Any:
+            if config is None:
+                config = {}
+            if "callbacks" not in config:
+                config["callbacks"] = []
+            config["callbacks"].append(callback)
+            return original_stream(input, config, **kwargs)
+        app.stream = wrapped_stream
+
+    logger.info("LangGraph app wrapped with NullRun tracking")
+    return app

nullrun/tracing.py

@@ -0,0 +1,155 @@
+"""
+Trace/span context management via Python contextvars.
+
+This module is the core of the new trace/span system (Phase 2 of
+the SDK cleanup plan). The previous `nullrun.context` module
+exposed loose `_trace_id` and `_span_id` contextvars — fine for
+attaching IDs to events, but it didn't model the parent/child
+hierarchy that a trace timeline needs.
+
+`SpanContext` is a structured value: a single contextvar holds
+the *current* span, and child spans are derived from it via
+`create_child_span(parent)`. This is the same pattern OpenTelemetry
+uses for its Python SDK (`opentelemetry.context.get_current`) and
+gives `@protect` (Commit 4) and `track_*` (Commit 5) a uniform
+way to attach `trace_id` / `span_id` / `parent_span_id` / `depth`
+to every emitted event.
+
+Thread/async safety: `ContextVar` is thread-local by default but
+PEP 567 guarantees the right value is restored across `await`
+boundaries in asyncio, so concurrent coroutines each see their
+own current span.
+
+What this module does NOT do:
+  - It does not emit events. `SpanContext` is a pure data
+    structure. The runtime's `track_event()` is what actually
+    posts `span_start` / `span_end` events to the backend. See
+    `_emit_span_start` / `_emit_span_end` in `nullrun.decorators`
+    for the wiring.
+  - It does not implement OTel-style attributes, status, or
+    exception recording. We keep the surface minimal — a span
+    is just an ID tuple, the dashboard reconstructs the rest
+    from the event stream.
+"""
+from __future__ import annotations
+
+import uuid
+from contextvars import ContextVar
+from dataclasses import dataclass
+
+
+def _new_id() -> str:
+    """Generate a fresh span/trace id.
+
+    Returns a real UUID4 with dashes (e.g. ``95ca7c0b-...-2788803ef3b8``)
+    so the backend's `Uuid::parse_str` accepts it on the wire. Earlier
+    we shipped `uuid.uuid4().hex` (32 hex chars, no dashes) which the
+    backend silently dropped to NULL.
+    """
+    return str(uuid.uuid4())
+
+
+@dataclass(frozen=True)
+class SpanContext:
+    """
+    One span in the call tree.
+
+    Attributes:
+        trace_id:    Stable across the whole trace (root + all descendants).
+        span_id:     Unique to this span. Children reference it as
+                     `parent_span_id`.
+        parent_span_id: The parent's `span_id`, or None for the root span.
+        depth:       0 for the root, parent.depth + 1 for each child.
+                     Useful for the waterfall UI's indentation.
+    """
+
+    trace_id: str
+    span_id: str
+    parent_span_id: str | None = None
+    depth: int = 0
+
+
+# The currently-active span. `None` means "no trace in progress" — track_*
+# will fall back to creating a synthetic root on each call so events are
+# still attributed to *something*.
+_current_span: ContextVar[SpanContext | None] = ContextVar(
+    "nullrun_span", default=None
+)
+
+
+def get_current_span() -> SpanContext | None:
+    """
+    Return the active span, or None if no `@protect` / manual `set_span`
+    has put us inside a trace.
+    """
+    return _current_span.get()
+
+
+def create_child_span(parent: SpanContext) -> SpanContext:
+    """
+    Derive a new child span from `parent`.
+
+    The child inherits `parent.trace_id` and increments `parent.depth`.
+    `parent_span_id` is set to `parent.span_id` so the tree is fully
+    reconstructable from the event stream.
+
+    Raises:
+        ValueError: if `parent` is ``None``. The function does NOT
+            silently degrade to creating a root span — that would
+            hide bugs in the caller where a parent was expected.
+            Sprint 2.6 (B5): pre-fix this raised
+            ``TypeError: unsupported operand for None + 1`` on
+            ``parent.depth + 1`` which crashed the entire
+            ``@protect`` / track_* pipeline. Raise a clear
+            ``ValueError`` instead so the caller can fix the bug.
+    """
+    if parent is None:
+        raise ValueError(
+            "create_child_span requires a non-None parent SpanContext. "
+            "If you want a root span, use create_root_span() instead."
+        )
+    return SpanContext(
+        trace_id=parent.trace_id,
+        span_id=_new_id(),
+        parent_span_id=parent.span_id,
+        depth=parent.depth + 1,
+    )
+
+
+def create_root_span() -> SpanContext:
+    """
+    Start a new trace. Returns a SpanContext with no parent and depth 0.
+    """
+    tid = _new_id()
+    return SpanContext(
+        trace_id=tid,
+        span_id=_new_id(),
+        parent_span_id=None,
+        depth=0,
+    )
+
+
+def set_span(ctx: SpanContext):
+    """
+    Make `ctx` the current span. Returns a token that MUST be passed
+    back to `reset_span` in a `finally` block to restore the previous
+    context (which may itself be None).
+
+    Usage:
+        span = create_root_span()
+        token = set_span(span)
+        try:
+            ...
+        finally:
+            reset_span(token)
+    """
+    return _current_span.set(ctx)
+
+
+def reset_span(token) -> None:
+    """
+    Restore the context that was active before the matching `set_span`.
+    Pair with `set_span` — never call reset_span with a token from a
+    different context.
+    """
+    _current_span.reset(token)