synth-ai 0.2.4.dev5__py3-none-any.whl → 0.2.4.dev7__py3-none-any.whl

synth_ai/tracing_v3/session_tracer.py

@@ -1,15 +1,21 @@
 """Main SessionTracer class for tracing v3."""
 
 import asyncio
-from datetime import datetime
-from typing import Dict, List, Optional, Any, Union
 from contextlib import asynccontextmanager
-
-from .abstractions import SessionTrace, SessionTimeStep, BaseEvent, SessionEventMarkovBlanketMessage, TimeRecord
-from .decorators import set_session_id, set_turn_number, set_session_tracer, SessionContext
-from .turso.manager import AsyncSQLTraceManager
+from datetime import datetime
+from typing import Any
+
+from .abstractions import (
+    BaseEvent,
+    SessionEventMarkovBlanketMessage,
+    SessionTimeStep,
+    SessionTrace,
+    TimeRecord,
+)
 from .config import CONFIG
-from .hooks import HookManager, GLOBAL_HOOKS
+from .decorators import set_session_id, set_session_tracer, set_turn_number
+from .hooks import GLOBAL_HOOKS, HookManager
+from .turso.manager import AsyncSQLTraceManager
 from .utils import generate_session_id
 
 
@@ -18,8 +24,8 @@ class SessionTracer:
 
     def __init__(
         self,
-        hooks: Optional[HookManager] = None,
-        db_url: Optional[str] = None,
+        hooks: HookManager | None = None,
+        db_url: str | None = None,
         auto_save: bool = True,
     ):
         """Initialize session tracer.
@@ -30,20 +36,20 @@ class SessionTracer:
             auto_save: Whether to automatically save sessions on end
         """
         self.hooks = hooks or GLOBAL_HOOKS
-        self._current_trace: Optional[SessionTrace] = None
+        self._current_trace: SessionTrace | None = None
         self._lock = asyncio.Lock()
         self.db_url = db_url or CONFIG.db_url
-        self.db: Optional[AsyncSQLTraceManager] = None
+        self.db: AsyncSQLTraceManager | None = None
         self.auto_save = auto_save
-        self._current_step: Optional[SessionTimeStep] = None
+        self._current_step: SessionTimeStep | None = None
 
     @property
-    def current_session(self) -> Optional[SessionTrace]:
+    def current_session(self) -> SessionTrace | None:
         """Get the current session trace."""
         return self._current_trace
 
     @property
-    def current_step(self) -> Optional[SessionTimeStep]:
+    def current_step(self) -> SessionTimeStep | None:
         """Get the current timestep."""
         return self._current_step
 
@@ -54,14 +60,14 @@ class SessionTracer:
             await self.db.initialize()
 
     async def start_session(
-        self, session_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None
+        self, session_id: str | None = None, metadata: dict[str, Any] | None = None
     ) -> str:
         """Start a new session.
-        
+
         Creates a new tracing session and sets up the necessary context variables.
         This method is thread-safe and will raise an error if a session is already
         active to prevent accidental session mixing.
-        
+
         The session ID is propagated through asyncio context variables, allowing
         nested async functions to access the tracer without explicit passing.
 
@@ -76,7 +82,7 @@ class SessionTracer:
 
         Returns:
             The session ID (either provided or generated)
-            
+
         Raises:
             RuntimeError: If a session is already active
         """
@@ -111,8 +117,8 @@ class SessionTracer:
     async def start_timestep(
         self,
         step_id: str,
-        turn_number: Optional[int] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        turn_number: int | None = None,
+        metadata: dict[str, Any] | None = None,
     ) -> SessionTimeStep:
         """Start a new timestep.
 
@@ -148,7 +154,7 @@ class SessionTracer:
 
         return step
 
-    async def end_timestep(self, step_id: Optional[str] = None):
+    async def end_timestep(self, step_id: str | None = None):
         """End the current or specified timestep."""
         if self._current_trace is None:
             raise RuntimeError("No active session")
@@ -199,9 +205,9 @@ class SessionTracer:
         self,
         content: str,
         message_type: str,
-        event_time: Optional[float] = None,
-        message_time: Optional[int] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        event_time: float | None = None,
+        message_time: int | None = None,
+        metadata: dict[str, Any] | None = None,
     ):
         """Record a message.
 
@@ -281,8 +287,8 @@ class SessionTracer:
     @asynccontextmanager
     async def session(
         self,
-        session_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        session_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
         save: bool = None,
     ):
         """Context manager for a session.
@@ -302,8 +308,8 @@ class SessionTracer:
     async def timestep(
         self,
         step_id: str,
-        turn_number: Optional[int] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        turn_number: int | None = None,
+        metadata: dict[str, Any] | None = None,
     ):
         """Context manager for a timestep.
 
