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

trodo/client.py

@@ -2,21 +2,36 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, Union
+import functools
+from typing import Any, Callable, Dict, List, Optional, Union
 
 from .api.http_client import HttpClient
 from .session.session_manager import SessionManager
-from .managers.people_manager import PeopleManager
-from .managers.group_manager import GroupManager, GroupProfile
+from .managers.group_manager import GroupProfile
 from .queue.event_queue import EventQueue
 from .queue.batch_flusher import BatchFlusher
 from .auto.auto_event_manager import AutoEventManager
 from .user_context import UserContext
+from .otel.processor import TrodoSpanProcessor
+from .otel.wrap_agent import (
+    wrap_agent as wrap_agent_ctx,
+    span as span_ctx,
+    join_run as join_run_ctx,
+    current_run_id as _current_run_id,
+    current_span_id as _current_span_id,
+)
+from .otel.auto_instrument import enable_auto_instrument
+from .otel.helpers import (
+    tool as tool_decorator,
+    track_llm_call as track_llm_call_fn,
+    fastapi_middleware as fastapi_middleware_fn,
+    propagation_headers as propagation_headers_fn,
+)
 from .types import (
     ApiResult,
-    IdentifyResult,
     ResetResult,
     WalletAddressResult,
+    FeedbackProps,
 )
 
 
@@ -33,6 +48,7 @@ class TrodoClient:
         auto_events: bool = False,
         on_error: Optional[Any] = None,
         debug: bool = False,
+        auto_instrument: bool = True,
     ) -> None:
         if not site_id:
             raise ValueError("trodo-python: site_id is required")
@@ -69,6 +85,10 @@ class TrodoClient:
 
         self._user_context_cache: Dict[str, UserContext] = {}
 
+        self._span_processor = TrodoSpanProcessor(http_client=self._http)
+        if auto_instrument:
+            enable_auto_instrument(self._span_processor)
+
     # --------------------------------------------------------------------------
     # Primary pattern: for_user()
     # --------------------------------------------------------------------------
@@ -107,8 +127,12 @@ class TrodoClient:
     ) -> None:
         self.for_user(distinct_id).track(event_name, properties, category)
 
-    def identify(self, distinct_id: str, identify_id: str) -> IdentifyResult:
-        return self.for_user(distinct_id).identify(identify_id)
+    def identify(self, identify_id: str, session_id: Optional[str] = None) -> "UserContext":
+        if identify_id in self._user_context_cache:
+            return self._user_context_cache[identify_id]
+        ctx = self.for_user(identify_id, session_id)
+        ctx.identify(identify_id)
+        return ctx
 
     def wallet_address(self, distinct_id: str, wallet_addr: str) -> WalletAddressResult:
         return self.for_user(distinct_id).wallet_address(wallet_addr)
@@ -187,9 +211,189 @@ class TrodoClient:
     def flush(self) -> None:
         if self._batch_flusher:
             self._batch_flusher.flush()
+        self._span_processor.force_flush()
 
     def shutdown(self) -> None:
         self._auto_event_manager.disable()
         if self._batch_flusher:
             self._batch_flusher.stop()
             self._batch_flusher.flush()
