agentos-python 0.1.0__py3-none-any.whl

agentos/decorators.py

@@ -0,0 +1,372 @@
+"""Decorators for automatic tracing and event capture.
+
+Provides ``@trace_agent``, ``@trace_llm_call``, ``@trace_tool``,
+and Langfuse-compatible ``@observe``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import time
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from agentos.tracing import (
+    get_current_context,
+    span,
+    trace,
+)
+
+logger = logging.getLogger("agentos.decorators")
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _get_global_client() -> Any:
+    """Lazily import the global client to avoid circular imports."""
+    from agentos.client import get_client
+
+    return get_client()
+
+
+def _is_async(fn: Callable[..., Any]) -> bool:
+    return asyncio.iscoroutinefunction(fn)
+
+
+# ---------------------------------------------------------------------------
+# @observe — Langfuse-compatible decorator
+# ---------------------------------------------------------------------------
+
+
+def observe(
+    *,
+    name: str | None = None,
+    as_type: str | None = None,
+) -> Callable[[F], F]:
+    """Langfuse-compatible decorator for automatic tracing.
+
+    Creates a span (or trace if none exists) around the decorated function.
+    Nesting is automatic via contextvars.
+
+    Args:
+        name: Span name. Defaults to the function name.
+        as_type: If ``"generation"``, emits an ``agent.llm_call`` event.
+            Otherwise emits timing as a span.
+    """
+
+    def decorator(fn: F) -> F:
+        span_name = name or fn.__name__
+
+        if _is_async(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                parent = get_current_context()
+                ctx_mgr = trace(agent_id=span_name) if parent is None else span()
+                start = time.monotonic()
+                error_info = None
+
+                with ctx_mgr as ctx:
+                    try:
+                        result = await fn(*args, **kwargs)
+                        return result
+                    except Exception as exc:
+                        error_info = {"type": type(exc).__name__, "message": str(exc)}
+                        raise
+                    finally:
+                        duration_ms = (time.monotonic() - start) * 1000
+                        client = _get_global_client()
+                        if client is not None and as_type == "generation":
+                            # Pull any updates from trace_context helper
+                            updates = ctx._observation_updates
+                            client.llm_call(
+                                ctx.agent_id,
+                                model=updates.get("model", "unknown"),
+                                system=updates.get("system", "unknown"),
+                                duration_ms=duration_ms,
+                                error=error_info,
+                                trace_id=ctx.trace_id,
+                                span_id=ctx.span_id,
+                                parent_span_id=ctx.parent_span_id,
+                                **{
+                                    k: v for k, v in updates.items() if k not in ("model", "system")
+                                },
+                            )
+
+            return async_wrapper  # type: ignore[return-value]
+
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                parent = get_current_context()
+                ctx_mgr = trace(agent_id=span_name) if parent is None else span()
+                start = time.monotonic()
+                error_info = None
+
+                with ctx_mgr as ctx:
+                    try:
+                        result = fn(*args, **kwargs)
+                        return result
+                    except Exception as exc:
+                        error_info = {"type": type(exc).__name__, "message": str(exc)}
+                        raise
+                    finally:
+                        duration_ms = (time.monotonic() - start) * 1000
+                        client = _get_global_client()
+                        if client is not None and as_type == "generation":
+                            updates = ctx._observation_updates
+                            client.llm_call(
+                                ctx.agent_id,
+                                model=updates.get("model", "unknown"),
+                                system=updates.get("system", "unknown"),
+                                duration_ms=duration_ms,
+                                error=error_info,
+                                trace_id=ctx.trace_id,
+                                span_id=ctx.span_id,
+                                parent_span_id=ctx.parent_span_id,
+                                **{
+                                    k: v for k, v in updates.items() if k not in ("model", "system")
+                                },
+                            )
+
+            return sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+# ---------------------------------------------------------------------------
+# trace_context — Langfuse-compatible context helper
+# ---------------------------------------------------------------------------
+
+
+class trace_context:
+    """Module-level helper for updating the current trace/observation.
+
+    Mirrors ``langfuse_context`` from the Langfuse SDK.
+    """
+
+    @staticmethod
+    def update_current_observation(
+        *,
+        model: str | None = None,
+        system: str | None = None,
+        usage: dict[str, Any] | None = None,
+        input: Any | None = None,
+        output: Any | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Update the current span/generation with additional data."""
+        ctx = get_current_context()
+        if ctx is None:
+            return
+        if model:
+            ctx._observation_updates["model"] = model
+        if system:
+            ctx._observation_updates["system"] = system
+        if usage:
+            if "input" in usage:
+                ctx._observation_updates["input_tokens"] = usage["input"]
+            if "output" in usage:
+                ctx._observation_updates["output_tokens"] = usage["output"]
+            if "total" in usage:
+                ctx._observation_updates["total_tokens"] = usage["total"]
+        if input is not None:
+            ctx._observation_updates["input"] = input
+        if output is not None:
+            ctx._observation_updates["output"] = output
+        ctx._observation_updates.update(kwargs)
+
+    @staticmethod
+    def update_current_trace(
+        *,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """Update the current trace metadata."""
+        ctx = get_current_context()
+        if ctx is None:
+            return
+        if user_id:
+            ctx._trace_updates["user_id"] = user_id
+        if session_id:
+            ctx._trace_updates["session_id"] = session_id
+        if metadata:
+            ctx._trace_updates["metadata"] = metadata
+
+    @staticmethod
+    def score_current_trace(*, name: str, value: float, comment: str | None = None) -> None:
+        """Attach a score to the current trace."""
+        ctx = get_current_context()
+        if ctx is None:
+            return
+        score = {"name": name, "value": value}
+        if comment:
+            score["comment"] = comment
+        ctx._scores.append(score)
+
+        client = _get_global_client()
+        if client is not None:
+            client.eval(
+                ctx.agent_id,
+                eval_name=name,
+                score=value,
+                trace_id=ctx.trace_id,
+                span_id=ctx.span_id,
+            )
+
+    @staticmethod
+    def score_current_observation(*, name: str, value: float, comment: str | None = None) -> None:
+        """Attach a score to the current observation/span."""
+        trace_context.score_current_trace(name=name, value=value, comment=comment)
+
+
+# ---------------------------------------------------------------------------
+# Native decorators
+# ---------------------------------------------------------------------------
+
+
+def trace_agent(
+    agent_id: str,
+    *,
+    task_run_id: str | None = None,
+    user_id: str | None = None,
+    session_id: str | None = None,
+) -> Callable[[F], F]:
+    """Decorator that wraps a function in a trace context.
+
+    Usage::
+
+        @trace_agent("my-agent")
+        def run_agent(query: str) -> str:
+            ...
+    """
+
+    def decorator(fn: F) -> F:
+        if _is_async(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with trace(
+                    agent_id=agent_id,
+                    task_run_id=task_run_id,
+                    user_id=user_id,
+                    session_id=session_id,
+                ):
+                    return await fn(*args, **kwargs)
+
+            return async_wrapper  # type: ignore[return-value]
+
+        @functools.wraps(fn)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with trace(
+                agent_id=agent_id,
+                task_run_id=task_run_id,
+                user_id=user_id,
+                session_id=session_id,
+            ):
+                return fn(*args, **kwargs)
+
+        return sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def trace_tool(
+    tool_name: str,
+) -> Callable[[F], F]:
+    """Decorator that captures tool call events automatically.
+
+    Captures input (function args), output (return value), duration, and errors.
+
+    Usage::
+
+        @trace_tool("web-search")
+        def search(query: str) -> list:
+            ...
+    """
+
+    def decorator(fn: F) -> F:
+        if _is_async(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with span() as ctx:
+                    start = time.monotonic()
+                    try:
+                        result = await fn(*args, **kwargs)
+                        duration_ms = (time.monotonic() - start) * 1000
+                        client = _get_global_client()
+                        if client:
+                            client.tool_call(
+                                ctx.agent_id,
+                                tool_name=tool_name,
+                                status="success",
+                                input=kwargs or None,
+                                duration_ms=duration_ms,
+                                trace_id=ctx.trace_id,
+                                span_id=ctx.span_id,
+                                parent_span_id=ctx.parent_span_id,
+                            )
+                        return result
+                    except Exception as exc:
+                        duration_ms = (time.monotonic() - start) * 1000
+                        client = _get_global_client()
+                        if client:
+                            client.tool_call(
+                                ctx.agent_id,
+                                tool_name=tool_name,
+                                status="error",
+                                duration_ms=duration_ms,
+                                error={"type": type(exc).__name__, "message": str(exc)},
+                                trace_id=ctx.trace_id,
+                                span_id=ctx.span_id,
+                                parent_span_id=ctx.parent_span_id,
+                            )
+                        raise
+
+            return async_wrapper  # type: ignore[return-value]
+
+        @functools.wraps(fn)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with span() as ctx:
+                start = time.monotonic()
+                try:
+                    result = fn(*args, **kwargs)
+                    duration_ms = (time.monotonic() - start) * 1000
+                    client = _get_global_client()
+                    if client:
+                        client.tool_call(
+                            ctx.agent_id,
+                            tool_name=tool_name,
+                            status="success",
+                            input=kwargs or None,
+                            duration_ms=duration_ms,
+                            trace_id=ctx.trace_id,
+                            span_id=ctx.span_id,
+                            parent_span_id=ctx.parent_span_id,
+                        )
+                    return result
+                except Exception as exc:
+                    duration_ms = (time.monotonic() - start) * 1000
+                    client = _get_global_client()
+                    if client:
+                        client.tool_call(
+                            ctx.agent_id,
+                            tool_name=tool_name,
+                            status="error",
+                            duration_ms=duration_ms,
+                            error={"type": type(exc).__name__, "message": str(exc)},
+                            trace_id=ctx.trace_id,
+                            span_id=ctx.span_id,
+                            parent_span_id=ctx.parent_span_id,
+                        )
+                    raise
+
+        return sync_wrapper  # type: ignore[return-value]
+
+    return decorator