synth-ai 0.2.6.dev1__py3-none-any.whl → 0.4.3__py3-none-any.whl

synth_ai/{tracing_v3 → core/tracing_v3}/decorators.py

@@ -1,4 +1,3 @@
-from __future__ import annotations
 """Async-aware decorators for tracing v3.
 
 This module provides decorators and context management utilities for the tracing
@@ -23,12 +22,15 @@ 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 Callable
-from typing import Any, TypeVar
+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
@@ -37,13 +39,13 @@ from .utils import calculate_cost, detect_provider
 # 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", default=None
+    "session_id"
 )
 _turn_number_ctx: contextvars.ContextVar[int | None] = contextvars.ContextVar(
-    "turn_number", default=None
+    "turn_number"
 )
 _session_tracer_ctx: contextvars.ContextVar[Any | None] = contextvars.ContextVar(
-    "session_tracer", default=None
+    "session_tracer"
 )
 
 
@@ -91,6 +93,16 @@ def get_session_tracer() -> Any:
 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.
 
@@ -112,25 +124,31 @@ def with_session(require: bool = True):
         ```
     """
 
-    def decorator(fn: Callable[..., T]) -> Callable[..., T]:
+    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, **kwargs):
+            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 {fn.__name__}")
-                return await fn(*args, **kwargs)
+                    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, **kwargs):
+            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 {fn.__name__}")
-                return fn(*args, **kwargs)
+                    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
 
@@ -138,7 +156,7 @@ def with_session(require: bool = True):
 
 
 def trace_llm_call(
-    model_name: str = None,
+    model_name: str | None = None,
     system_id: str = "llm",
     extract_tokens: bool = True,
     extract_cost: bool = True,
@@ -171,31 +189,36 @@ def trace_llm_call(
         ```
     """
 
-    def decorator(fn: Callable[..., T]) -> Callable[..., T]:
+    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, **kwargs):
+            async def async_wrapper(*args: Any, **kwargs: Any) -> T:
                 tracer = get_session_tracer()
                 if not tracer:
-                    return await fn(*args, **kwargs)
+                    return await async_fn(*args, **kwargs)
 
                 start_time = time.time()
                 system_state_before = kwargs.get("state_before", {})
 
                 try:
-                    result = await fn(*args, **kwargs)
+                    result = await async_fn(*args, **kwargs)
 
                     # Extract metrics from result - this assumes the result follows
                     # common LLM API response formats (OpenAI, Anthropic, etc.)
-                    if extract_tokens and isinstance(result, dict):
-                        input_tokens = result.get("usage", {}).get("prompt_tokens")
-                        output_tokens = result.get("usage", {}).get("completion_tokens")
-                        total_tokens = result.get("usage", {}).get("total_tokens")
-                        actual_model = result.get("model", model_name)
-                    else:
-                        input_tokens = output_tokens = total_tokens = None
-                        actual_model = model_name
+                    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)
 
@@ -208,14 +231,16 @@ def trace_llm_call(
                         input_tokens=input_tokens,
                         output_tokens=output_tokens,
                         total_tokens=total_tokens,
-                        cost_usd=calculate_cost(actual_model, input_tokens or 0, output_tokens or 0)
+                        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": fn.__name__,
+                            "function": getattr(fn, "__name__", "unknown"),
                             "step_id": kwargs.get("step_id"),
                         },
                     )
@@ -234,7 +259,7 @@ def trace_llm_call(
                             provider=detect_provider(model_name),
                             latency_ms=int((time.time() - start_time) * 1000),
                             metadata={
-                                "function": fn.__name__,
+                                "function": getattr(fn, "__name__", "unknown"),
                                 "error": str(e),
                                 "error_type": type(e).__name__,
                             },
@@ -249,7 +274,7 @@ def trace_llm_call(
     return decorator
 
 
-def trace_method(event_type: str = "runtime", system_id: str = None):
+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
@@ -269,43 +294,51 @@ def trace_method(event_type: str = "runtime", system_id: str = None):
         ```
     """
 
-    def decorator(fn: Callable[..., T]) -> Callable[..., T]:
+    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(self, *args, **kwargs):
+            async def async_wrapper(*args: Any, **kwargs: Any) -> T:
                 tracer = get_session_tracer()
                 if not tracer:
-                    return await fn(self, *args, **kwargs)
+                    return await async_fn(*args, **kwargs)
 
                 from .abstractions import RuntimeEvent
 
                 # Use class name as system_id if not provided
-                actual_system_id = system_id or self.__class__.__name__
+                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": fn.__name__,
+                        "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 fn(self, *args, **kwargs)
+                return await async_fn(*args, **kwargs)
 
             return async_wrapper
         else:
 
             @functools.wraps(fn)
-            def sync_wrapper(self, *args, **kwargs):
+            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
-                return fn(self, *args, **kwargs)
+                sync_fn = cast(Callable[..., T], fn)
+                return sync_fn(*args, **kwargs)
 
             return sync_wrapper
 
@@ -335,11 +368,11 @@ class SessionContext:
         ```
     """
 
-    def __init__(self, session_id: str, tracer=None):
+    def __init__(self, session_id: str, tracer: Any | None = None):
         self.session_id = session_id
         self.tracer = tracer
-        self._token = None
-        self._tracer_token = None
+        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
@@ -350,8 +383,9 @@ class SessionContext:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         # Restore previous context - this is crucial for proper isolation
-        _session_id_ctx.reset(self._token)
-        if self._tracer_token:
+        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):
@@ -361,6 +395,7 @@ class SessionContext:
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        _session_id_ctx.reset(self._token)
-        if self._tracer_token:
+        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/{tracing_v3 → core/tracing_v3}/examples/basic_usage.py

@@ -2,13 +2,14 @@
 
 import asyncio
 import time
+from typing import Any
 
-from synth_ai.tracing_v3 import SessionTracer
-from synth_ai.tracing_v3.abstractions import EnvironmentEvent, LMCAISEvent, RuntimeEvent, TimeRecord
-from synth_ai.tracing_v3.turso.daemon import SqldDaemon
+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:
+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
 
@@ -133,6 +134,9 @@ async def main():
         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)
