synth-ai 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

synth_ai/task/server.py

@@ -6,6 +6,7 @@ import asyncio
 import inspect
 import os
 from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
@@ -34,6 +35,10 @@ InstanceProvider = Callable[[Sequence[int]], Iterable[TaskInfo] | Awaitable[Iter
 RolloutExecutor = Callable[[RolloutRequest, Request], Any | Awaitable[Any]]
 
 
+def _default_app_state() -> dict[str, Any]:
+    return {}
+
+
 @dataclass(slots=True)
 class RubricBundle:
     """Optional rubrics advertised by the task app."""
@@ -69,7 +74,7 @@ class TaskAppConfig:
     proxy: ProxyConfig | None = None
     routers: Sequence[APIRouter] = field(default_factory=tuple)
     middleware: Sequence[Middleware] = field(default_factory=tuple)
-    app_state: Mapping[str, Any] = field(default_factory=dict)
+    app_state: MutableMapping[str, Any] = field(default_factory=_default_app_state)
     require_api_key: bool = True
     expose_debug_env: bool = True
     cors_origins: Sequence[str] | None = None
@@ -260,17 +265,19 @@ def create_task_app(config: TaskAppConfig) -> FastAPI:
             return _maybe_await(hook(app))  # type: ignore[misc]
         return _maybe_await(hook())
 
-    @app.on_event("startup")
-    async def _startup() -> None:  # pragma: no cover - FastAPI lifecycle
+    @asynccontextmanager
+    async def lifespan(_: FastAPI):
         normalize_environment_api_key()
         normalize_vendor_keys()
         for hook in cfg.startup_hooks:
             await _call_hook(hook)
+        try:
+            yield
+        finally:
+            for hook in cfg.shutdown_hooks:
+                await _call_hook(hook)
 
-    @app.on_event("shutdown")
-    async def _shutdown() -> None:  # pragma: no cover - FastAPI lifecycle
-        for hook in cfg.shutdown_hooks:
-            await _call_hook(hook)
+    app.router.lifespan_context = lifespan
 
     @app.get("/")
     async def root() -> Mapping[str, Any]:

synth_ai/tracing_v3/decorators.py

@@ -28,8 +28,8 @@ 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 typing import Any, TypeVar, cast, overload
 
 from .abstractions import LMCAISEvent, TimeRecord
 from .utils import calculate_cost, detect_provider
@@ -88,6 +88,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.
 
@@ -109,29 +119,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 {getattr(fn, '__name__', 'unknown')}"
                     )
-                return await fn(*args, **kwargs)
+                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 {getattr(fn, '__name__', 'unknown')}"
                     )
-                return fn(*args, **kwargs)
+                sync_fn = cast(Callable[..., T], fn)
+                return sync_fn(*args, **kwargs)
 
             return sync_wrapper
 
@@ -172,31 +184,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)
 
@@ -272,19 +289,26 @@ def trace_method(event_type: str = "runtime", system_id: str | None = 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,
@@ -298,17 +322,18 @@ def trace_method(event_type: str = "runtime", system_id: str | None = None):
                 )
 
                 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
 

synth_ai/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 ---")

synth_ai/tracing_v3/llm_call_record_helpers.py