@@ -318,7 +324,7 @@ class SessionTracer:
         finally:
             await self.end_timestep(step_id)
 
-    async def get_session_history(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
+    async def get_session_history(self, limit: int | None = None) -> list[dict[str, Any]]:
         """Get recent session history from database."""
         if self.db is None:
             await self.initialize()

synth_ai/tracing_v3/storage/__init__.py

@@ -1,8 +1,8 @@
 """Storage abstraction layer for tracing v3."""
 
 from .base import TraceStorage
-from .factory import create_storage
 from .config import StorageConfig
+from .factory import create_storage
 from .types import EventType, MessageType, Provider
 
 __all__ = [

synth_ai/tracing_v3/storage/base.py

@@ -2,7 +2,8 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any
+
 import pandas as pd
 
 from ..abstractions import SessionTrace
@@ -29,7 +30,7 @@ class TraceStorage(ABC):
         pass
 
     @abstractmethod
-    async def get_session_trace(self, session_id: str) -> Optional[Dict[str, Any]]:
+    async def get_session_trace(self, session_id: str) -> dict[str, Any] | None:
         """Retrieve a session trace by ID.
 
         Args:
@@ -41,7 +42,7 @@ class TraceStorage(ABC):
         pass
 
     @abstractmethod
-    async def query_traces(self, query: str, params: Dict[str, Any] = None) -> pd.DataFrame:
+    async def query_traces(self, query: str, params: dict[str, Any] = None) -> pd.DataFrame:
         """Execute a query and return results as DataFrame.
 
         Args:
@@ -92,7 +93,7 @@ class TraceStorage(ABC):
         experiment_id: str,
         name: str,
         description: str = None,
-        configuration: Dict[str, Any] = None,
+        configuration: dict[str, Any] = None,
     ) -> str:
         """Create a new experiment."""
         raise NotImplementedError("Experiment management not supported by this backend")
@@ -103,14 +104,14 @@ class TraceStorage(ABC):
 
     async def get_sessions_by_experiment(
         self, experiment_id: str, limit: int = None
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Get all sessions for an experiment."""
         raise NotImplementedError("Experiment management not supported by this backend")
 
     # Batch operations
     async def batch_insert_sessions(
-        self, traces: List[SessionTrace], batch_size: int = 1000
-    ) -> List[str]:
+        self, traces: list[SessionTrace], batch_size: int = 1000
+    ) -> list[str]:
         """Batch insert multiple session traces.
 
         Default implementation calls insert_session_trace for each trace.

synth_ai/tracing_v3/storage/config.py

@@ -1,9 +1,9 @@
 """Storage configuration for tracing v3."""
 
-from dataclasses import dataclass
-from typing import Optional, Dict, Any
 import os
+from dataclasses import dataclass
 from enum import Enum
+from typing import Any
 
 
 class StorageBackend(str, Enum):
@@ -19,7 +19,7 @@ class StorageConfig:
     """Configuration for storage backend."""
 
     backend: StorageBackend = StorageBackend.TURSO
-    connection_string: Optional[str] = None
+    connection_string: str | None = None
 
     # Turso-specific settings
     turso_url: str = os.getenv("TURSO_DATABASE_URL", "sqlite+libsql://http://127.0.0.1:8080")
@@ -48,7 +48,7 @@ class StorageConfig:
         else:
             raise ValueError(f"Unknown backend: {self.backend}")
 
-    def get_backend_config(self) -> Dict[str, Any]:
+    def get_backend_config(self) -> dict[str, Any]:
         """Get backend-specific configuration."""
         if self.backend == StorageBackend.TURSO:
             config = {}

synth_ai/tracing_v3/storage/factory.py

@@ -1,12 +1,12 @@
 """Factory for creating storage instances."""
 
-from typing import Optional
-from .base import TraceStorage
-from .config import StorageConfig, StorageBackend
+
 from ..turso.manager import AsyncSQLTraceManager
+from .base import TraceStorage
+from .config import StorageBackend, StorageConfig
 
 
-def create_storage(config: Optional[StorageConfig] = None) -> TraceStorage:
+def create_storage(config: StorageConfig | None = None) -> TraceStorage:
     """Create a storage instance based on configuration.
 
     Args:

synth_ai/tracing_v3/storage/utils.py

@@ -1,10 +1,10 @@
 """Utility functions for storage layer."""
 
 import asyncio
-from typing import Any, Dict, List, Optional, TypeVar, Callable
 import functools
 import time
-
+from collections.abc import Callable
+from typing import Any, TypeVar
 
 T = TypeVar("T")
 
@@ -43,8 +43,8 @@ def retry_async(max_attempts: int = 3, delay: float = 1.0, backoff: float = 2.0)
 
 
 async def batch_process(
-    items: List[Any], processor: Callable, batch_size: int = 100, max_concurrent: int = 5
-) -> List[Any]:
+    items: list[Any], processor: Callable, batch_size: int = 100, max_concurrent: int = 5
+) -> list[Any]:
     """Process items in batches with concurrency control.
 
     Args:
@@ -88,7 +88,7 @@ def sanitize_json(data: Any) -> Any:
         return {k: sanitize_json(v) for k, v in data.items()}
     elif isinstance(data, list):
         return [sanitize_json(item) for item in data]
-    elif isinstance(data, (str, int, float, bool, type(None))):
+    elif isinstance(data, str | int | float | bool | type(None)):
         return data
     else:
         return str(data)
@@ -101,7 +101,7 @@ def estimate_size(obj: Any) -> int:
     """
     if isinstance(obj, str):
         return len(obj.encode("utf-8"))
-    elif isinstance(obj, (int, float)):
+    elif isinstance(obj, int | float):
         return 8
     elif isinstance(obj, bool):
         return 1
@@ -120,10 +120,10 @@ class StorageMetrics:
     """Track storage operation metrics."""
 
     def __init__(self):
-        self.operations: Dict[str, Dict[str, Any]] = {}
+        self.operations: dict[str, dict[str, Any]] = {}
 
     def record_operation(
-        self, operation: str, duration: float, success: bool, size: Optional[int] = None
+        self, operation: str, duration: float, success: bool, size: int | None = None
     ):
         """Record a storage operation."""
         if operation not in self.operations:
@@ -142,7 +142,7 @@ class StorageMetrics:
         if size:
             stats["total_size"] += size
 
-    def get_stats(self, operation: Optional[str] = None) -> Dict[str, Any]:
+    def get_stats(self, operation: str | None = None) -> dict[str, Any]:
         """Get statistics for operations."""
         if operation:
             stats = self.operations.get(operation, {})

synth_ai/tracing_v3/turso/__init__.py

@@ -3,11 +3,11 @@
 from .manager import AsyncSQLTraceManager
 from .models import (
     Base,
-    SessionTrace,
-    SessionTimestep,
     Event,
-    Message,
     Experiment,
+    Message,
+    SessionTimestep,
+    SessionTrace,
     System,
     SystemVersion,
 )

synth_ai/tracing_v3/turso/daemon.py

@@ -1,12 +1,12 @@
 """sqld daemon management utilities."""
 
-import subprocess
 import pathlib
 import shutil
-import sys
+import subprocess
 import time
+
 import requests
-from typing import Optional
+
 from ..config import CONFIG
 
 
@@ -15,23 +15,23 @@ class SqldDaemon:
 
     def __init__(self, db_path: str = None, http_port: int = None, binary_path: str = None):
         """Initialize sqld daemon manager.
-        
+
         Args:
             db_path: Path to database file (uses config default if not provided)
-            http_port: HTTP port for daemon (uses config default if not provided)  
+            http_port: HTTP port for daemon (uses config default if not provided)
             binary_path: Path to sqld binary (auto-detected if not provided)
         """
         self.db_path = db_path or CONFIG.sqld_db_path
         self.http_port = http_port or CONFIG.sqld_http_port
         self.binary_path = binary_path or self._find_binary()
-        self.process: Optional[subprocess.Popen] = None
+        self.process: subprocess.Popen | None = None
 
     def _find_binary(self) -> str:
         """Find sqld binary in PATH."""
         binary = shutil.which(CONFIG.sqld_binary) or shutil.which("libsql-server")
         if not binary:
             raise RuntimeError(
-                f"sqld binary not found in PATH. Install with: brew install turso-tech/tools/sqld"
+                "sqld binary not found in PATH. Install with: brew install turso-tech/tools/sqld"
             )
         return binary
 
@@ -117,7 +117,7 @@ class SqldDaemon:
 
 
 # Convenience functions
-_daemon: Optional[SqldDaemon] = None
+_daemon: SqldDaemon | None = None
 
 
 def start_sqld(db_path: str = None, port: int = None) -> SqldDaemon:
@@ -139,6 +139,6 @@ def stop_sqld():
         _daemon = None
 
 
-def get_daemon() -> Optional[SqldDaemon]:
+def get_daemon() -> SqldDaemon | None:
     """Get the global daemon instance."""
     return _daemon

synth_ai/tracing_v3/turso/manager.py

@@ -21,37 +21,43 @@ Performance Considerations:
 """
 
 import asyncio
+import logging
 from contextlib import asynccontextmanager
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any
+
 import pandas as pd
-from sqlalchemy import select, insert, update, delete, text, and_, or_, func
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine, AsyncSession
-from sqlalchemy.orm import sessionmaker, selectinload, joinedload
-from sqlalchemy.pool import NullPool
+from sqlalchemy import select, text, update
 from sqlalchemy.exc import IntegrityError
-import logging
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
+from sqlalchemy.orm import selectinload, sessionmaker
+from sqlalchemy.pool import NullPool
 
-from ..config import CONFIG
 from ..abstractions import (
-    SessionTrace,
-    SessionTimeStep,
-    BaseEvent,
-    LMCAISEvent,
     EnvironmentEvent,
+    LMCAISEvent,
     RuntimeEvent,
+    SessionTrace,
 )
-from ..utils import json_dumps
+from ..config import CONFIG
 from .models import (
     Base,
-    SessionTrace as DBSessionTrace,
-    SessionTimestep as DBSessionTimestep,
+    analytics_views,
+)
+from .models import (
     Event as DBEvent,
-    Message as DBMessage,
+)
+from .models import (
     Experiment as DBExperiment,
-    System as DBSystem,
-    SystemVersion as DBSystemVersion,
-    analytics_views,
+)
+from .models import (
+    Message as DBMessage,
+)
+from .models import (
+    SessionTimestep as DBSessionTimestep,
+)
+from .models import (
+    SessionTrace as DBSessionTrace,
 )
 
 logger = logging.getLogger(__name__)
@@ -59,10 +65,10 @@ logger = logging.getLogger(__name__)
 
 class AsyncSQLTraceManager:
     """Async trace storage manager using SQLAlchemy and Turso/sqld.
-    
+
     Handles all database operations for the tracing system. Designed to work
     with both local SQLite (via aiosqlite) and remote Turso databases.
-    
+
     The manager handles:
     - Connection lifecycle management
     - Schema creation and verification
@@ -71,22 +77,22 @@ class AsyncSQLTraceManager:
     - Analytics view creation
     """
 
-    def __init__(self, db_url: Optional[str] = None):
+    def __init__(self, db_url: str | None = None):
         self.db_url = db_url or CONFIG.db_url
-        self.engine: Optional[AsyncEngine] = None
-        self.SessionLocal: Optional[sessionmaker] = None
+        self.engine: AsyncEngine | None = None
+        self.SessionLocal: sessionmaker | None = None
         self._schema_lock = asyncio.Lock()
         self._schema_ready = False
 
     async def initialize(self):
         """Initialize the database connection and schema.
-        
+
         This method is idempotent and thread-safe. It:
         1. Creates the async engine with appropriate settings
         2. Verifies database file exists (for SQLite)
         3. Creates schema if needed
         4. Sets up analytics views
-        
+
         The schema lock ensures only one worker creates the schema in
         concurrent scenarios.
         """
@@ -136,7 +142,7 @@ class AsyncSQLTraceManager:
 
     async def _ensure_schema(self):
         """Ensure database schema is created.
-        
+
         Uses a lock to prevent race conditions when multiple workers start
         simultaneously. The checkfirst=True parameter handles cases where
         another worker already created the schema.
@@ -154,7 +160,7 @@ class AsyncSQLTraceManager:
                     await conn.run_sync(
                         lambda sync_conn: Base.metadata.create_all(sync_conn, checkfirst=True)
                     )
-                    #logger.info("✅ Database schema created/verified successfully")
+                    # logger.info("✅ Database schema created/verified successfully")
                 except Exception as e:
                     # If tables already exist, that's fine - another worker created them
                     if "already exists" not in str(e):
@@ -183,7 +189,7 @@ class AsyncSQLTraceManager:
                             logger.warning(f"Could not create view {view_name}: {e}")
 
             self._schema_ready = True
-            #logger.debug("🎯 Database ready for use!")
+            # logger.debug("🎯 Database ready for use!")
 
     @asynccontextmanager
     async def session(self):
@@ -195,18 +201,18 @@ class AsyncSQLTraceManager:
 
     async def insert_session_trace(self, trace: SessionTrace) -> str:
         """Insert a complete session trace.
-        
+
         This method handles the complex task of inserting a complete session
         with all its timesteps, events, and messages. It uses a single
         transaction for atomicity and flushes after timesteps to get their
         auto-generated IDs for foreign keys.
-        
+
         Args:
             trace: The complete session trace to store
-            
+
         Returns:
             The session ID
-            
+
         Raises:
             IntegrityError: If session ID already exists (handled gracefully)
         """
@@ -214,7 +220,7 @@ class AsyncSQLTraceManager:
             try:
                 # Convert to cents for cost storage - avoids floating point
                 # precision issues and allows for integer arithmetic
-                def to_cents(cost: Optional[float]) -> Optional[int]:
+                def to_cents(cost: float | None) -> int | None:
                     return int(cost * 100) if cost is not None else None
 
                 # Insert session
@@ -230,7 +236,7 @@ class AsyncSQLTraceManager:
 
                 # Track timestep IDs for foreign keys - we need these to link
                 # events and messages to their respective timesteps
-                step_id_map: Dict[str, int] = {}
+                step_id_map: dict[str, int] = {}
 
                 # Insert timesteps
                 for step in trace.session_time_steps:
@@ -270,8 +276,9 @@ class AsyncSQLTraceManager:
                         call_records_data = None
                         if event.call_records:
                             from dataclasses import asdict
+
                             call_records_data = [asdict(record) for record in event.call_records]
-                        
+
                         event_data.update(
                             {
                                 "event_type": "cais",
@@ -340,7 +347,7 @@ class AsyncSQLTraceManager:
                     return trace.session_id  # Return existing ID
                 raise
 
-    async def get_session_trace(self, session_id: str) -> Optional[Dict[str, Any]]:
+    async def get_session_trace(self, session_id: str) -> dict[str, Any] | None:
         """Retrieve a session trace by ID."""
         async with self.session() as sess:
             result = await sess.execute(
@@ -377,7 +384,9 @@ class AsyncSQLTraceManager:
                 ],
             }
 
-    async def query_traces(self, query: str, params: Optional[Dict[str, Any]] = None) -> pd.DataFrame:
+    async def query_traces(
+        self, query: str, params: dict[str, Any] | None = None
+    ) -> pd.DataFrame:
         """Execute a query and return results as DataFrame."""
         async with self.session() as sess:
             result = await sess.execute(text(query), params or {})
@@ -385,7 +394,10 @@ class AsyncSQLTraceManager:
             return pd.DataFrame(rows)
 
     async def get_model_usage(
-        self, start_date: Optional[datetime] = None, end_date: Optional[datetime] = None, model_name: Optional[str] = None
+        self,
+        start_date: datetime | None = None,
+        end_date: datetime | None = None,
+        model_name: str | None = None,
     ) -> pd.DataFrame:
         """Get model usage statistics."""
         query = """
@@ -414,8 +426,8 @@ class AsyncSQLTraceManager:
         self,
         experiment_id: str,
         name: str,
-        description: Optional[str] = None,
-        configuration: Optional[Dict[str, Any]] = None,
+        description: str | None = None,
+        configuration: dict[str, Any] | None = None,
     ) -> str:
         """Create a new experiment."""
         async with self.session() as sess:
@@ -440,18 +452,18 @@ class AsyncSQLTraceManager:
             await sess.commit()
 
     async def batch_insert_sessions(
-        self, traces: List[SessionTrace], batch_size: Optional[int] = None
-    ) -> List[str]:
+        self, traces: list[SessionTrace], batch_size: int | None = None
+    ) -> list[str]:
         """Batch insert multiple session traces.
-        
+
         Processes traces in batches to balance memory usage and performance.
         Each batch is inserted in a separate transaction to avoid holding
         locks for too long.
-        
+
         Args:
             traces: List of session traces to insert
             batch_size: Number of traces per batch (defaults to config)
-            
+
         Returns:
             List of inserted session IDs
         """
@@ -470,8 +482,8 @@ class AsyncSQLTraceManager:
         return inserted_ids
 
     async def get_sessions_by_experiment(
-        self, experiment_id: str, limit: Optional[int] = None
-    ) -> List[Dict[str, Any]]:
+        self, experiment_id: str, limit: int | None = None
+    ) -> list[dict[str, Any]]:
         """Get all sessions for an experiment."""
         async with self.session() as sess:
             query = (
@@ -515,7 +527,7 @@ class AsyncSQLTraceManager:
 
     async def close(self):
         """Close the database connection.
-        
+
         Properly disposes of the engine and all connections. This is important
         for cleanup, especially with SQLite which can leave lock files.
         """