trodo-python 1.2.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

trodo/otel/wrap_agent.py

@@ -0,0 +1,508 @@
+"""wrap_agent + join_run + span — user-facing surface for agent tracking.
+
+Users never generate IDs. The SDK creates run_id on wrap_agent and threads
+span_id through contextvars. Any nested ``span()`` call (or OTel
+auto-instrumented span from Anthropic/OpenAI/LangChain/etc.) inherits the
+current span as its parent.
+
+Usage (root service):
+    with trodo.wrap_agent('customer-support',
+                          distinct_id=uid, conversation_id=conv) as run:
+        result = agent.run(query)
+        run.set_output(result)
+
+Usage (downstream microservice — becomes a SPAN on the caller's run):
+    with trodo.join_run(run_id=headers['X-Trodo-Run-Id'],
+                        parent_span_id=headers['X-Trodo-Parent-Span-Id']):
+        # LLM calls here are auto-captured as spans under the caller's run
+        ...
+
+Decorator form:
+    @trodo.agent('nightly-report')
+    def nightly_report(): ...
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, Optional
+
+from .context import ActiveSpanContext, get_active_context, run_with_context
+from .processor import TrodoSpanProcessor, TrodoRun, TrodoSpan
+
+
+def _now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+
+
+def _truncate(value: Any, max_len: int = 64_000) -> Optional[str]:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        s = value
+    else:
+        try:
+            s = json.dumps(value, default=str)
+        except Exception:
+            s = str(value)
+    return s[:max_len] if len(s) > max_len else s
+
+
+def current_run_id() -> Optional[str]:
+    """Return the run_id of the currently active agent run, if any."""
+    ctx = get_active_context()
+    return ctx.run_id if ctx else None
+
+
+def current_span_id() -> Optional[str]:
+    """Return the span_id of the currently active span, if any."""
+    ctx = get_active_context()
+    return ctx.span_id if ctx else None
+
+
+def _aggregate(spans: list[TrodoSpan]) -> Dict[str, Any]:
+    total_in = 0
+    total_out = 0
+    total_cost = 0.0
+    tool_count = 0
+    error_count = 0
+    for s in spans:
+        if s.input_tokens:
+            total_in += int(s.input_tokens)
+        if s.output_tokens:
+            total_out += int(s.output_tokens)
+        if s.cost:
+            total_cost += float(s.cost)
+        if s.kind == "tool":
+            tool_count += 1
+        if s.status == "error":
+            error_count += 1
+    return {
+        "total_tokens_in": total_in,
+        "total_tokens_out": total_out,
+        "total_cost": round(total_cost, 8) if total_cost else 0.0,
+        "span_count": len(spans),
+        "tool_count": tool_count,
+        "error_count": error_count,
+    }
+
+
+class RunHandle:
+    """Handle returned by wrap_agent for setting input/output and getting run_id."""
+
+    def __init__(self, run_id: str, agent_name: str) -> None:
+        self.run_id = run_id
+        self.agent_name = agent_name
+        self.input: Optional[str] = None
+        self.output: Optional[str] = None
+        self.metadata: Dict[str, Any] = {}
+
+    def set_input(self, value: Any) -> None:
+        self.input = _truncate(value)
+
+    def set_output(self, value: Any) -> None:
+        self.output = _truncate(value)
+
+    def set_metadata(self, **kwargs: Any) -> None:
+        self.metadata.update(kwargs)
+
+
+class SpanHandle:
+    """Handle returned by span context manager for setting output/attrs."""
+
+    def __init__(self, span_id: str, name: str) -> None:
+        self.span_id = span_id
+        self.name = name
+        self.input: Optional[str] = None
+        self.output: Optional[str] = None
+        self.attributes: Dict[str, Any] = {}
+        self.model: Optional[str] = None
+        self.provider: Optional[str] = None
+        self.input_tokens: Optional[int] = None
+        self.output_tokens: Optional[int] = None
+        self.cost: Optional[float] = None
+        self.temperature: Optional[float] = None
+        self.tool_name: Optional[str] = None
+
+    def set_input(self, value: Any) -> None:
+        self.input = _truncate(value)
+
+    def set_output(self, value: Any) -> None:
+        self.output = _truncate(value)
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        self.attributes[key] = value
+
+    def set_llm(
+        self,
+        *,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
+        input_tokens: Optional[int] = None,
+        output_tokens: Optional[int] = None,
+        cost: Optional[float] = None,
+        temperature: Optional[float] = None,
+    ) -> None:
+        if model is not None:
+            self.model = model
+        if provider is not None:
+            self.provider = provider
+        if input_tokens is not None:
+            self.input_tokens = int(input_tokens)
+        if output_tokens is not None:
+            self.output_tokens = int(output_tokens)
+        if cost is not None:
+            self.cost = float(cost)
+        if temperature is not None:
+            self.temperature = float(temperature)
+
+    def set_tool(self, tool_name: str) -> None:
+        self.tool_name = tool_name
+
+
+def start_run(
+    *,
+    processor: TrodoSpanProcessor,
+    agent_name: str,
+    run_id: Optional[str] = None,
+    distinct_id: Optional[str] = None,
+    conversation_id: Optional[str] = None,
+    parent_run_id: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    input: Any = None,
+) -> str:
+    """Open a Run record without holding a context manager.
+
+    Pairs with :func:`end_run` for sessions that span multiple processes or
+    HTTP requests (e.g. an MCP server where ``initialize`` opens a run and
+    later ``tools/call`` requests append spans before a final close).
+
+    Returns the ``run_id`` (caller-supplied or freshly minted UUID). Between
+    ``start_run`` and ``end_run`` any process can use ``join_run(run_id, ...)``
+    to add spans — they flush incrementally via ``append_spans``.
+    """
+    rid = run_id or str(uuid.uuid4())
+    run = TrodoRun(
+        run_id=rid,
+        agent_name=agent_name,
+        distinct_id=distinct_id,
+        conversation_id=conversation_id,
+        parent_run_id=parent_run_id,
+        status="running",
+        input=_truncate(input),
+        started_at=_now_iso(),
+        metadata=metadata,
+    )
+    processor.mark_joined(rid)
+    processor.start_run(run)
+    return rid
+
+
+def end_run(
+    run_id: str,
+    *,
+    processor: TrodoSpanProcessor,
+    output: Any = None,
+    status: str = "ok",
+    error_summary: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Finalise a Run opened by :func:`start_run`.
+
+    Aggregates any locally-buffered spans for ``run_id``, POSTs to the
+    ``/runs/{id}/end`` endpoint, and unmarks the run as joined. Idempotent
+    on the local-state side; the backend treats a second call as a row update.
+    """
+    pending = processor.get_pending(run_id)
+    agg = _aggregate(pending)
+    payload: Dict[str, Any] = {
+        "ended_at": _now_iso(),
+        "status": status,
+        **agg,
+    }
+    if output is not None:
+        payload["output"] = _truncate(output)
+    if error_summary is not None:
+        payload["error_summary"] = _truncate(error_summary, 4_000)
+    if metadata is not None:
+        payload["metadata"] = metadata
+    processor.end_run(run_id, payload)
+    processor.unmark_joined(run_id)
+
+
+class wrap_agent:
+    """Context manager wrapping an agent run.
+
+    Every span created inside the ``with`` block (whether by the user's
+    manual ``span()`` calls or by OTel auto-instrumentation of the LLM SDK)
+    is automatically parented to this run via contextvars.
+    """
+
+    def __init__(
+        self,
+        processor: TrodoSpanProcessor,
+        team_site_id: str,
+        agent_name: str,
+        distinct_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        parent_run_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self._processor = processor
+        self._team_site_id = team_site_id
+        self._agent_name = agent_name
+        self._distinct_id = distinct_id
+        self._conversation_id = conversation_id
+        self._parent_run_id = parent_run_id
+        self._metadata = metadata
+        self._ctx_mgr: Optional[run_with_context] = None
+        self._started_ms: float = 0.0
+        self._started_iso: str = ""
+        self.handle: Optional[RunHandle] = None
+
+    def __enter__(self) -> RunHandle:
+        run_id = str(uuid.uuid4())
+        root_span_id = str(uuid.uuid4())
+        self._started_iso = _now_iso()
+        self._started_ms = time.time() * 1000.0
+
+        self.handle = RunHandle(run_id, self._agent_name)
+        ctx = ActiveSpanContext(
+            run_id=run_id,
+            span_id=root_span_id,
+            parent_span_id=None,
+            team_site_id=self._team_site_id,
+            processor=self._processor,
+        )
+        self._ctx_mgr = run_with_context(ctx)
+        self._ctx_mgr.__enter__()
+        return self.handle
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        assert self.handle is not None
+        ended_iso = _now_iso()
+        duration_ms = int(time.time() * 1000.0 - self._started_ms)
+        status = "error" if exc is not None else "ok"
+        error_summary = None
+        if exc is not None:
+            error_summary = _truncate(str(exc), 4_000)
+
+        pending = self._processor.get_pending(self.handle.run_id)
+        agg = _aggregate(pending)
+
+        run = TrodoRun(
+            run_id=self.handle.run_id,
+            agent_name=self._agent_name,
+            distinct_id=self._distinct_id,
+            conversation_id=self._conversation_id,
+            parent_run_id=self._parent_run_id,
+            status=status,
+            input=self.handle.input,
+            output=self.handle.output,
+            started_at=self._started_iso,
+            ended_at=ended_iso,
+            duration_ms=duration_ms,
+            error_summary=error_summary,
+            metadata={**(self._metadata or {}), **self.handle.metadata} or None,
+            total_tokens_in=agg["total_tokens_in"],
+            total_tokens_out=agg["total_tokens_out"],
+            total_cost=agg["total_cost"],
+            span_count=agg["span_count"],
+            tool_count=agg["tool_count"],
+            error_count=agg["error_count"],
+        )
+        try:
+            self._processor.ingest_run(run)
+        finally:
+            if self._ctx_mgr is not None:
+                self._ctx_mgr.__exit__(exc_type, exc, tb)
+        return None
+
+
+class join_run:
+    """Join an existing agent run owned by a remote service.
+
+    Opens a new SPAN (not a run) in the context of the caller's run_id.
+    Every span produced inside the ``with`` block — including OTel
+    auto-instrumented LLM spans — is parented to ``parent_span_id`` and
+    flushed incrementally via ``append_spans`` (the remote service owns
+    run lifecycle; we never call ``ingest_run`` here).
+    """
+
+    def __init__(
+        self,
+        processor: TrodoSpanProcessor,
+        team_site_id: str,
+        run_id: str,
+        parent_span_id: Optional[str] = None,
+        name: str = "remote.handler",
+        kind: str = "agent",
+        input: Any = None,
+        attributes: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self._processor = processor
+        self._team_site_id = team_site_id
+        self._run_id = run_id
+        self._parent_span_id = parent_span_id
+        self._name = name
+        self._kind = kind
+        self._input = _truncate(input) if input is not None else None
+        self._attributes = attributes
+        self._ctx_mgr: Optional[run_with_context] = None
+        self._started_ms: float = 0.0
+        self._started_iso: str = ""
+        self._span_id: str = ""
+        self.handle: Optional[SpanHandle] = None
+
+    def __enter__(self) -> SpanHandle:
+        self._span_id = str(uuid.uuid4())
+        self._started_iso = _now_iso()
+        self._started_ms = time.time() * 1000.0
+        self.handle = SpanHandle(self._span_id, self._name)
+        if self._input is not None:
+            self.handle.input = self._input
+        if self._attributes:
+            self.handle.attributes.update(self._attributes)
+
+        self._processor.mark_joined(self._run_id)
+        ctx = ActiveSpanContext(
+            run_id=self._run_id,
+            span_id=self._span_id,
+            parent_span_id=self._parent_span_id,
+            team_site_id=self._team_site_id,
+            processor=self._processor,
+        )
+        self._ctx_mgr = run_with_context(ctx)
+        self._ctx_mgr.__enter__()
+        return self.handle
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        if self.handle is None:
+            return None
+        ended_iso = _now_iso()
+        duration_ms = int(time.time() * 1000.0 - self._started_ms)
+        status = "error" if exc is not None else "ok"
+        error_type = exc_type.__name__ if exc_type else None
+        error_message = _truncate(str(exc), 4_000) if exc else None
+
+        trodo_span = TrodoSpan(
+            span_id=self._span_id,
+            run_id=self._run_id,
+            parent_span_id=self._parent_span_id,
+            kind=self._kind,
+            name=self._name,
+            status=status,
+            started_at=self._started_iso,
+            ended_at=ended_iso,
+            duration_ms=duration_ms,
+            input=self.handle.input,
+            output=self.handle.output,
+            error_type=error_type,
+            error_message=error_message,
+            model=self.handle.model,
+            provider=self.handle.provider,
+            input_tokens=self.handle.input_tokens,
+            output_tokens=self.handle.output_tokens,
+            cost=self.handle.cost,
+            temperature=self.handle.temperature,
+            tool_name=self.handle.tool_name,
+            attributes=self.handle.attributes or None,
+        )
+        try:
+            self._processor.append_spans(self._run_id, [trodo_span])
+        finally:
+            self._processor.unmark_joined(self._run_id)
+            if self._ctx_mgr is not None:
+                self._ctx_mgr.__exit__(exc_type, exc, tb)
+        return None
+
+
+class span:
+    """Context manager for a nested span under the active run.
+
+    Outside of a wrap_agent / join_run this is a no-op that still runs the
+    inner code.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        kind: str = "generic",
+        input: Any = None,
+        attributes: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self._name = name
+        self._kind = kind
+        self._input = _truncate(input) if input is not None else None
+        self._attributes = attributes
+        self._ctx_mgr: Optional[run_with_context] = None
+        self._started_ms: float = 0.0
+        self._started_iso: str = ""
+        self._active: Optional[ActiveSpanContext] = None
+        self._span_id: str = ""
+        self.handle: Optional[SpanHandle] = None
+
+    def __enter__(self) -> SpanHandle:
+        self._active = get_active_context()
+        self._span_id = str(uuid.uuid4())
+        self.handle = SpanHandle(self._span_id, self._name)
+        if self._input is not None:
+            self.handle.input = self._input
+        if self._attributes:
+            self.handle.attributes.update(self._attributes)
+        if self._active is None:
+            # No active run — still execute user code, just don't track.
+            return self.handle
+        child = ActiveSpanContext(
+            run_id=self._active.run_id,
+            span_id=self._span_id,
+            parent_span_id=self._active.span_id,
+            team_site_id=self._active.team_site_id,
+            processor=self._active.processor,
+        )
+        self._ctx_mgr = run_with_context(child)
+        self._ctx_mgr.__enter__()
+        self._started_iso = _now_iso()
+        self._started_ms = time.time() * 1000.0
+        return self.handle
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        if self._active is None or self.handle is None:
+            return None
+        ended_iso = _now_iso()
+        duration_ms = int(time.time() * 1000.0 - self._started_ms)
+        status = "error" if exc is not None else "ok"
+        error_type = exc_type.__name__ if exc_type else None
+        error_message = _truncate(str(exc), 4_000) if exc else None
+
+        trodo_span = TrodoSpan(
+            span_id=self._span_id,
+            run_id=self._active.run_id,
+            parent_span_id=self._active.span_id,
+            kind=self._kind,
+            name=self._name,
+            status=status,
+            started_at=self._started_iso,
+            ended_at=ended_iso,
+            duration_ms=duration_ms,
+            input=self.handle.input,
+            output=self.handle.output,
+            error_type=error_type,
+            error_message=error_message,
+            model=self.handle.model,
+            provider=self.handle.provider,
+            input_tokens=self.handle.input_tokens,
+            output_tokens=self.handle.output_tokens,
+            cost=self.handle.cost,
+            temperature=self.handle.temperature,
+            tool_name=self.handle.tool_name,
+            attributes=self.handle.attributes or None,
+        )
+        processor: TrodoSpanProcessor = self._active.processor  # type: ignore[assignment]
+        processor.enqueue_span(trodo_span)
+        if self._ctx_mgr is not None:
+            self._ctx_mgr.__exit__(exc_type, exc, tb)
+        return None

trodo/types.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Literal, Optional
 
 # ----------------------------------------------------------------------------
 # Configuration
@@ -21,6 +21,7 @@ class TrodoConfig:
     auto_events: bool = False
     on_error: Optional[Any] = None  # Callable[[Exception], None]
     debug: bool = False
+    auto_instrument: bool = False
 
 
 # ----------------------------------------------------------------------------
@@ -80,75 +81,18 @@ ResetResult = Dict[str, Any]
 
 
 # ----------------------------------------------------------------------------
-# Agent Analytics
+# Agent Runs (Lemma-style)
 # ----------------------------------------------------------------------------
 
-@dataclass
-class AgentCallProps:
-    agent_id: str
-    conversation_id: str
-    message_id: str
-    distinct_id: Optional[str] = None
-    prompt: Optional[str] = None
-    model: Optional[str] = None
-    temperature: Optional[float] = None
-    system_prompt_version: Optional[str] = None
-    provider: Optional[str] = None
-    timestamp: Optional[str] = None
-    properties: Dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class ToolUseProps:
-    agent_id: str
-    conversation_id: str
-    message_id: str
-    tool_name: str
-    distinct_id: Optional[str] = None
-    input: Optional[Any] = None
-    output: Optional[Any] = None
-    latency_ms: Optional[int] = None
-    status: Optional[str] = None      # 'success' | 'failure'
-    timestamp: Optional[str] = None
-    properties: Dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class AgentResponseProps:
-    agent_id: str
-    conversation_id: str
-    message_id: str
-    distinct_id: Optional[str] = None
-    output: Optional[str] = None
-    model: Optional[str] = None
-    completion_tokens: Optional[int] = None
-    prompt_tokens: Optional[int] = None
-    total_tokens: Optional[int] = None
-    finish_reason: Optional[str] = None
-    timestamp: Optional[str] = None
-    properties: Dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class AgentErrorProps:
-    agent_id: str
-    conversation_id: str
-    message_id: str
-    distinct_id: Optional[str] = None
-    error_type: Optional[str] = None
-    error_message: Optional[str] = None
-    failed_tool: Optional[str] = None
-    traceback: Optional[str] = None
-    timestamp: Optional[str] = None
-    properties: Dict[str, Any] = field(default_factory=dict)
-
-
 @dataclass
 class FeedbackProps:
-    agent_id: str
-    conversation_id: str
-    message_id: str
-    feedback: str                       # 'positive' | 'negative' | 'unreact'
+    """Feedback attached to a run. ``run_id`` is returned by ``wrap_agent``.
+
+    At least one of ``satisfaction`` / ``rating`` / ``comment`` must be provided.
+    """
+    satisfaction: Optional[Literal["positive", "negative"]] = None
+    rating: Optional[float] = None
+    comment: Optional[str] = None
+    feedback: Optional[str] = None  # alias for comment
     distinct_id: Optional[str] = None
-    timestamp: Optional[str] = None
-    properties: Dict[str, Any] = field(default_factory=dict)
+    metadata: Optional[Dict[str, Any]] = None