trodo-python 1.0.1__py3-none-any.whl → 2.1.0__py3-none-any.whl

trodo/otel/helpers.py

@@ -0,0 +1,407 @@
+"""Modular instrumentation helpers: span helpers (trace / tool / llm /
+retrieval), one-shot ``track_llm_call``, and FastAPI middleware for
+cross-service run joining.
+
+These are the SDK's customer-facing surface for custom code: one dual-form
+wrapper per span kind, usable either as ``helper("name", fn)`` or as a
+``@helper("name")`` decorator. All forms auto-capture args/kwargs as
+``input``, return value as ``output``, exception as ``error``.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import Any, Awaitable, Callable, Dict, Optional, Tuple, TypeVar
+
+from .context import get_active_context
+from .wrap_agent import SpanHandle, join_run, span as span_ctx
+
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _make_input(fn: Callable[..., Any], args: tuple, kwargs: dict) -> Dict[str, Any]:
+    """Bind positional+keyword args to their parameter names.
+
+    Named kwargs are the common case; positional args fall back to 'args'
+    if signature binding fails.
+    """
+    payload: Dict[str, Any] = {}
+    if args:
+        try:
+            sig = inspect.signature(fn)
+            bound = sig.bind_partial(*args, **kwargs)
+            payload.update(bound.arguments)
+        except Exception:
+            payload["args"] = list(args)
+            payload["kwargs"] = dict(kwargs)
+    else:
+        payload.update(kwargs)
+    return payload
+
+
+def _wrap_fn(
+    fn: F,
+    *,
+    name: str,
+    kind: str,
+    on_result: Optional[Callable[[SpanHandle, Any], None]] = None,
+    extra_set: Optional[Callable[[SpanHandle], None]] = None,
+) -> F:
+    """Wrap fn so each call opens a span (sync + async aware).
+
+    ``extra_set`` runs on the span handle before the call (e.g. set_llm).
+    ``on_result`` runs on the span handle with the return value (e.g.
+    extracting OpenAI/Gemini usage into the LLM span).
+    """
+    is_async = inspect.iscoroutinefunction(fn)
+
+    if is_async:
+
+        @functools.wraps(fn)
+        async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with span_ctx(name, kind=kind, input=_make_input(fn, args, kwargs)) as s:
+                if kind == "tool":
+                    s.set_tool(name)
+                if extra_set is not None:
+                    extra_set(s)
+                result = await fn(*args, **kwargs)
+                if on_result is not None:
+                    try:
+                        on_result(s, result)
+                    except Exception:
+                        pass
+                s.set_output(result)
+                return result
+
+        return _async_wrapper  # type: ignore[return-value]
+
+    @functools.wraps(fn)
+    def _sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        with span_ctx(name, kind=kind, input=_make_input(fn, args, kwargs)) as s:
+            if kind == "tool":
+                s.set_tool(name)
+            if extra_set is not None:
+                extra_set(s)
+            result = fn(*args, **kwargs)
+            if on_result is not None:
+                try:
+                    on_result(s, result)
+                except Exception:
+                    pass
+            s.set_output(result)
+            return result
+
+    return _sync_wrapper  # type: ignore[return-value]
+
+
+def _dual_form(default_kind: str):
+    """Build a dual-form helper: ``h("name", fn)`` OR ``@h("name")``."""
+
+    def helper(
+        name: Any = None,
+        fn: Optional[Callable[..., Any]] = None,
+        *,
+        kind: str = default_kind,
+        on_result: Optional[Callable[[SpanHandle, Any], None]] = None,
+        extra_set: Optional[Callable[[SpanHandle], None]] = None,
+    ) -> Any:
+        # Case: h(fn) — bare decorator without parens; `name` is the fn.
+        if callable(name) and fn is None:
+            target = name
+            return _wrap_fn(
+                target,
+                name=target.__name__,
+                kind=kind,
+                on_result=on_result,
+                extra_set=extra_set,
+            )
+        # Case: h("name", fn) — helper form; wrap immediately.
+        if callable(fn):
+            return _wrap_fn(
+                fn,
+                name=name or fn.__name__,
+                kind=kind,
+                on_result=on_result,
+                extra_set=extra_set,
+            )
+        # Case: h("name") / h(name="name") / h() — return decorator.
+        def decorator(target: F) -> F:
+            return _wrap_fn(
+                target,
+                name=name or target.__name__,
+                kind=kind,
+                on_result=on_result,
+                extra_set=extra_set,
+            )
+
+        return decorator
+
+    return helper
+
+
+def tool(
+    name: Any = None,
+    fn: Optional[Callable[..., Any]] = None,
+    *,
+    kind: str = "tool",
+) -> Any:
+    """Wrap a function as a tool span — dual-form helper + decorator.
+
+    Helper form (uselemma-compatible)::
+
+        run_funnel_query = trodo.tool('run_funnel_query', run_funnel_query)
+        result = run_funnel_query(team_id=1, preset='day7')
+
+    Decorator form (backward-compatible)::
+
+        @trodo.tool()
+        def run_funnel_query(team_id, preset): ...
+
+        @trodo.tool(name='custom-name')
+        async def fetch(...): ...
+
+        @trodo.tool  # bare (no parens)
+        def do_thing(): ...
+    """
+    return _dual_form("tool")(name, fn, kind=kind)
+
+
+def trace(
+    name: Any = None,
+    fn: Optional[Callable[..., Any]] = None,
+) -> Any:
+    """Wrap a function as a generic span — dual-form helper + decorator.
+
+    Usage::
+
+        trodo.trace('prepare', prepare_fn)({'q': 'hi'})
+
+        @trodo.trace('step')
+        def step(): ...
+    """
+    return _dual_form("generic")(name, fn, kind="generic")
+
+
+def retrieval(
+    name: Any = None,
+    fn: Optional[Callable[..., Any]] = None,
+) -> Any:
+    """Wrap a retriever / vector search as a ``kind='retrieval'`` span."""
+    return _dual_form("retrieval")(name, fn, kind="retrieval")
+
+
+def _default_usage_extractor(result: Any) -> Tuple[Optional[int], Optional[int]]:
+    """Auto-detect OpenAI ``usage`` and Gemini ``usageMetadata`` shapes."""
+    if result is None:
+        return (None, None)
+    # OpenAI / OpenAI-compat: {"usage": {"prompt_tokens", "completion_tokens"}}
+    usage = None
+    if isinstance(result, dict):
+        usage = result.get("usage")
+    else:
+        usage = getattr(result, "usage", None)
+    if usage is not None:
+        get = (lambda k: usage.get(k)) if isinstance(usage, dict) else (lambda k: getattr(usage, k, None))
+        pt = get("prompt_tokens")
+        ct = get("completion_tokens")
+        if pt is not None or ct is not None:
+            return (
+                int(pt) if pt is not None else None,
+                int(ct) if ct is not None else None,
+            )
+        # Anthropic-style nested under `usage`:
+        it = get("input_tokens")
+        ot = get("output_tokens")
+        if it is not None or ot is not None:
+            return (
+                int(it) if it is not None else None,
+                int(ot) if ot is not None else None,
+            )
+    # Gemini: {"usageMetadata": {"promptTokenCount", "candidatesTokenCount"}}
+    if isinstance(result, dict):
+        meta = result.get("usageMetadata")
+        if isinstance(meta, dict):
+            pt = meta.get("promptTokenCount")
+            ct = meta.get("candidatesTokenCount")
+            return (
+                int(pt) if pt is not None else None,
+                int(ct) if ct is not None else None,
+            )
+    return (None, None)
+
+
+def llm(
+    name: Any = None,
+    fn: Optional[Callable[..., Any]] = None,
+    *,
+    model: Optional[str] = None,
+    provider: Optional[str] = None,
+    temperature: Optional[float] = None,
+    extract_usage: Optional[Callable[[Any], Tuple[Optional[int], Optional[int]]]] = None,
+) -> Any:
+    """Wrap an LLM call as a ``kind='llm'`` span with auto token extraction.
+
+    The helper records ``model``/``provider`` on entry; on return it inspects
+    the response for the common usage shapes (OpenAI ``usage.prompt_tokens``,
+    Anthropic ``usage.input_tokens``, Gemini ``usageMetadata.promptTokenCount``)
+    and records tokens. Pass ``extract_usage=lambda r: (in, out)`` to override.
+
+    Usage::
+
+        answer = trodo.llm(
+            'answer', call_openai, model='gpt-4o-mini', provider='openai',
+        )(messages)
+
+        @trodo.llm('plan', model='claude-haiku-4-5', provider='anthropic')
+        def plan(messages): ...
+    """
+    extractor = extract_usage or _default_usage_extractor
+
+    def _set_llm(s: SpanHandle) -> None:
+        if model or provider or temperature is not None:
+            s.set_llm(
+                model=model,
+                provider=provider,
+                temperature=temperature,
+            )
+
+    def _on_result(s: SpanHandle, result: Any) -> None:
+        try:
+            pt, ct = extractor(result)
+        except Exception:
+            pt, ct = (None, None)
+        if pt is not None or ct is not None:
+            s.set_llm(
+                model=model,
+                provider=provider,
+                input_tokens=pt,
+                output_tokens=ct,
+                temperature=temperature,
+            )
+
+    return _dual_form("llm")(
+        name, fn, kind="llm", extra_set=_set_llm, on_result=_on_result
+    )
+
+
+def track_llm_call(
+    *,
+    model: Optional[str] = None,
+    provider: Optional[str] = None,
+    input_tokens: Optional[int] = None,
+    output_tokens: Optional[int] = None,
+    prompt: Any = None,
+    completion: Any = None,
+    temperature: Optional[float] = None,
+    cost: Optional[float] = None,
+    name: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Record a one-shot LLM span for a raw-HTTP caller.
+
+    Opens and immediately closes a ``span(kind='llm')`` populated with the
+    model + token counts + prompt/completion. No-op outside an active run
+    context.
+
+    Usage:
+        resp = httpx.post(url, json=body).json()
+        trodo.track_llm_call(
+            model='gemini-2.5-flash',
+            provider='google',
+            input_tokens=resp['usageMetadata']['promptTokenCount'],
+            output_tokens=resp['usageMetadata']['candidatesTokenCount'],
+            prompt=body,
+            completion=resp,
+        )
+    """
+    if get_active_context() is None:
+        return
+    span_name = name or (f"llm.{provider}.{model}" if model else "llm")
+    with span_ctx(span_name, kind="llm", input=prompt, attributes=metadata) as s:
+        s.set_llm(
+            model=model,
+            provider=provider,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cost=cost,
+            temperature=temperature,
+        )
+        if completion is not None:
+            s.set_output(completion)
+
+
+# ---------------------------------------------------------------------------
+# FastAPI / Starlette middleware
+# ---------------------------------------------------------------------------
+
+_TRODO_RUN_HEADER = "x-trodo-run-id"
+_TRODO_PARENT_SPAN_HEADER = "x-trodo-parent-span-id"
+_TRODO_AGENT_HEADER = "x-trodo-agent-name"
+
+
+def fastapi_middleware(client: Any) -> Callable:
+    """Return a FastAPI/Starlette middleware that auto-joins remote runs.
+
+    If inbound request has X-Trodo-Run-Id, every span created while
+    handling the request nests under the caller's run. Otherwise the
+    handler runs with no active context (no-op tracking).
+
+    Usage:
+        from fastapi import FastAPI
+        import trodo
+        trodo.init(site_id='...')
+        app = FastAPI()
+        app.middleware('http')(trodo.fastapi_middleware(trodo._client))
+    """
+    processor = client._span_processor
+    site_id = client.site_id
+
+    async def _middleware(request: Any, call_next: Callable[[Any], Awaitable[Any]]):
+        headers = {k.lower(): v for k, v in request.headers.items()}
+        run_id = headers.get(_TRODO_RUN_HEADER)
+        parent_span_id = headers.get(_TRODO_PARENT_SPAN_HEADER)
+        name = (
+            headers.get(_TRODO_AGENT_HEADER)
+            or f"http.{request.method}.{request.url.path}"
+        )
+        if not run_id:
+            return await call_next(request)
+
+        input_payload = {
+            "method": request.method,
+            "path": request.url.path,
+        }
+        with join_run(
+            processor=processor,
+            team_site_id=site_id,
+            run_id=run_id,
+            parent_span_id=parent_span_id,
+            name=name,
+            kind="agent",
+            input=input_payload,
+        ) as s:
+            response = await call_next(request)
+            try:
+                s.set_output({"status_code": getattr(response, "status_code", None)})
+            except Exception:
+                pass
+            return response
+
+    return _middleware
+
+
+def propagation_headers() -> Dict[str, str]:
+    """Return HTTP headers that carry the current run/span context.
+
+    Use when making outbound HTTP calls to downstream services so they
+    can ``join_run`` instead of creating their own runs.
+    """
+    ctx = get_active_context()
+    if ctx is None:
+        return {}
+    headers: Dict[str, str] = {"X-Trodo-Run-Id": ctx.run_id}
+    if ctx.span_id:
+        headers["X-Trodo-Parent-Span-Id"] = ctx.span_id
+    return headers

