synth-ai 0.2.9.dev11__py3-none-any.whl → 0.4.1__py3-none-any.whl

synth_ai/core/tracing_v3/decorators.py

@@ -0,0 +1,401 @@
+"""Async-aware decorators for tracing v3.
+
+This module provides decorators and context management utilities for the tracing
+system. The key innovation is the use of asyncio's ContextVar for propagating
+tracing context across async boundaries without explicit parameter passing.
+
+Context Variables:
+-----------------
+We use ContextVar instead of threading.local because:
+1. ContextVars work correctly with asyncio tasks and coroutines
+2. They automatically propagate across async boundaries
+3. They're isolated between concurrent tasks
+4. They work with sync code too (unlike pure async solutions)
+
+Key Decorators:
+--------------
+- @with_session: Ensures a session is active before function execution
+- @trace_llm_call: Automatically traces LLM API calls with metrics
+- @trace_method: Generic method tracing for any class method
+
+The decorators support both sync and async functions where appropriate,
+though async is preferred for consistency with the rest of the system.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import functools
+import time
+from collections.abc import Awaitable, Callable, Mapping
+from contextvars import Token
+from typing import Any, TypeVar, cast, overload
+
+from .abstractions import LMCAISEvent, TimeRecord
+from .utils import calculate_cost, detect_provider
+
+# Context variables for session and turn tracking
+# These variables automatically propagate across async call boundaries,
+# allowing deeply nested code to access tracing context without explicit passing
+_session_id_ctx: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "session_id"
+)
+_turn_number_ctx: contextvars.ContextVar[int | None] = contextvars.ContextVar(
+    "turn_number"
+)
+_session_tracer_ctx: contextvars.ContextVar[Any | None] = contextvars.ContextVar(
+    "session_tracer"
+)
+
+
+def set_session_id(session_id: str | None) -> None:
+    """Set the current session ID in context.
+
+    This ID will be available to all async tasks spawned from the current context.
+    Setting to None clears the session context.
+
+    Args:
+        session_id: The session ID to set, or None to clear
+    """
+    _session_id_ctx.set(session_id)
+
+
+def get_session_id() -> str | None:
+    """Get the current session ID from context.
+
+    Returns:
+        The current session ID if one is set, None otherwise
+    """
+    return _session_id_ctx.get()
+
+
+def set_turn_number(turn: int | None) -> None:
+    """Set the current turn number in context."""
+    _turn_number_ctx.set(turn)
+
+
+def get_turn_number() -> int | None:
+    """Get the current turn number from context."""
+    return _turn_number_ctx.get()
+
+
+def set_session_tracer(tracer: Any) -> None:
+    """Set the current session tracer in context."""
+    _session_tracer_ctx.set(tracer)
+
+
+def get_session_tracer() -> Any:
+    """Get the current session tracer from context."""
+    return _session_tracer_ctx.get()
+
+
+T = TypeVar("T")
+
+
+@overload
+def with_session(require: bool = True) -> Callable[[Callable[..., Awaitable[T]]], Callable[..., Awaitable[T]]]:
+    ...
+
+
+@overload
+def with_session(require: bool = True) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    ...
+
+
+def with_session(require: bool = True):
+    """Decorator that ensures a session is active.
+
+    This decorator checks if a session is active before allowing the decorated
+    function to execute. It supports both sync and async functions.
+
+    Args:
+        require: If True, raises RuntimeError when no session is active.
+                If False, allows execution without a session (useful for
+                optional tracing).
+
+    Example:
+        ```python
+        @with_session()
+        async def process_message(content: str):
+            # This will only run if a session is active
+            tracer = get_session_tracer()
+            await tracer.record_message(content, "user")
+        ```
+    """
+
+    def decorator(fn: Callable[..., Awaitable[T]] | Callable[..., T]) -> Callable[..., Awaitable[T]] | Callable[..., T]:
+        if asyncio.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> T:
+                session_id = get_session_id()
+                if require and session_id is None:
+                    raise RuntimeError(
+                        f"No active session for {getattr(fn, '__name__', 'unknown')}"
+                    )
+                async_fn = cast(Callable[..., Awaitable[T]], fn)
+                return await async_fn(*args, **kwargs)
+
+            return async_wrapper
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> T:
+                session_id = get_session_id()
+                if require and session_id is None:
+                    raise RuntimeError(
+                        f"No active session for {getattr(fn, '__name__', 'unknown')}"
+                    )
+                sync_fn = cast(Callable[..., T], fn)
+                return sync_fn(*args, **kwargs)
+
+            return sync_wrapper
+
+    return decorator
+
+
+def trace_llm_call(
+    model_name: str | None = None,
+    system_id: str = "llm",
+    extract_tokens: bool = True,
+    extract_cost: bool = True,
+):
+    """Decorator to trace LLM API calls.
+
+    Automatically records LLM API calls as LMCAISEvent instances. Extracts token
+    counts, calculates costs, and measures latency. Only works with async functions.
+
+    Args:
+        model_name: Model name to record (can be overridden by actual response)
+        system_id: System identifier for the event (default: "llm")
+        extract_tokens: Whether to extract token counts from response
+        extract_cost: Whether to calculate USD cost from token counts
+
+    Expected Response Format:
+        The decorated function should return a dict with:
+        - 'usage': dict with 'prompt_tokens', 'completion_tokens', 'total_tokens'
+        - 'model': actual model name (optional, falls back to model_name param)
+
+    Example:
+        ```python
+        @trace_llm_call(model_name="gpt-4")
+        async def call_openai(prompt: str) -> dict:
+            response = await openai_client.completions.create(
+                model="gpt-4",
+                messages=[{"role": "user", "content": prompt}]
+            )
+            return response.model_dump()
+        ```
+    """
+
+    def decorator(fn: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
+        if asyncio.iscoroutinefunction(fn):
+            async_fn: Callable[..., Awaitable[T]] = fn
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> T:
+                tracer = get_session_tracer()
+                if not tracer:
+                    return await async_fn(*args, **kwargs)
+
+                start_time = time.time()
+                system_state_before = kwargs.get("state_before", {})
+
+                try:
+                    result = await async_fn(*args, **kwargs)
+
+                    # Extract metrics from result - this assumes the result follows
+                    # common LLM API response formats (OpenAI, Anthropic, etc.)
+                    input_tokens = output_tokens = total_tokens = None
+                    actual_model = model_name
+                    if extract_tokens and isinstance(result, Mapping):
+                        result_mapping = cast(Mapping[str, Any], result)
+                        usage = result_mapping.get("usage")
+                        if isinstance(usage, Mapping):
+                            input_tokens = usage.get("prompt_tokens")
+                            output_tokens = usage.get("completion_tokens")
+                            total_tokens = usage.get("total_tokens")
+                        value = result_mapping.get("model")
+                        if isinstance(value, str):
+                            actual_model = value
+
+                    latency_ms = int((time.time() - start_time) * 1000)
+
+                    # Create event with all extracted metrics
+                    event = LMCAISEvent(
+                        system_instance_id=system_id,
+                        time_record=TimeRecord(event_time=time.time()),
+                        model_name=actual_model or "unknown",
+                        provider=detect_provider(actual_model),
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                        total_tokens=total_tokens,
+                        cost_usd=calculate_cost(
+                            actual_model or "unknown", input_tokens or 0, output_tokens or 0
+                        )
+                        if extract_cost
+                        else None,
+                        latency_ms=latency_ms,
+                        system_state_before=system_state_before,
+                        system_state_after=kwargs.get("state_after", {}),
+                        metadata={
+                            "function": getattr(fn, "__name__", "unknown"),
+                            "step_id": kwargs.get("step_id"),
+                        },
+                    )
+
+                    await tracer.record_event(event)
+                    return result
+
+                except Exception as e:
+                    # Record error event - we still want to track failed LLM calls
+                    # for debugging and cost analysis
+                    if tracer:
+                        event = LMCAISEvent(
+                            system_instance_id=system_id,
+                            time_record=TimeRecord(event_time=time.time()),
+                            model_name=model_name or "unknown",
+                            provider=detect_provider(model_name),
+                            latency_ms=int((time.time() - start_time) * 1000),
+                            metadata={
+                                "function": getattr(fn, "__name__", "unknown"),
+                                "error": str(e),
+                                "error_type": type(e).__name__,
+                            },
+                        )
+                        await tracer.record_event(event)
+                    raise
+
+            return async_wrapper
+        else:
+            raise ValueError("trace_llm_call only supports async functions")
+
+    return decorator
+
+
+def trace_method(event_type: str = "runtime", system_id: str | None = None):
+    """Generic method tracing decorator.
+
+    Traces any method call by recording it as a RuntimeEvent. Supports both
+    sync and async methods, though async is preferred.
+
+    Args:
+        event_type: Type of event to create (default: "runtime")
+        system_id: System identifier (defaults to class name)
+
+    Example:
+        ```python
+        class Agent:
+            @trace_method(system_id="my_agent")
+            async def take_action(self, observation):
+                # Method execution is automatically traced
+                return self.policy(observation)
+        ```
+    """
+
+    def decorator(
+        fn: Callable[..., Awaitable[T]] | Callable[..., T]
+    ) -> Callable[..., Awaitable[T]] | Callable[..., T]:
+        if asyncio.iscoroutinefunction(fn):
+            async_fn = cast(Callable[..., Awaitable[T]], fn)
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> T:
+                tracer = get_session_tracer()
+                if not tracer:
+                    return await async_fn(*args, **kwargs)
+
+                from .abstractions import RuntimeEvent
+
+                # Use class name as system_id if not provided
+                self_obj = args[0] if args else None
+                inferred_system_id = (
+                    self_obj.__class__.__name__ if self_obj is not None else "unknown"
+                )
+                actual_system_id = system_id or inferred_system_id
+
+                event = RuntimeEvent(
+                    system_instance_id=actual_system_id,
+                    time_record=TimeRecord(event_time=time.time()),
+                    actions=[],  # Can be overridden in metadata
+                    metadata={
+                        "method": getattr(fn, "__name__", "unknown"),
+                        "args": str(args)[:100],  # Truncate for safety
+                        "step_id": kwargs.get("step_id"),
+                    },
+                )
+
+                await tracer.record_event(event)
+                return await async_fn(*args, **kwargs)
+
+            return async_wrapper
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> T:
+                # For sync methods, we can't easily trace without blocking
+                # the event loop. This is a limitation of the async-first design.
+                # Consider converting to async or using a different approach
+                sync_fn = cast(Callable[..., T], fn)
+                return sync_fn(*args, **kwargs)
+
+            return sync_wrapper
+
+    return decorator
+
+
+class SessionContext:
+    """Context manager for session tracking.
+
+    Provides a way to temporarily set session context, useful for testing
+    or when you need to manually manage context outside of SessionTracer.
+
+    This context manager properly handles both sync and async contexts,
+    and ensures the previous context is restored on exit.
+
+    Example:
+        ```python
+        # Sync usage
+        with SessionContext("test_session_123", tracer):
+            # Code here sees the test session
+            process_data()
+
+        # Async usage
+        async with SessionContext("test_session_123", tracer):
+            # Async code here sees the test session
+            await process_data_async()
+        ```
+    """
+
+    def __init__(self, session_id: str, tracer: Any | None = None):
+        self.session_id = session_id
+        self.tracer = tracer
+        self._token: Token[str | None] | None = None
+        self._tracer_token: Token[Any] | None = None
+
+    def __enter__(self):
+        # Store tokens to restore previous context on exit
+        self._token = _session_id_ctx.set(self.session_id)
+        if self.tracer:
+            self._tracer_token = _session_tracer_ctx.set(self.tracer)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Restore previous context - this is crucial for proper isolation
+        if self._token is not None:
+            _session_id_ctx.reset(self._token)
+        if self._tracer_token is not None:
+            _session_tracer_ctx.reset(self._tracer_token)
+
+    async def __aenter__(self):
+        self._token = _session_id_ctx.set(self.session_id)
+        if self.tracer:
+            self._tracer_token = _session_tracer_ctx.set(self.tracer)
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._token is not None:
+            _session_id_ctx.reset(self._token)
+        if self._tracer_token is not None:
+            _session_tracer_ctx.reset(self._tracer_token)

synth_ai/core/tracing_v3/examples/basic_usage.py

@@ -0,0 +1,194 @@
+"""Basic usage example for tracing v3."""
+
+import asyncio
+import time
+from typing import Any
+
+from .. import SessionTracer
+from ..abstractions import EnvironmentEvent, LMCAISEvent, RuntimeEvent, TimeRecord
+from ..turso.daemon import SqldDaemon
+
+
+async def simulate_llm_call(model: str, prompt: str) -> dict[str, Any]:
+    """Simulate an LLM API call."""
+    await asyncio.sleep(0.1)  # Simulate network latency
+
+    # Simulate response
+    tokens = len(prompt.split()) * 3
+    return {
+        "model": model,
+        "response": f"Response to: {prompt[:50]}...",
+        "usage": {
+            "prompt_tokens": len(prompt.split()) * 2,
+            "completion_tokens": tokens,
+            "total_tokens": len(prompt.split()) * 2 + tokens,
+        },
+    }
+
+
+async def main():
+    """Demonstrate basic tracing v3 usage."""
+    print("Starting tracing v3 example...")
+
+    # Option 1: Start sqld daemon programmatically
+    with SqldDaemon():
+        print("✓ Started sqld daemon")
+
+        # Wait for daemon to be ready
+        await asyncio.sleep(1)
+
+        # Create tracer
+        tracer = SessionTracer()
+        await tracer.initialize()
+        print("✓ Initialized tracer")
+
+        # Example 1: Basic session with events
+        print("\n--- Example 1: Basic Session ---")
+        async with tracer.session(metadata={"example": "basic"}) as session_id:
+            print(f"Started session: {session_id}")
+
+            # Timestep 1: LLM interaction
+            async with tracer.timestep("llm_step", turn_number=1):
+                # Simulate LLM call
+                result = await simulate_llm_call("gpt-4", "What is the capital of France?")
+
+                # Record LLM event
+                event = LMCAISEvent(
+                    system_instance_id="llm_system",
+                    time_record=TimeRecord(event_time=time.time()),
+                    model_name=result["model"],
+                    input_tokens=result["usage"]["prompt_tokens"],
+                    output_tokens=result["usage"]["completion_tokens"],
+                    total_tokens=result["usage"]["total_tokens"],
+                    cost_usd=0.003,  # $0.003
+                    latency_ms=100,
+                    metadata={"prompt": "What is the capital of France?"},
+                )
+                await tracer.record_event(event)
+
+                # Record messages
+                await tracer.record_message(
+                    content="What is the capital of France?", message_type="user"
+                )
+                await tracer.record_message(content=result["response"], message_type="assistant")
+                print("✓ Recorded LLM interaction")
+
+            # Timestep 2: Environment interaction
+            async with tracer.timestep("env_step", turn_number=2):
+                # Record environment event
+                env_event = EnvironmentEvent(
+                    system_instance_id="environment",
+                    time_record=TimeRecord(event_time=time.time()),
+                    reward=0.8,
+                    terminated=False,
+                    system_state_before={"position": [0, 0]},
+                    system_state_after={"position": [1, 0]},
+                )
+                await tracer.record_event(env_event)
+                print("✓ Recorded environment event")
+
+            # Timestep 3: Runtime action
+            async with tracer.timestep("runtime_step", turn_number=3):
+                # Record runtime event
+                runtime_event = RuntimeEvent(
+                    system_instance_id="agent",
+                    time_record=TimeRecord(event_time=time.time()),
+                    actions=[1, 0, 0, 1],  # Example action vector
+                    metadata={"action_type": "move_right"},
+                )
+                await tracer.record_event(runtime_event)
+                print("✓ Recorded runtime event")
+
+        print(f"✓ Session {session_id} saved\n")
+
+        # Example 2: Concurrent sessions
+        print("--- Example 2: Concurrent Sessions ---")
+
+        async def run_concurrent_session(session_num: int):
+            """Run a session concurrently."""
+            async with tracer.session(
+                metadata={"example": "concurrent", "session_num": session_num}
+            ) as sid:
+                for i in range(3):
+                    async with tracer.timestep(f"step_{i}", turn_number=i):
+                        # Simulate some work
+                        await asyncio.sleep(0.05)
+
+                        # Record event
+                        event = RuntimeEvent(
+                            system_instance_id=f"worker_{session_num}",
+                            time_record=TimeRecord(event_time=time.time()),
+                            actions=[i],
+                            metadata={"iteration": i},
+                        )
+                        await tracer.record_event(event)
+
+                return sid
+
+        # Run 5 concurrent sessions
+        tasks = [run_concurrent_session(i) for i in range(5)]
+        session_ids = await asyncio.gather(*tasks)
+        print(f"✓ Completed {len(session_ids)} concurrent sessions")
+
+        # Example 3: Query stored data
+        print("\n--- Example 3: Querying Data ---")
+
+        # Get model usage statistics
+        if tracer.db is None:
+            raise RuntimeError("Tracer database backend is not initialized")
+
+        model_usage = await tracer.db.get_model_usage()
+        print("\nModel Usage:")
+        print(model_usage)
+
+        # Query recent sessions
+        recent_sessions = await tracer.get_session_history(limit=5)
+        print(f"\nRecent Sessions: {len(recent_sessions)} found")
+        for session in recent_sessions:
+            print(
+                f"  - {session['session_id']}: "
+                f"{session['num_events']} events, "
+                f"{session['num_messages']} messages"
+            )
+
+        # Get specific session details
+        if recent_sessions:
+            session_detail = await tracer.db.get_session_trace(recent_sessions[0]["session_id"])
+            if session_detail:
+                print(f"\nSession Detail for {session_detail['session_id']}:")
+                print(f"  Created: {session_detail['created_at']}")
+                print(f"  Timesteps: {len(session_detail['timesteps'])}")
+
+        # Example 4: Using hooks
+        print("\n--- Example 4: Hooks ---")
+
+        # Add a custom hook
+        call_count = {"count": 0}
+
+        async def count_events(event, **kwargs):
+            call_count["count"] += 1
+            print(f"  Hook: Event #{call_count['count']} recorded")
+
+        tracer.hooks.register("event_recorded", count_events, name="event_counter")
+
+        async with (
+            tracer.session(metadata={"example": "hooks"}) as session_id,
+            tracer.timestep("hook_test"),
+        ):
+            for i in range(3):
+                event = RuntimeEvent(
+                    system_instance_id="hook_test",
+                    time_record=TimeRecord(event_time=time.time()),
+                    actions=[i],
+                )
+                await tracer.record_event(event)
+
+        print(f"✓ Hook called {call_count['count']} times")
+
+        # Cleanup
+        await tracer.close()
+        print("\n✓ Example completed successfully!")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())