@@ -4,11 +4,14 @@ This module provides utilities to convert vendor responses to LLMCallRecord
 format and compute aggregates from call records.
 """
 
+from __future__ import annotations
+
 import uuid
+from dataclasses import dataclass, field
 from datetime import UTC, datetime
-from typing import Any
+from typing import Any, TypedDict, cast
 
-from synth_ai.tracing_v3.lm_call_record_abstractions import (
+from .lm_call_record_abstractions import (
     LLMCallRecord,
     LLMChunk,
     LLMContentPart,
@@ -17,7 +20,21 @@ from synth_ai.tracing_v3.lm_call_record_abstractions import (
     LLMUsage,
     ToolCallSpec,
 )
-from synth_ai.v0.lm.vendors.base import BaseLMResponse
+
+BaseLMResponse = Any
+
+
+class _UsageDict(TypedDict, total=False):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    reasoning_tokens: int
+    cost_usd: float
+    duration_ms: int
+    reasoning_input_tokens: int
+    reasoning_output_tokens: int
+    cache_write_tokens: int
+    cache_read_tokens: int
 
 
 def create_llm_call_record_from_response(
@@ -110,9 +127,10 @@ def create_llm_call_record_from_response(
         )
 
     # Extract tool calls if present
-    output_tool_calls = []
-    if hasattr(response, "tool_calls") and response.tool_calls:
-        for idx, tool_call in enumerate(response.tool_calls):
+    output_tool_calls: list[ToolCallSpec] = []
+    tool_calls_data = cast(list[dict[str, Any]] | None, getattr(response, "tool_calls", None))
+    if tool_calls_data:
+        for idx, tool_call in enumerate(tool_calls_data):
             if isinstance(tool_call, dict):
                 output_tool_calls.append(
                     ToolCallSpec(
@@ -125,18 +143,19 @@ def create_llm_call_record_from_response(
 
     # Extract usage information
     usage = None
-    if hasattr(response, "usage") and response.usage:
+    usage_data = cast(_UsageDict | None, getattr(response, "usage", None))
+    if usage_data:
         usage = LLMUsage(
-            input_tokens=response.usage.get("input_tokens"),
-            output_tokens=response.usage.get("output_tokens"),
-            total_tokens=response.usage.get("total_tokens"),
-            cost_usd=response.usage.get("cost_usd"),
+            input_tokens=usage_data.get("input_tokens"),
+            output_tokens=usage_data.get("output_tokens"),
+            total_tokens=usage_data.get("total_tokens"),
+            cost_usd=usage_data.get("cost_usd"),
             # Additional token accounting if available
-            reasoning_tokens=response.usage.get("reasoning_tokens"),
-            reasoning_input_tokens=response.usage.get("reasoning_input_tokens"),
-            reasoning_output_tokens=response.usage.get("reasoning_output_tokens"),
-            cache_write_tokens=response.usage.get("cache_write_tokens"),
-            cache_read_tokens=response.usage.get("cache_read_tokens"),
+            reasoning_tokens=usage_data.get("reasoning_tokens"),
+            reasoning_input_tokens=usage_data.get("reasoning_input_tokens"),
+            reasoning_output_tokens=usage_data.get("reasoning_output_tokens"),
+            cache_write_tokens=usage_data.get("cache_write_tokens"),
+            cache_read_tokens=usage_data.get("cache_read_tokens"),
         )
 
     # Build request parameters
@@ -188,7 +207,45 @@ def create_llm_call_record_from_response(
     return record
 
 
-def compute_aggregates_from_call_records(call_records: list[LLMCallRecord]) -> dict[str, Any]:
+@dataclass
+class _AggregateAccumulator:
+    """Mutable accumulator for call record aggregates."""
+
+    call_count: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+    reasoning_tokens: int = 0
+    cost_usd: float = 0.0
+    latency_ms: int = 0
+    models_used: set[str] = field(default_factory=set)
+    providers_used: set[str] = field(default_factory=set)
+    tool_calls_count: int = 0
+    error_count: int = 0
+    success_count: int = 0
+
+
+class AggregateSummary(TypedDict, total=False):
+    """Aggregate metrics derived from call records."""
+
+    call_count: int
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+    reasoning_tokens: int
+    cost_usd: float
+    latency_ms: int
+    models_used: list[str]
+    providers_used: list[str]
+    tool_calls_count: int
+    error_count: int
+    success_count: int
+    avg_latency_ms: float
+    avg_input_tokens: float
+    avg_output_tokens: float
+
+
+def compute_aggregates_from_call_records(call_records: list[LLMCallRecord]) -> AggregateSummary:
     """Compute aggregate statistics from a list of LLMCallRecord instances.
 
     Args:
@@ -197,65 +254,62 @@ def compute_aggregates_from_call_records(call_records: list[LLMCallRecord]) -> d
     Returns:
         Dictionary containing aggregated statistics
     """
-    aggregates = {
-        "input_tokens": 0,
-        "output_tokens": 0,
-        "total_tokens": 0,
-        "reasoning_tokens": 0,
-        "cost_usd": 0.0,
-        "latency_ms": 0,
-        "models_used": set(),
-        "providers_used": set(),
-        "tool_calls_count": 0,
-        "error_count": 0,
-        "success_count": 0,
-        "call_count": len(call_records),
-    }
+    aggregates = _AggregateAccumulator(call_count=len(call_records))
 
     for record in call_records:
         # Token aggregation
         if record.usage:
             if record.usage.input_tokens:
-                aggregates["input_tokens"] += record.usage.input_tokens
+                aggregates.input_tokens += record.usage.input_tokens
             if record.usage.output_tokens:
-                aggregates["output_tokens"] += record.usage.output_tokens
+                aggregates.output_tokens += record.usage.output_tokens
             if record.usage.total_tokens:
-                aggregates["total_tokens"] += record.usage.total_tokens
+                aggregates.total_tokens += record.usage.total_tokens
             if record.usage.reasoning_tokens:
-                aggregates["reasoning_tokens"] += record.usage.reasoning_tokens
+                aggregates.reasoning_tokens += record.usage.reasoning_tokens
             if record.usage.cost_usd:
-                aggregates["cost_usd"] += record.usage.cost_usd
+                aggregates.cost_usd += record.usage.cost_usd
 
         # Latency aggregation
-        if record.latency_ms:
-            aggregates["latency_ms"] += record.latency_ms
+        if record.latency_ms is not None:
+            aggregates.latency_ms += record.latency_ms
 
         # Model and provider tracking
         if record.model_name:
-            aggregates["models_used"].add(record.model_name)
+            aggregates.models_used.add(record.model_name)
         if record.provider:
-            aggregates["providers_used"].add(record.provider)
+            aggregates.providers_used.add(record.provider)
 
         # Tool calls
-        aggregates["tool_calls_count"] += len(record.output_tool_calls)
+        aggregates.tool_calls_count += len(record.output_tool_calls)
 
         # Success/error tracking
         if record.outcome == "error":
-            aggregates["error_count"] += 1
+            aggregates.error_count += 1
         elif record.outcome == "success":
-            aggregates["success_count"] += 1
-
-    # Convert sets to lists for JSON serialization
-    aggregates["models_used"] = list(aggregates["models_used"])
-    aggregates["providers_used"] = list(aggregates["providers_used"])
+            aggregates.success_count += 1
+
+    summary: AggregateSummary = {
+        "call_count": aggregates.call_count,
+        "input_tokens": aggregates.input_tokens,
+        "output_tokens": aggregates.output_tokens,
+        "total_tokens": aggregates.total_tokens,
+        "reasoning_tokens": aggregates.reasoning_tokens,
+        "cost_usd": aggregates.cost_usd,
+        "latency_ms": aggregates.latency_ms,
+        "models_used": list(aggregates.models_used),
+        "providers_used": list(aggregates.providers_used),
+        "tool_calls_count": aggregates.tool_calls_count,
+        "error_count": aggregates.error_count,
+        "success_count": aggregates.success_count,
+    }
 
-    # Compute averages
-    if aggregates["call_count"] > 0:
-        aggregates["avg_latency_ms"] = aggregates["latency_ms"] / aggregates["call_count"]
-        aggregates["avg_input_tokens"] = aggregates["input_tokens"] / aggregates["call_count"]
-        aggregates["avg_output_tokens"] = aggregates["output_tokens"] / aggregates["call_count"]
+    if aggregates.call_count > 0:
+        summary["avg_latency_ms"] = aggregates.latency_ms / aggregates.call_count
+        summary["avg_input_tokens"] = aggregates.input_tokens / aggregates.call_count
+        summary["avg_output_tokens"] = aggregates.output_tokens / aggregates.call_count
 
-    return aggregates
+    return summary
 
 
 def create_llm_call_record_from_streaming(

synth_ai/tracing_v3/replica_sync.py

@@ -26,6 +26,7 @@ application to continue without blocking on sync operations.
 
 import asyncio
 import logging
+from typing import Any
 
 import libsql
 
@@ -66,8 +67,8 @@ class ReplicaSync:
         self.sync_url = sync_url or CONFIG.sync_url
         self.auth_token = auth_token or CONFIG.auth_token
         self.sync_interval = sync_interval or CONFIG.sync_interval
-        self._sync_task: asyncio.Task | None = None
-        self._conn: libsql.Connection | None = None
+        self._sync_task: asyncio.Task[Any] | None = None
+        self._conn: Any | None = None
 
     def _ensure_connection(self):
         """Ensure libsql connection is established.