trodo/otel/processor.py

@@ -0,0 +1,165 @@
+"""TrodoSpanProcessor — buffers spans and ships them to the Trodo ingest API.
+
+Mirrors the Node SDK processor: when a run finalises we POST /runs/ingest
+with the run and all pending spans for that run. Standalone spans (no run)
+flush asynchronously via append_spans.
+"""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass, asdict
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class TrodoRun:
+    run_id: str
+    agent_name: str
+    distinct_id: Optional[str] = None
+    conversation_id: Optional[str] = None
+    parent_run_id: Optional[str] = None
+    status: str = "ok"  # 'running' | 'ok' | 'error'
+    input: Optional[str] = None
+    output: Optional[str] = None
+    started_at: Optional[str] = None
+    ended_at: Optional[str] = None
+    duration_ms: Optional[int] = None
+    error_summary: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+    # Aggregates summed from child spans at finalisation.
+    total_tokens_in: Optional[int] = None
+    total_tokens_out: Optional[int] = None
+    total_cost: Optional[float] = None
+    span_count: Optional[int] = None
+    tool_count: Optional[int] = None
+    error_count: Optional[int] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {k: v for k, v in asdict(self).items() if v is not None}
+
+
+@dataclass
+class TrodoSpan:
+    span_id: str
+    run_id: str
+    parent_span_id: Optional[str]
+    kind: str = "generic"  # 'llm' | 'tool' | 'agent' | 'retrieval' | 'generic'
+    name: str = ""
+    status: str = "ok"
+    started_at: Optional[str] = None
+    ended_at: Optional[str] = None
+    duration_ms: Optional[int] = None
+    input: Optional[str] = None
+    output: Optional[str] = None
+    error_type: Optional[str] = None
+    error_message: Optional[str] = None
+    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
+    tool_name: Optional[str] = None
+    attributes: Optional[Dict[str, Any]] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {k: v for k, v in asdict(self).items() if v is not None}
+
+
+class TrodoSpanProcessor:
+    """Buffers spans per run_id; flushes to /runs/ingest when the run ends.
+
+    Spans emitted while a ``join_run`` context is active are forwarded via
+    ``append_spans`` immediately on each enqueue (because the owning service
+    is remote and won't finalise the run).
+    """
+
+    def __init__(
+        self,
+        http_client: Any,
+        flush_interval_s: float = 5.0,
+        max_spans_per_run: int = 1000,
+    ) -> None:
+        self._http = http_client
+        self._flush_interval = flush_interval_s
+        self._max_spans = max_spans_per_run
+        self._pending: Dict[str, List[TrodoSpan]] = {}
+        self._lock = threading.Lock()
+        self._shutdown = False
+        # run_ids that are "joined" (owned by a remote service) — their spans
+        # are flushed incrementally via append_spans, never batched for a
+        # local ingest_run call.
+        self._joined_runs: set[str] = set()
+
+    def mark_joined(self, run_id: str) -> None:
+        with self._lock:
+            self._joined_runs.add(run_id)
+
+    def unmark_joined(self, run_id: str) -> None:
+        with self._lock:
+            self._joined_runs.discard(run_id)
+            self._pending.pop(run_id, None)
+
+    def enqueue_span(self, span: TrodoSpan) -> None:
+        """Buffer a completed span under its run_id.
+
+        Joined runs flush immediately (append_spans) instead of batching.
+        """
+        with self._lock:
+            joined = span.run_id in self._joined_runs
+            if joined:
+                pass
+            else:
+                bucket = self._pending.setdefault(span.run_id, [])
+                if len(bucket) < self._max_spans:
+                    bucket.append(span)
+        if joined:
+            try:
+                self._http.post_spans_append(span.run_id, [span.to_dict()])
+            except Exception:
+                pass
+
+    def get_pending(self, run_id: str) -> List[TrodoSpan]:
+        with self._lock:
+            return list(self._pending.get(run_id, []))
+
+    def ingest_run(self, run: TrodoRun) -> None:
+        """Flush a run + all its buffered spans in one call."""
+        with self._lock:
+            spans = self._pending.pop(run.run_id, [])
+        payload: Dict[str, Any] = {"run": run.to_dict()}
+        if spans:
+            payload["spans"] = [s.to_dict() for s in spans]
+        try:
+            self._http.post_run_ingest(payload)
+        except Exception:
+            pass
+
+    def append_spans(self, run_id: str, spans: List[TrodoSpan]) -> None:
+        """Stream spans for a long-running or joined run without finalising."""
+        if not spans:
+            return
+        try:
+            self._http.post_spans_append(run_id, [s.to_dict() for s in spans])
+        except Exception:
+            pass
+
+    def force_flush(self) -> None:
+        """Flush any buffered spans that don't have a run finalise yet."""
+        with self._lock:
+            pending = list(self._pending.items())
+            self._pending.clear()
+        for run_id, spans in pending:
+            if not spans:
+                continue
+            try:
+                self._http.post_spans_append(
+                    run_id, [s.to_dict() for s in spans]
+                )
+            except Exception:
+                pass
+
+    def shutdown(self) -> None:
+        self._shutdown = True
+        self.force_flush()