synth-ai 0.2.4.dev6__py3-none-any.whl → 0.2.4.dev8__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
         """
@@ -101,6 +107,10 @@ class SessionTracer:
             if self.auto_save and self.db is None:
                 await self.initialize()
 
+            # Ensure session row exists for incremental writes
+            if self.db:
+                await self.db.ensure_session(session_id, created_at=self._current_trace.created_at, metadata=metadata or {})
+
             # Trigger hooks
             await self.hooks.trigger(
                 "session_start", session_id=session_id, metadata=metadata or {}
@@ -111,8 +121,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.
 
@@ -146,9 +156,20 @@ class SessionTracer:
             "timestep_start", step=step, session_id=self._current_trace.session_id
         )
 
+        # Ensure timestep row exists in DB for incremental linkage
+        if self.db:
+            await self.db.ensure_timestep(
+                self._current_trace.session_id,
+                step_id=step.step_id,
+                step_index=step.step_index,
+                turn_number=turn_number,
+                started_at=step.timestamp,
+                metadata=metadata or {},
+            )
+
         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")
@@ -174,7 +195,7 @@ class SessionTracer:
         if step == self._current_step:
             self._current_step = None
 
-    async def record_event(self, event: BaseEvent):
+    async def record_event(self, event: BaseEvent) -> int | None:
         """Record an event.
 
         Args:
@@ -195,14 +216,54 @@ class SessionTracer:
         if self._current_step:
             self._current_step.events.append(event)
 
+        # Persist incrementally if DB is available; return DB event id
+        if self.db:
+            timestep_db_id = None
+            if self._current_step:
+                # ensure timestep exists and get id
+                timestep_db_id = await self.db.ensure_timestep(
+                    self._current_trace.session_id,
+                    step_id=self._current_step.step_id,
+                    step_index=self._current_step.step_index,
+                    turn_number=self._current_step.turn_number,
+                    started_at=self._current_step.timestamp,
+                    completed_at=self._current_step.completed_at,
+                    metadata=self._current_step.step_metadata,
+                )
+            event_id = await self.db.insert_event_row(
+                self._current_trace.session_id,
+                timestep_db_id=timestep_db_id,
+                event=event,
+            )
+            # Auto-insert an event reward if EnvironmentEvent carries reward
+            try:
+                from .abstractions import EnvironmentEvent  # local import to avoid cycles
+
+                if isinstance(event, EnvironmentEvent) and event.reward is not None:
+                    await self.record_event_reward(
+                        event_id=event_id,
+                        message_id=None,
+                        turn_number=self._current_step.turn_number if self._current_step else None,
+                        reward_value=float(event.reward),
+                        reward_type="sparse",
+                        key=None,
+                        annotation=getattr(event, "event_metadata", None),
+                        source="environment",
+                    )
+            except Exception:
+                # Do not fail tracing if reward recording fails
+                pass
+            return event_id
+        return None
+
     async def record_message(
         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,
+    ) -> int | None:
         """Record a message.
 
         Args:
@@ -236,6 +297,31 @@ class SessionTracer:
         if self._current_step:
             self._current_step.markov_blanket_messages.append(msg)
 
+        # Persist incrementally and return DB message id
+        if self.db:
+            timestep_db_id = None
+            if self._current_step:
+                timestep_db_id = await self.db.ensure_timestep(
+                    self._current_trace.session_id,
+                    step_id=self._current_step.step_id,
+                    step_index=self._current_step.step_index,
+                    turn_number=self._current_step.turn_number,
+                    started_at=self._current_step.timestamp,
+                    completed_at=self._current_step.completed_at,
+                    metadata=self._current_step.step_metadata,
+                )
+            message_id = await self.db.insert_message_row(
+                self._current_trace.session_id,
+                timestep_db_id=timestep_db_id,
+                message_type=message_type,
+                content=content,
+                event_time=msg.time_record.event_time,
+                message_time=msg.time_record.message_time,
+                metadata=msg.metadata,
+            )
+            return message_id
+        return None
+
     async def end_session(self, save: bool = None) -> SessionTrace:
         """End the current session.
 
@@ -281,8 +367,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 +388,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 +404,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()
@@ -335,3 +421,44 @@ class SessionTracer:
         if self.db:
             await self.db.close()
             self.db = None
+
+    # -------------------------------
+    # Reward recording helpers
+    # -------------------------------
+
+    async def record_outcome_reward(self, *, total_reward: int, achievements_count: int, total_steps: int) -> int | None:
+        """Record an episode-level outcome reward for the current session."""
+        if self._current_trace is None:
+            raise RuntimeError("No active session")
+        if self.db is None:
+            await self.initialize()
+        if self.db:
+            return await self.db.insert_outcome_reward(
+                self._current_trace.session_id,
+                total_reward=total_reward,
+                achievements_count=achievements_count,
+                total_steps=total_steps,
+            )
+        return None
+
+    # StepMetrics removed in favor of event_rewards; use record_event_reward for per-turn shaped values
+
+    async def record_event_reward(self, *, event_id: int, message_id: int | None = None, turn_number: int | None = None, reward_value: float = 0.0, reward_type: str | None = None, key: str | None = None, annotation: dict[str, Any] | None = None, source: str | None = None) -> int | None:
+        """Record a first-class event-level reward with optional annotations."""
+        if self._current_trace is None:
+            raise RuntimeError("No active session")
+        if self.db is None:
+            await self.initialize()
+        if self.db:
+            return await self.db.insert_event_reward(
+                self._current_trace.session_id,
+                event_id=event_id,
+                message_id=message_id,
+                turn_number=turn_number,
+                reward_value=reward_value,
+                reward_type=reward_type,
+                key=key,
+                annotation=annotation,
+                source=source,
+            )
+        return None

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