spanforge 2.0.0__py3-none-any.whl

spanforge/_trace.py

@@ -0,0 +1,334 @@
+"""spanforge._trace — :class:`Trace` object and :func:`start_trace` entry point.
+
+A :class:`Trace` wraps a *root* :class:`~spanforge._span.AgentRunContextManager`
+and gives callers a convenient imperative handle to the entire agent execution::
+
+    trace = spanforge.start_trace("research_agent")
+
+    with trace.llm_call(model="gpt-4o") as s:
+        s.set_attribute("query", "latest AI papers")
+
+    with trace.tool_call("search") as s:
+        s.set_attribute("results_count", 5)
+
+    trace.end()
+
+The :class:`Trace` object also provides convenience methods for serialisation
+(:meth:`to_json`, :meth:`save`) that will become richer in later phases when
+debug utilities (:meth:`print_tree`, :meth:`summary`) are added.
+
+Design notes
+------------
+* :func:`start_trace` opens the underlying :class:`AgentRunContextManager`
+  immediately (``__enter__``), so all subsequently created spans automatically
+  inherit the trace's ``trace_id``.
+* :meth:`Trace.end` closes that context manager (``__exit__``).  Callers
+  *must* call :meth:`end` or use the :class:`Trace` as a context manager::
+
+      with spanforge.start_trace("my-agent") as trace:
+          ...
+          # trace.end() called automatically on exit
+
+* :class:`Trace` collects emitted :class:`~spanforge._span.Span` instances via
+  :meth:`_record_span` which is called by the stream module when a
+  ``_TRACE_COLLECTOR`` attribute is present on the active
+  :class:`~spanforge._span.AgentRunContext`.  This allows :meth:`to_json` and
+  :meth:`save` to operate on in-memory data without re-reading from a file.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from spanforge._span import (
+    AgentRunContext,
+    AgentRunContextManager,
+    Span,
+    SpanContextManager,
+    _span_id,
+    _trace_id,
+    _now_ns,
+)
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+__all__ = ["Trace", "start_trace"]
+
+
+@dataclass
+class Trace:
+    """A handle to a complete agent trace.
+
+    Created by :func:`start_trace`; do not construct directly.
+
+    Attributes:
+        agent_name:   Name of the agent being traced.
+        trace_id:     32-char hex OTel-compatible trace ID.
+        start_time:   Unix timestamp (float seconds) when the trace started.
+    """
+
+    agent_name: str
+    trace_id: str
+    start_time: float
+    _run_ctx: AgentRunContext = field(repr=False)
+    _run_cm: AgentRunContextManager = field(repr=False)
+    attributes: dict[str, Any] = field(default_factory=dict)
+    _spans: list[Span] = field(default_factory=list, repr=False)
+    _ended: bool = field(default=False, init=False, repr=False)
+
+    # ------------------------------------------------------------------
+    # Span convenience methods
+    # ------------------------------------------------------------------
+
+    def llm_call(
+        self,
+        model: str | None = None,
+        *,
+        operation: str = "chat",
+        temperature: float | None = None,
+        top_p: float | None = None,
+        max_tokens: int | None = None,
+        attributes: dict[str, Any] | None = None,
+    ) -> SpanContextManager:
+        """Open a child span for an LLM call within this trace.
+
+        Example::
+
+            with trace.llm_call(model="gpt-4o", temperature=0.7) as s:
+                s.set_attribute("prompt_tokens", 512)
+        """
+        name = f"llm_call:{model}" if model else "llm_call"
+        return SpanContextManager(
+            name=name,
+            model=model,
+            operation=operation,
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+            attributes=attributes,
+        )
+
+    def tool_call(
+        self,
+        tool_name: str,
+        *,
+        attributes: dict[str, Any] | None = None,
+    ) -> SpanContextManager:
+        """Open a child span for a tool call within this trace.
+
+        Example::
+
+            with trace.tool_call("search") as s:
+                results = search(query)
+        """
+        return SpanContextManager(
+            name=f"tool_call:{tool_name}",
+            operation="tool",
+            attributes=attributes,
+        )
+
+    def span(
+        self,
+        name: str,
+        *,
+        model: str | None = None,
+        operation: str = "chat",
+        temperature: float | None = None,
+        top_p: float | None = None,
+        max_tokens: int | None = None,
+        attributes: dict[str, Any] | None = None,
+    ) -> SpanContextManager:
+        """Open a generic child span within this trace."""
+        return SpanContextManager(
+            name=name,
+            model=model,
+            operation=operation,
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+            attributes=attributes,
+        )
+
+    # ------------------------------------------------------------------
+    # Internal span collection
+    # ------------------------------------------------------------------
+
+    def _record_span(self, span: Span) -> None:
+        """Called by _stream.emit_span when this trace is the active run context."""
+        self._spans.append(span)
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def end(self) -> None:
+        """Close the trace and emit the agent-run completion event.
+
+        Idempotent — subsequent calls are no-ops.
+        """
+        if self._ended:
+            return
+        self._ended = True
+        self._run_cm.__exit__(None, None, None)
+
+    # ------------------------------------------------------------------
+    # Context manager protocol (allows ``with start_trace(...) as trace:``)
+    # ------------------------------------------------------------------
+
+    def __enter__(self) -> "Trace":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        if not self._ended:
+            self._run_cm.__exit__(exc_type, exc_val, exc_tb)
+            self._ended = True
+        return False
+
+    async def __aenter__(self) -> "Trace":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        return self.__exit__(exc_type, exc_val, exc_tb)
+
+    # ------------------------------------------------------------------
+    # Serialisation
+    # ------------------------------------------------------------------
+
+    def to_json(self, *, indent: int | None = None) -> str:
+        """Serialise the trace's collected spans to a JSON string.
+
+        Returns a JSON object with ``trace_id``, ``agent_name``, ``start_time``,
+        and a ``spans`` array of span payload dicts.
+
+        Args:
+            indent: Optional JSON indentation level (``None`` = compact).
+        """
+        return json.dumps(self._to_dict(), indent=indent, sort_keys=True, default=str)
+
+    def save(self, path: str) -> None:
+        """Write the trace as NDJSON (one span per line) to *path*.
+
+        Args:
+            path: File path to write to.  The file is created or overwritten.
+        """
+        lines = [
+            json.dumps(span.to_span_payload().to_dict(), sort_keys=True, default=str)
+            for span in self._spans
+        ]
+        with open(path, "w", encoding="utf-8") as fh:
+            fh.write("\n".join(lines))
+            if lines:
+                fh.write("\n")
+
+    def _to_dict(self) -> dict[str, Any]:
+        d: dict[str, Any] = {
+            "trace_id": self.trace_id,
+            "agent_name": self.agent_name,
+            "start_time": self.start_time,
+            "spans": [span.to_span_payload().to_dict() for span in self._spans],
+        }
+        if self.attributes:
+            d["attributes"] = self.attributes
+        return d
+
+    # Placeholder methods for Phase 3 debug utilities.
+    def print_tree(self, *, file: Any = None) -> None:
+        """Pretty-print a hierarchical tree of spans to stdout.
+
+        Delegates to :func:`spanforge.debug.print_tree`.
+        Requires the trace to have ended (or have accumulated spans).
+        """
+        from spanforge.debug import print_tree  # noqa: PLC0415
+        print_tree(self._spans, file=file)
+
+    def summary(self) -> dict[str, Any]:
+        """Return an aggregated statistics dict for the trace's spans.
+
+        Delegates to :func:`spanforge.debug.summary`.
+        """
+        from spanforge.debug import summary  # noqa: PLC0415
+        return summary(self._spans)
+
+    def visualize(self, *, output: str = "html", path: str | None = None) -> str:
+        """Generate a self-contained HTML Gantt timeline for this trace.
+
+        Delegates to :func:`spanforge.debug.visualize`.
+
+        Args:
+            output: Output format — currently only ``"html"``.
+            path:   Optional file path to write the HTML to.
+
+        Returns:
+            HTML string.
+        """
+        from spanforge.debug import visualize  # noqa: PLC0415
+        return visualize(self._spans, output=output, path=path)
+
+
+# ---------------------------------------------------------------------------
+# start_trace()
+# ---------------------------------------------------------------------------
+
+
+def start_trace(agent_name: str, **attributes: Any) -> Trace:
+    """Start a new agent trace and return a :class:`Trace` handle.
+
+    Opens a root :class:`~spanforge._span.AgentRunContextManager` so all spans
+    created within the trace automatically inherit the correct ``trace_id``.
+
+    Must be closed by calling :meth:`Trace.end` or by using the returned
+    object as a context manager::
+
+        # Imperative style
+        trace = spanforge.start_trace("research_agent")
+        with trace.llm_call(model="gpt-4o"):
+            ...
+        trace.end()
+
+        # Context manager style (recommended)
+        with spanforge.start_trace("research_agent") as trace:
+            with trace.llm_call(model="gpt-4o"):
+                ...
+
+    Args:
+        agent_name:  Name of the agent being executed.
+        **attributes: Optional key-value attributes stored on the
+                      :attr:`Trace.attributes` dict and included in
+                      :meth:`Trace.to_json` output.
+
+    Returns:
+        A :class:`Trace` object that acts as the root execution context.
+    """
+    if not isinstance(agent_name, str) or not agent_name:
+        raise ValueError("start_trace: agent_name must be a non-empty string")
+
+    cm = AgentRunContextManager(agent_name=agent_name)
+    run_ctx = cm.__enter__()
+
+    # Attach a back-reference so _stream can call _record_span on this Trace.
+    trace = Trace(
+        agent_name=agent_name,
+        trace_id=run_ctx.trace_id,
+        start_time=time.time(),
+        _run_ctx=run_ctx,
+        _run_cm=cm,
+        attributes=dict(attributes) if attributes else {},
+    )
+    # Store the Trace on the AgentRunContext so _stream.emit_span can find it.
+    run_ctx._trace_collector = trace  # type: ignore[attr-defined]
+
+    return trace

spanforge/_tracer.py

@@ -0,0 +1,253 @@
+"""spanforge._tracer — :class:`Tracer` class and module-level ``tracer`` singleton.
+
+The :class:`Tracer` is the primary entry point for instrumenting code with
+SpanForge.  Import the module-level singleton ``tracer`` and use its context
+managers to create spans and agent traces::
+
+    from spanforge import tracer, configure
+
+    configure(exporter="console")
+
+    with tracer.span("chat", model="gpt-4o") as s:
+        s.set_attribute("prompt_tokens", 512)
+
+    with tracer.agent_run("research-agent") as run:
+        with tracer.agent_step("web-search") as step:
+            step.set_attribute("query", "what is RAG?")
+        with tracer.agent_step("summarize"):
+            pass
+
+    # Imperative trace API (Phase 1)
+    trace = tracer.start_trace("research-agent")
+    with trace.llm_call(model="gpt-4o"):
+        ...
+    trace.end()
+
+All context managers support both ``with`` and ``async with``.  Async usage
+is safe because the span stack is backed by :mod:`contextvars` so each asyncio
+task sees its own stack slice.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from spanforge._span import (
+    AgentRunContextManager,
+    AgentStepContextManager,
+    SpanContextManager,
+)
+from spanforge._trace import Trace, start_trace as _start_trace
+from spanforge.trace import trace as _trace_decorator, _TraceDecorator
+
+__all__ = ["Tracer", "tracer"]
+
+
+class Tracer:
+    """The SpanForge tracing façade.
+
+    A single module-level instance is created as :data:`tracer` and is the
+    recommended way to instrument code.  Creating additional :class:`Tracer`
+    instances is supported but shares the same thread-local context stacks.
+
+    All ``span``/``agent_run``/``agent_step`` methods return context managers
+    that push the new context onto the thread-local stack on ``__enter__`` and
+    pop it (and emit the event) on ``__exit__``.
+    """
+
+    # ------------------------------------------------------------------
+    # Span API  (Phase 2)
+    # ------------------------------------------------------------------
+
+    def span(
+        self,
+        name: str,
+        *,
+        model: str | None = None,
+        operation: str = "chat",
+        temperature: float | None = None,
+        top_p: float | None = None,
+        max_tokens: int | None = None,
+        attributes: dict[str, Any] | None = None,
+    ) -> SpanContextManager:
+        """Create a new :class:`~spanforge._span.SpanContextManager`.
+
+        Use as a context manager::
+
+            with tracer.span("llm-call", model="gpt-4o", temperature=0.7) as s:
+                s.set_attribute("prompt_tokens", 512)
+
+        Args:
+            name:        Human-readable span name (non-empty string).
+            model:       Model name string (e.g. ``"gpt-4o"``).  Used to infer
+                         the provider when no integration has set
+                         :attr:`~spanforge._span.Span.token_usage`.
+            operation:   GenAI operation name (default ``"chat"``).  Any
+                         :class:`~spanforge.namespaces.trace.GenAIOperationName`
+                         value or a custom string.
+            temperature: Sampling temperature forwarded to :class:`SpanPayload`.
+            top_p:       Nucleus sampling ``top_p`` value.
+            max_tokens:  Maximum token limit for this LLM call.
+            attributes:  Initial key-value attributes.  Additional attributes
+                         can be added inside the block via
+                         :meth:`~spanforge._span.Span.set_attribute`.
+
+        Returns:
+            A :class:`~spanforge._span.SpanContextManager` that yields a
+            :class:`~spanforge._span.Span` on ``__enter__``.
+        """
+        return SpanContextManager(
+            name=name,
+            model=model,
+            operation=operation,
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+            attributes=attributes,
+        )
+
+    # ------------------------------------------------------------------
+    # Agent API  (Phase 4)
+    # ------------------------------------------------------------------
+
+    def agent_run(self, agent_name: str) -> AgentRunContextManager:
+        """Create a root agent-run context manager.
+
+        Use as an outer context that wraps one or more ``agent_step`` calls::
+
+            with tracer.agent_run("my-agent") as run:
+                with tracer.agent_step("step-1"):
+                    ...
+
+        On exit, emits an
+        :data:`~spanforge.types.EventType.TRACE_AGENT_COMPLETED` event with
+        aggregated totals across all child steps.
+
+        Args:
+            agent_name: Name of the agent (non-empty string).
+
+        Returns:
+            :class:`~spanforge._span.AgentRunContextManager`
+        """
+        return AgentRunContextManager(agent_name=agent_name)
+
+    def agent_step(
+        self,
+        step_name: str,
+        *,
+        operation: str = "invoke_agent",
+        attributes: dict[str, Any] | None = None,
+    ) -> AgentStepContextManager:
+        """Create a single agent-step context manager.
+
+        Must be used inside an ``agent_run`` block::
+
+            with tracer.agent_run("my-agent"):
+                with tracer.agent_step("search") as step:
+                    step.set_attribute("query", "hello")
+
+        On exit, emits an
+        :data:`~spanforge.types.EventType.TRACE_AGENT_STEP` event.
+
+        Args:
+            step_name:  Human-readable step name.
+            operation:  GenAI operation name (default ``"invoke_agent"``).
+            attributes: Initial key-value attributes.
+
+        Returns:
+            :class:`~spanforge._span.AgentStepContextManager`
+
+        Raises:
+            RuntimeError: If called outside an ``agent_run`` context.
+        """
+        return AgentStepContextManager(
+            step_name=step_name,
+            operation=operation,
+            attributes=attributes,
+        )
+
+    # ------------------------------------------------------------------
+    # Trace API  (Phase 1)
+    # ------------------------------------------------------------------
+
+    def start_trace(self, agent_name: str, **attributes: Any) -> Trace:
+        """Start a new agent trace and return a :class:`~spanforge._trace.Trace` handle.
+
+        Convenience wrapper around the module-level :func:`~spanforge._trace.start_trace`.
+        The returned :class:`Trace` must be closed with :meth:`~spanforge._trace.Trace.end`
+        or used as a context manager::
+
+            with tracer.start_trace("research-agent") as trace:
+                with trace.llm_call(model="gpt-4o") as s:
+                    ...
+
+        Args:
+            agent_name:   Name of the agent being executed.
+            **attributes: Optional key-value attributes for the root run context.
+
+        Returns:
+            :class:`~spanforge._trace.Trace`
+        """
+        return _start_trace(agent_name, **attributes)
+
+    def trace(
+        self,
+        fn: Any = None,
+        *,
+        name: str | None = None,
+        model: str | None = None,
+        operation: str = "chat",
+        tool: bool = False,
+        capture_args: bool = False,
+        capture_return: bool = False,
+        attributes: dict[str, Any] | None = None,
+    ) -> Any:
+        """Decorator that instruments a function as an SpanForge span.
+
+        Delegates to :func:`~spanforge.trace.trace`.  Provided as a convenience
+        method so callers who already hold a :class:`Tracer` reference do not
+        need a separate import::
+
+            @tracer.trace(name="my-step", model="gpt-4o")
+            def call_llm(prompt: str) -> str: ...
+
+            @tracer.trace(name="async-step")
+            async def async_step(x: int) -> dict: ...
+
+        Args:
+            fn:             Function to wrap when used bare (``@tracer.trace``).
+            name:           Span name (defaults to ``fn.__qualname__``).
+            model:          Model identifier string.
+            operation:      GenAI operation name (default ``"chat"``).
+            tool:           Mark as tool call; sets operation to ``"execute_tool"``.
+            capture_args:   Record call arguments as span attributes.
+            capture_return: Record return value as a span attribute.
+            attributes:     Static key-value attributes on every span.
+
+        Returns:
+            Decorated callable, or a single-argument decorator.
+        """
+        return _trace_decorator(
+            fn,
+            name=name,
+            model=model,
+            operation=operation,
+            tool=tool,
+            capture_args=capture_args,
+            capture_return=capture_return,
+            attributes=attributes,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Module-level singleton — ``from spanforge import tracer``
+# ---------------------------------------------------------------------------
+
+#: The default :class:`Tracer` singleton.
+#:
+#: Import this directly for convenience::
+#:
+#:     from spanforge import tracer
+#:     with tracer.span("my-span"):
+#:         ...
+tracer: Tracer = Tracer()