@@ -150,9 +154,10 @@ async def main():
         # Get specific session details
         if recent_sessions:
             session_detail = await tracer.db.get_session_trace(recent_sessions[0]["session_id"])
-            print(f"\nSession Detail for {session_detail['session_id']}:")
-            print(f"  Created: {session_detail['created_at']}")
-            print(f"  Timesteps: {len(session_detail['timesteps'])}")
+            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 ---")
@@ -166,8 +171,9 @@ async def main():
 
         tracer.hooks.register("event_recorded", count_events, name="event_counter")
 
-        async with tracer.session(metadata={"example": "hooks"}) as session_id, tracer.timestep(
-            "hook_test"
+        async with (
+            tracer.session(metadata={"example": "hooks"}) as session_id,
+            tracer.timestep("hook_test"),
         ):
             for i in range(3):
                 event = RuntimeEvent(

synth_ai/{tracing_v3 → core/tracing_v3}/hooks.py

@@ -1,4 +1,3 @@
-from __future__ import annotations
 """Hook system for extending tracing functionality.
 
 The hook system provides a flexible way to extend the tracing system without
@@ -33,6 +32,8 @@ Common Use Cases:
 - Custom filtering and sampling
 """
 
+from __future__ import annotations
+
 import asyncio
 from collections.abc import Callable
 from dataclasses import dataclass
@@ -88,9 +89,9 @@ class HookManager:
         self,
         event: str,
         callback: Callable,
-        name: str = None,
+        name: str | None = None,
         priority: int = 0,
-        event_types: list[str] = None,
+        event_types: list[str] | None = None,
     ) -> Hook:
         """Register a new hook.
 
@@ -114,7 +115,7 @@ class HookManager:
             raise ValueError(f"Unknown hook event: {event}")
 
         hook = Hook(
-            name=name or callback.__name__,
+            name=name or getattr(callback, "__name__", "unknown"),
             callback=callback,
             event_types=event_types,
             priority=priority,
@@ -202,6 +203,7 @@ def create_default_hooks() -> HookManager:
     # Example: Log session starts - useful for debugging and monitoring
     async def log_session_start(session_id: str, metadata: dict[str, Any]):
         import os
+
         if os.getenv("SYNTH_TRACE_VERBOSE", "0") in ("1", "true", "True"):
             print(f"Session started: {session_id}")