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

trodo/client.py

@@ -2,24 +2,37 @@
 
 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,
+    start_run as start_run_fn,
+    end_run as end_run_fn,
+    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,
     ResetResult,
     WalletAddressResult,
-    AgentCallProps,
-    ToolUseProps,
-    AgentResponseProps,
-    AgentErrorProps,
     FeedbackProps,
 )
 
@@ -37,6 +50,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")
@@ -73,6 +87,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()
     # --------------------------------------------------------------------------
@@ -195,124 +213,234 @@ 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 Analytics
+    # Agent Runs — Lemma-style seamless wrapping
     # --------------------------------------------------------------------------
 
-    def _build_agent_base(
+    def wrap_agent(
         self,
-        agent_id: str,
-        conversation_id: str,
-        message_id: str,
-        event_type: str,
+        agent_name: str,
         distinct_id: Optional[str] = None,
-        timestamp: Optional[str] = None,
-    ) -> Dict[str, Any]:
-        payload: Dict[str, Any] = {
-            "event_type": event_type,
-            "agent_id": agent_id,
-            "conversation_id": conversation_id,
-            "message_id": message_id,
-        }
-        if distinct_id is not None:
-            payload["distinct_id"] = distinct_id
-        if timestamp is not None:
-            payload["timestamp"] = timestamp
-        return payload
-
-    def track_agent_call(self, props: AgentCallProps) -> None:
-        """Track an LLM invocation / inbound message."""
-        payload = self._build_agent_base(
-            props.agent_id, props.conversation_id, props.message_id,
-            "agent_call", props.distinct_id, props.timestamp,
+        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,
         )
-        if props.prompt is not None:
-            payload["prompt"] = props.prompt
-        if props.model is not None:
-            payload["model"] = props.model
-        if props.temperature is not None:
-            payload["temperature"] = props.temperature
-        if props.system_prompt_version is not None:
-            payload["system_prompt_version"] = props.system_prompt_version
-        if props.provider is not None:
-            payload["provider"] = props.provider
-        if props.properties:
-            payload.update(props.properties)
-        self._http.post_agent_event(payload)
-
-    def track_tool_use(self, props: ToolUseProps) -> None:
-        """Track a tool invocation within an agent turn."""
-        payload = self._build_agent_base(
-            props.agent_id, props.conversation_id, props.message_id,
-            "tool_use", props.distinct_id, props.timestamp,
+
+    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 start_run(
+        self,
+        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 outside a context manager. Returns the run_id.
+
+        Use ``end_run`` to finalise, ``join_run`` from any process to add spans.
+        """
+        return start_run_fn(
+            processor=self._span_processor,
+            agent_name=agent_name,
+            run_id=run_id,
+            distinct_id=distinct_id,
+            conversation_id=conversation_id,
+            parent_run_id=parent_run_id,
+            metadata=metadata,
+            input=input,
+        )
+
+    def end_run(
+        self,
+        run_id: str,
+        *,
+        output: Any = None,
+        status: str = "ok",
+        error_summary: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Finalise a Run previously opened by :meth:`start_run`."""
+        end_run_fn(
+            run_id,
+            processor=self._span_processor,
+            output=output,
+            status=status,
+            error_summary=error_summary,
+            metadata=metadata,
         )
-        payload["tool_name"] = props.tool_name
-        if props.input is not None:
-            payload["input"] = props.input
-        if props.output is not None:
-            payload["output"] = props.output
-        if props.latency_ms is not None:
-            payload["latency_ms"] = props.latency_ms
-        if props.status is not None:
-            payload["status"] = props.status
-        if props.properties:
-            payload.update(props.properties)
-        self._http.post_agent_event(payload)
-
-    def track_agent_response(self, props: AgentResponseProps) -> None:
-        """Track an LLM response / completion."""
-        payload = self._build_agent_base(
-            props.agent_id, props.conversation_id, props.message_id,
-            "agent_response", props.distinct_id, props.timestamp,
+
+    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,
         )
-        if props.output is not None:
-            payload["output"] = props.output
-        if props.model is not None:
-            payload["model"] = props.model
-        if props.completion_tokens is not None:
-            payload["completion_tokens"] = props.completion_tokens
-        if props.prompt_tokens is not None:
-            payload["prompt_tokens"] = props.prompt_tokens
-        if props.total_tokens is not None:
-            payload["total_tokens"] = props.total_tokens
-        if props.finish_reason is not None:
-            payload["finish_reason"] = props.finish_reason
-        if props.properties:
-            payload.update(props.properties)
-        self._http.post_agent_event(payload)
-
-    def track_agent_error(self, props: AgentErrorProps) -> None:
-        """Track an error during an agent turn."""
-        payload = self._build_agent_base(
-            props.agent_id, props.conversation_id, props.message_id,
-            "agent_error", props.distinct_id, props.timestamp,
+
+    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,
         )
-        if props.error_type is not None:
-            payload["error_type"] = props.error_type
-        if props.error_message is not None:
-            payload["error_message"] = props.error_message
-        if props.failed_tool is not None:
-            payload["failed_tool"] = props.failed_tool
-        if props.traceback is not None:
-            payload["traceback"] = props.traceback
-        if props.properties:
-            payload.update(props.properties)
-        self._http.post_agent_event(payload)
-
-    def track_feedback(self, props: FeedbackProps) -> None:
-        """Track a user feedback reaction on an agent response."""
-        payload = self._build_agent_base(
-            props.agent_id, props.conversation_id, props.message_id,
-            "feedback", props.distinct_id, props.timestamp,
+
+    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,
         )
-        payload["feedback"] = props.feedback
-        if props.properties:
-            payload.update(props.properties)
-        self._http.post_agent_event(payload)

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",
+]