+        self._span_processor.shutdown()
+
+    # --------------------------------------------------------------------------
+    # Agent Runs — Lemma-style seamless wrapping
+    # --------------------------------------------------------------------------
+
+    def wrap_agent(
+        self,
+        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,
+    ) -> wrap_agent_ctx:
+        """Context manager that captures the wrapped block as an agent run.
+
+        Every OTel-instrumented call (Anthropic, OpenAI, LangChain…) made
+        inside the ``with`` is auto-captured as a nested span.
+        """
+        return wrap_agent_ctx(
+            processor=self._span_processor,
+            team_site_id=self.site_id,
+            agent_name=agent_name,
+            distinct_id=distinct_id,
+            conversation_id=conversation_id,
+            parent_run_id=parent_run_id,
+            metadata=metadata,
+        )
+
+    def agent(
+        self,
+        agent_name: str,
+        *,
+        distinct_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Decorator form of wrap_agent."""
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            @functools.wraps(fn)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                with self.wrap_agent(
+                    agent_name,
+                    distinct_id=distinct_id,
+                    conversation_id=conversation_id,
+                    metadata=metadata,
+                ) as run:
+                    result = fn(*args, **kwargs)
+                    run.set_output(result)
+                    return result
+            return wrapper
+
+        return decorator
+
+    def span(
+        self,
+        name: str,
+        *,
+        kind: str = "generic",
+        input: Any = None,
+        attributes: Optional[Dict[str, Any]] = None,
+    ) -> span_ctx:
+        """Create a nested span inside the current run."""
+        return span_ctx(name=name, kind=kind, input=input, attributes=attributes)
+
+    def join_run(
+        self,
+        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,
+    ) -> join_run_ctx:
+        """Open a span on an existing run owned by a remote service.
+
+        Use this in downstream microservices that receive an inbound
+        request carrying ``X-Trodo-Run-Id`` / ``X-Trodo-Parent-Span-Id``.
+        Prefer :meth:`fastapi_middleware` to automate this.
+        """
+        return join_run_ctx(
+            processor=self._span_processor,
+            team_site_id=self.site_id,
+            run_id=run_id,
+            parent_span_id=parent_span_id,
+            name=name,
+            kind=kind,
+            input=input,
+            attributes=attributes,
+        )
+
+    def tool(
+        self,
+        name: Optional[str] = None,
+        *,
+        kind: str = "tool",
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Decorator that wraps a function call as a tool span."""
+        return tool_decorator(name=name, kind=kind)
+
+    def track_llm_call(
+        self,
+        *,
+        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 an LLM span for a raw-HTTP caller (no OTel adapter)."""
+        track_llm_call_fn(
+            model=model,
+            provider=provider,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            prompt=prompt,
+            completion=completion,
+            temperature=temperature,
+            cost=cost,
+            name=name,
+            metadata=metadata,
+        )
+
+    def fastapi_middleware(self) -> Callable:
+        """Return a FastAPI/Starlette middleware that auto-joins runs."""
+        return fastapi_middleware_fn(self)
+
+    def propagation_headers(self) -> Dict[str, str]:
+        """Return outbound HTTP headers carrying the current run/span id."""
+        return propagation_headers_fn()
+
+    def current_run_id(self) -> Optional[str]:
+        return _current_run_id()
+
+    def current_span_id(self) -> Optional[str]:
+        return _current_span_id()
+
+    def feedback(
+        self,
+        run_id: str,
+        satisfaction: Optional[str] = None,
+        rating: Optional[float] = None,
+        comment: Optional[str] = None,
+        feedback: Optional[str] = None,
+        distinct_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> ApiResult:
+        """Attach feedback to a completed run."""
+        if satisfaction is None and rating is None and not (comment or feedback):
+            raise ValueError(
+                "At least one of 'satisfaction', 'rating', or 'comment' is required"
+            )
+        payload: Dict[str, Any] = {
+            "satisfaction": satisfaction,
+            "rating": rating,
+            "comment": comment if comment is not None else feedback,
+            "distinct_id": distinct_id,
+            "attributes": metadata or {},
+        }
+        return self._http.post_run_feedback(run_id, payload)
+
+    def feedback_props(self, run_id: str, props: FeedbackProps) -> ApiResult:
+        """Convenience — pass a FeedbackProps dataclass."""
+        return self.feedback(
+            run_id,
+            satisfaction=props.satisfaction,
+            rating=props.rating,
+            comment=props.comment,
+            feedback=props.feedback,
+            distinct_id=props.distinct_id,
+            metadata=props.metadata,
+        )

trodo/otel/__init__.py

@@ -0,0 +1,21 @@
+"""OpenTelemetry-compatible run/span pipeline for trodo-python."""
+
+from .context import ActiveSpanContext, get_active_context, run_with_context
+from .processor import TrodoSpanProcessor, TrodoRun, TrodoSpan
+from .wrap_agent import wrap_agent, span, SpanHandle, RunHandle
+from .auto_instrument import enable_auto_instrument, otel_span_to_trodo_span
+
+__all__ = [
+    "ActiveSpanContext",
+    "get_active_context",
+    "run_with_context",
+    "TrodoSpanProcessor",
+    "TrodoRun",
+    "TrodoSpan",
+    "wrap_agent",
+    "span",
+    "SpanHandle",
+    "RunHandle",
+    "enable_auto_instrument",
+    "otel_span_to_trodo_span",
+]

trodo/otel/auto_instrument.py

@@ -0,0 +1,281 @@
+"""Optional: register OpenTelemetry auto-instrumentation that feeds spans
+into our TrodoSpanProcessor. Call ``enable_auto_instrument()`` once at
+startup to pick up Anthropic / OpenAI / LangChain / LlamaIndex / Google
+Generative AI / Vertex AI and others — without any more code.
+
+Each underlying OTel package is an optional dependency — we import them
+lazily and skip anything that isn't installed.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, Iterable, List, Optional
+
+from .context import get_active_context
+from .processor import TrodoSpan, TrodoSpanProcessor
+
+
+def _hr_to_iso(nanos: Optional[int]) -> Optional[str]:
+    if nanos is None:
+        return None
+    return datetime.fromtimestamp(nanos / 1e9, tz=timezone.utc).isoformat().replace("+00:00", "Z")
+
+
+def _infer_kind(attrs: Dict[str, Any]) -> str:
+    if not attrs:
+        return "generic"
+    if attrs.get("gen_ai.tool.name"):
+        return "tool"
+    if attrs.get("gen_ai.operation.name") or attrs.get("gen_ai.request.model"):
+        return "llm"
+    if attrs.get("db.system") or attrs.get("retrieval.query"):
+        return "retrieval"
+    return "generic"
+
+
+def otel_span_to_trodo_span(otel_span: Any) -> Optional[TrodoSpan]:
+    """Translate an OTel ReadableSpan to our TrodoSpan using GenAI semconv."""
+    ctx = get_active_context()
+    if ctx is None:
+        return None  # span emitted outside of wrap_agent — drop
+    try:
+        span_ctx = (
+            otel_span.get_span_context()
+            if callable(getattr(otel_span, "get_span_context", None))
+            else otel_span.context
+        )
+    except Exception:
+        return None
+
+    span_id_int = getattr(span_ctx, "span_id", None)
+    if not span_id_int:
+        return None
+    span_id = f"{span_id_int:016x}" if isinstance(span_id_int, int) else str(span_id_int)
+
+    parent = getattr(otel_span, "parent", None)
+    parent_span_id: Optional[str] = None
+    if parent is not None:
+        pid = getattr(parent, "span_id", None)
+        parent_span_id = f"{pid:016x}" if isinstance(pid, int) else str(pid) if pid else None
+    if parent_span_id is None:
+        parent_span_id = ctx.span_id
+
+    attrs = dict(getattr(otel_span, "attributes", {}) or {})
+    kind = _infer_kind(attrs)
+
+    start_time = getattr(otel_span, "start_time", None)
+    end_time = getattr(otel_span, "end_time", None)
+    started_at = _hr_to_iso(start_time) or datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+    ended_at = _hr_to_iso(end_time)
+    duration_ms = None
+    if start_time and end_time:
+        duration_ms = max(0, int((end_time - start_time) / 1e6))
+
+    status = getattr(otel_span, "status", None)
+    status_code = getattr(status, "status_code", None)
+    ok = "ok"
+    try:
+        from opentelemetry.trace import StatusCode
+
+        ok = "error" if status_code == StatusCode.ERROR else "ok"
+    except Exception:
+        ok = "error" if (status_code and str(status_code).endswith("ERROR")) else "ok"
+
+    # Accept both stable and experimental GenAI semconv keys.
+    in_toks = (
+        attrs.get("gen_ai.usage.input_tokens")
+        or attrs.get("gen_ai.usage.prompt_tokens")
+        or attrs.get("llm.usage.prompt_tokens")
+    )
+    out_toks = (
+        attrs.get("gen_ai.usage.output_tokens")
+        or attrs.get("gen_ai.usage.completion_tokens")
+        or attrs.get("llm.usage.completion_tokens")
+    )
+    model = (
+        attrs.get("gen_ai.request.model")
+        or attrs.get("gen_ai.response.model")
+        or attrs.get("llm.request.model")
+    )
+    provider = attrs.get("gen_ai.system") or attrs.get("llm.vendor")
+    prompt = attrs.get("gen_ai.prompt") or attrs.get("llm.prompts")
+    completion = attrs.get("gen_ai.completion") or attrs.get("llm.completion")
+
+    return TrodoSpan(
+        span_id=span_id,
+        run_id=ctx.run_id,
+        parent_span_id=parent_span_id,
+        kind=kind,
+        name=getattr(otel_span, "name", kind),
+        status=ok,
+        started_at=started_at,
+        ended_at=ended_at,
+        duration_ms=duration_ms,
+        model=model,
+        provider=provider,
+        input_tokens=int(in_toks) if in_toks is not None else None,
+        output_tokens=int(out_toks) if out_toks is not None else None,
+        temperature=attrs.get("gen_ai.request.temperature"),
+        tool_name=attrs.get("gen_ai.tool.name"),
+        input=prompt,
+        output=completion,
+        attributes=attrs or None,
+    )
+
+
+class _OtelAdapter:
+    """Adapter matching opentelemetry.sdk.trace.SpanProcessor interface."""
+
+    def __init__(self, processor: TrodoSpanProcessor) -> None:
+        self._processor = processor
+
+    def on_start(self, span: Any, parent_context: Any = None) -> None:
+        pass
+
+    def on_end(self, span: Any) -> None:
+        trodo_span = otel_span_to_trodo_span(span)
+        if trodo_span is not None:
+            self._processor.enqueue_span(trodo_span)
+
+    def shutdown(self) -> None:
+        self._processor.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30_000) -> bool:
+        self._processor.force_flush()
+        return True
+
+
+_INSTRUMENTORS: List[tuple[str, Callable[[], Any]]] = []
+
+
+def _register_instrumentors() -> None:
+    """Build the list of known instrumentor factories.
+
+    Each entry is ``(framework_id, factory)`` where ``factory`` performs a
+    lazy import and returns an ``instrument()`` call. Failures are swallowed
+    so missing optional deps never break user code.
+    """
+    global _INSTRUMENTORS
+    if _INSTRUMENTORS:
+        return
+
+    def _anthropic() -> Any:
+        from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor  # type: ignore
+
+        AnthropicInstrumentor().instrument()
+
+    def _openai() -> Any:
+        from opentelemetry.instrumentation.openai import OpenAIInstrumentor  # type: ignore
+
+        OpenAIInstrumentor().instrument()
+
+    def _openai_v2() -> Any:
+        from opentelemetry.instrumentation.openai_v2 import OpenAIInstrumentor  # type: ignore
+
+        OpenAIInstrumentor().instrument()
+
+    def _langchain() -> Any:
+        from opentelemetry.instrumentation.langchain import LangChainInstrumentor  # type: ignore
+
+        LangChainInstrumentor().instrument()
+
+    def _llama_index() -> Any:
+        from opentelemetry.instrumentation.llama_index import LlamaIndexInstrumentor  # type: ignore
+
+        LlamaIndexInstrumentor().instrument()
+
+    def _google_generativeai() -> Any:
+        from opentelemetry.instrumentation.google_generativeai import (  # type: ignore
+            GoogleGenerativeAIInstrumentor,
+        )
+
+        GoogleGenerativeAIInstrumentor().instrument()
+
+    def _vertexai() -> Any:
+        from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor  # type: ignore
+
+        VertexAIInstrumentor().instrument()
+
+    def _bedrock() -> Any:
+        from opentelemetry.instrumentation.bedrock import BedrockInstrumentor  # type: ignore
+
+        BedrockInstrumentor().instrument()
+
+    def _cohere() -> Any:
+        from opentelemetry.instrumentation.cohere import CohereInstrumentor  # type: ignore
+
+        CohereInstrumentor().instrument()
+
+    def _mistralai() -> Any:
+        from opentelemetry.instrumentation.mistralai import MistralAiInstrumentor  # type: ignore
+
+        MistralAiInstrumentor().instrument()
+
+    def _haystack() -> Any:
+        from opentelemetry.instrumentation.haystack import HaystackInstrumentor  # type: ignore
+
+        HaystackInstrumentor().instrument()
+
+    def _httpx() -> Any:
+        from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor  # type: ignore
+
+        HTTPXClientInstrumentor().instrument()
+
+    def _requests() -> Any:
+        from opentelemetry.instrumentation.requests import RequestsInstrumentor  # type: ignore
+
+        RequestsInstrumentor().instrument()
+
+    _INSTRUMENTORS = [
+        ("anthropic", _anthropic),
+        ("openai", _openai),
+        ("openai_v2", _openai_v2),
+        ("langchain", _langchain),
+        ("llama_index", _llama_index),
+        ("google_generativeai", _google_generativeai),
+        ("vertexai", _vertexai),
+        ("bedrock", _bedrock),
+        ("cohere", _cohere),
+        ("mistralai", _mistralai),
+        ("haystack", _haystack),
+        ("httpx", _httpx),
+        ("requests", _requests),
+    ]
+
+
+def enable_auto_instrument(
+    processor: TrodoSpanProcessor,
+    disable: Optional[Iterable[str]] = None,
+) -> List[str]:
+    """Register OTel auto-instrumentations for installed packages.
+
+    Returns the list of framework ids that were actually instrumented.
+    Skipped silently if opentelemetry-sdk or the per-framework
+    instrumentation package isn't installed.
+    """
+    disabled = set(disable or [])
+    try:
+        from opentelemetry import trace  # type: ignore
+        from opentelemetry.sdk.trace import TracerProvider  # type: ignore
+    except Exception:
+        return []
+
+    provider = trace.get_tracer_provider()
+    if not isinstance(provider, TracerProvider):
+        provider = TracerProvider()
+        trace.set_tracer_provider(provider)
+
+    provider.add_span_processor(_OtelAdapter(processor))
+
+    _register_instrumentors()
+    active: List[str] = []
+    for name, register in _INSTRUMENTORS:
+        if name in disabled:
+            continue
+        try:
+            register()
+            active.append(name)
+        except Exception:
+            continue
+    return active

trodo/otel/context.py

@@ -0,0 +1,44 @@
+"""Active run/span context using contextvars — auto-propagates across awaits
+and threading.local-style sync code.
+"""
+
+from __future__ import annotations
+
+import contextvars
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class ActiveSpanContext:
+    run_id: str
+    span_id: str
+    parent_span_id: Optional[str]
+    team_site_id: str
+    processor: object  # TrodoSpanProcessor — avoid circular import
+
+
+_active: contextvars.ContextVar[Optional[ActiveSpanContext]] = contextvars.ContextVar(
+    "trodo_active_span", default=None,
+)
+
+
+def get_active_context() -> Optional[ActiveSpanContext]:
+    return _active.get()
+
+
+class run_with_context:
+    """Context manager that sets the active span context for its duration."""
+
+    def __init__(self, ctx: ActiveSpanContext) -> None:
+        self._ctx = ctx
+        self._token: Optional[contextvars.Token] = None
+
+    def __enter__(self) -> ActiveSpanContext:
+        self._token = _active.set(self._ctx)
+        return self._ctx
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        if self._token is not None:
+            _active.reset(self._token)
+            self._token = None