@@ -113,8 +114,11 @@ class ReplicaSync:
         """
         try:
             self._ensure_connection()
+            conn = self._conn
+            if conn is None:
+                raise RuntimeError("Replica sync connection is not available after initialization")
             # Run sync in thread pool since libsql sync is blocking
-            await asyncio.to_thread(self._conn.sync)
+            await asyncio.to_thread(conn.sync)
             logger.info("Successfully synced with remote Turso database")
             return True
         except Exception as e:
@@ -146,7 +150,7 @@ class ReplicaSync:
             # Sleep until next sync interval
             await asyncio.sleep(self.sync_interval)
 
-    def start_background_sync(self) -> asyncio.Task:
+    def start_background_sync(self) -> asyncio.Task[Any]:
         """Start the background sync task.
 
         Creates an asyncio task that runs the sync loop. The task is stored

synth_ai/tracing_v3/storage/utils.py

@@ -3,8 +3,8 @@
 import asyncio
 import functools
 import time
-from collections.abc import Callable
-from typing import Any, TypeVar
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar, cast
 
 T = TypeVar("T")
 
@@ -18,9 +18,9 @@ def retry_async(max_attempts: int = 3, delay: float = 1.0, backoff: float = 2.0)
         backoff: Backoff multiplier for each retry
     """
 
-    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
         @functools.wraps(func)
-        async def wrapper(*args, **kwargs):
+        async def wrapper(*args: Any, **kwargs: Any) -> T:
             last_exception: Exception | None = None
             current_delay = delay
 
@@ -171,13 +171,14 @@ STORAGE_METRICS = StorageMetrics()
 def track_metrics(operation: str):
     """Decorator to track storage operation metrics."""
 
-    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+    def decorator(func: Callable[..., Awaitable[T]] | Callable[..., T]) -> Callable[..., Awaitable[T]] | Callable[..., T]:
         @functools.wraps(func)
-        async def async_wrapper(*args, **kwargs):
+        async def async_wrapper(*args: Any, **kwargs: Any) -> T:
             start_time = time.time()
             success = False
             try:
-                result = await func(*args, **kwargs)
+                async_func = cast(Callable[..., Awaitable[T]], func)
+                result = await async_func(*args, **kwargs)
                 success = True
                 return result
             finally:
@@ -185,11 +186,12 @@ def track_metrics(operation: str):
                 STORAGE_METRICS.record_operation(operation, duration, success)
 
         @functools.wraps(func)
-        def sync_wrapper(*args, **kwargs):
+        def sync_wrapper(*args: Any, **kwargs: Any) -> T:
             start_time = time.time()
             success = False
             try:
-                result = func(*args, **kwargs)
+                sync_func = cast(Callable[..., T], func)
+                result = sync_func(*args, **kwargs)
                 success = True
                 return result
             finally:

synth_ai/tracing_v3/turso/__init__.py

@@ -0,0 +1,12 @@
+"""Turso integration package for tracing v3."""
+
+from .daemon import SqldDaemon, get_daemon, start_sqld, stop_sqld
+from .native_manager import NativeLibsqlTraceManager
+
+__all__ = [
+    "SqldDaemon",
+    "NativeLibsqlTraceManager",
+    "get_daemon",
+    "start_sqld",
+    "stop_sqld",
+]

synth_ai/tracing_v3/turso/daemon.py

@@ -6,6 +6,7 @@ import subprocess
 import time
 
 import requests
+from requests import RequestException
 
 from ..config import CONFIG
 
@@ -79,7 +80,7 @@ class SqldDaemon:
                 response = requests.get(health_url, timeout=1)
                 if response.status_code == 200:
                     return
-            except requests.exceptions.RequestException:
+            except RequestException:
                 pass
 
             # Check if process crashed