synth-ai 0.2.8.dev4__py3-none-any.whl → 0.2.23.dev3__py3-none-any.whl

synth_ai/tracing_v3/session_tracer.py

@@ -1,14 +1,17 @@
-from __future__ import annotations
 """Main SessionTracer class for tracing v3."""
 
+from __future__ import annotations
+
 import asyncio
+import json
 from contextlib import asynccontextmanager
-from datetime import datetime
+from datetime import UTC, datetime
 from typing import Any
 
 from .abstractions import (
     BaseEvent,
     SessionEventMarkovBlanketMessage,
+    SessionMessageContent,
     SessionTimeStep,
     SessionTrace,
     TimeRecord,
@@ -16,7 +19,9 @@ from .abstractions import (
 from .config import CONFIG
 from .decorators import set_session_id, set_session_tracer, set_turn_number
 from .hooks import GLOBAL_HOOKS, HookManager
-from .turso.manager import AsyncSQLTraceManager
+from .storage.base import TraceStorage
+from .storage.config import StorageConfig
+from .storage.factory import create_storage
 from .utils import generate_session_id
 
 
@@ -28,6 +33,8 @@ class SessionTracer:
         hooks: HookManager | None = None,
         db_url: str | None = None,
         auto_save: bool = True,
+        storage: TraceStorage | None = None,
+        storage_config: StorageConfig | None = None,
     ):
         """Initialize session tracer.
 
@@ -40,7 +47,8 @@ class SessionTracer:
         self._current_trace: SessionTrace | None = None
         self._lock = asyncio.Lock()
         self.db_url = db_url or CONFIG.db_url
-        self.db: AsyncSQLTraceManager | None = None
+        self._storage_config = storage_config
+        self.db: TraceStorage | None = storage
         self.auto_save = auto_save
         self._current_step: SessionTimeStep | None = None
 
@@ -57,7 +65,8 @@ class SessionTracer:
     async def initialize(self):
         """Initialize the database connection."""
         if self.db is None:
-            self.db = AsyncSQLTraceManager(self.db_url)
+            config = self._storage_config or StorageConfig(connection_string=self.db_url)
+            self.db = create_storage(config)
             await self.db.initialize()
 
     async def start_session(
@@ -97,7 +106,7 @@ class SessionTracer:
 
             self._current_trace = SessionTrace(
                 session_id=session_id,
-                created_at=datetime.utcnow(),
+                created_at=datetime.now(UTC),
                 session_time_steps=[],
                 event_history=[],
                 markov_blanket_message_history=[],
@@ -110,7 +119,9 @@ class SessionTracer:
 
             # 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 {})
+                await self.db.ensure_session(
+                    session_id, created_at=self._current_trace.created_at, metadata=metadata or {}
+                )
 
             # Trigger hooks
             await self.hooks.trigger(
@@ -141,7 +152,7 @@ class SessionTracer:
         step = SessionTimeStep(
             step_id=step_id,
             step_index=len(self._current_trace.session_time_steps),
-            timestamp=datetime.utcnow(),
+            timestamp=datetime.now(UTC),
             turn_number=turn_number,
             step_metadata=metadata or {},
         )
@@ -186,7 +197,7 @@ class SessionTracer:
             step = self._current_step
 
         if step and step.completed_at is None:
-            step.completed_at = datetime.utcnow()
+            step.completed_at = datetime.now(UTC)
 
             # Trigger hooks
             await self.hooks.trigger(
@@ -234,7 +245,7 @@ class SessionTracer:
             event_id = await self.db.insert_event_row(
                 self._current_trace.session_id,
                 timestep_db_id=timestep_db_id,
-                event=event,
+                event=event,  # type: ignore[arg-type]
             )
             # Auto-insert an event reward if EnvironmentEvent carries reward
             try:
@@ -259,7 +270,7 @@ class SessionTracer:
 
     async def record_message(
         self,
-        content: str,
+        content: Any,
         message_type: str,
         event_time: float | None = None,
         message_time: int | None = None,
@@ -277,11 +288,13 @@ class SessionTracer:
         if self._current_trace is None:
             raise RuntimeError("No active session")
 
+        normalised_content, content_str = self._normalise_message_content(content)
+
         msg = SessionEventMarkovBlanketMessage(
-            content=content,
+            content=normalised_content,
             message_type=message_type,
             time_record=TimeRecord(
-                event_time=event_time or datetime.utcnow().timestamp(), message_time=message_time
+                event_time=event_time or datetime.now(UTC).timestamp(), message_time=message_time
             ),
             metadata=metadata or {},
         )
@@ -315,7 +328,7 @@ class SessionTracer:
                 self._current_trace.session_id,
                 timestep_db_id=timestep_db_id,
                 message_type=message_type,
-                content=content,
+                content=content_str,
                 event_time=msg.time_record.event_time,
                 message_time=msg.time_record.message_time,
                 metadata=msg.metadata,
@@ -323,7 +336,23 @@ class SessionTracer:
             return message_id
         return None
 
-    async def end_session(self, save: bool = None) -> SessionTrace:
+    @staticmethod
+    def _normalise_message_content(content: Any) -> tuple[SessionMessageContent, str]:
+        if isinstance(content, SessionMessageContent):
+            return content, content.as_text()
+        if isinstance(content, str):
+            payload = SessionMessageContent(text=content)
+            return payload, payload.as_text()
+        try:
+            serialized = json.dumps(content, ensure_ascii=False)
+            payload = SessionMessageContent(json_payload=serialized)
+            return payload, serialized
+        except (TypeError, ValueError):
+            text = str(content)
+            payload = SessionMessageContent(text=text)
+            return payload, text
+
+    async def end_session(self, save: bool | None = None) -> SessionTrace:
         """End the current session.
 
         Args:
@@ -339,18 +368,28 @@ class SessionTracer:
             # End any open timesteps
             for step in self._current_trace.session_time_steps:
                 if step.completed_at is None:
-                    step.completed_at = datetime.utcnow()
+                    step.completed_at = datetime.now(UTC)
 
             # Trigger pre-save hooks
             await self.hooks.trigger("before_save", session=self._current_trace)
 
             # Save if requested
             should_save = save if save is not None else self.auto_save
+            
+            # Debug logging
+            import logging
+            _logger = logging.getLogger(__name__)
+            _logger.info(f"[TRACE_DEBUG] end_session: should_save={should_save}, self.db={self.db is not None}, auto_save={self.auto_save}")
+            
             if should_save and self.db:
+                _logger.info(f"[TRACE_DEBUG] Calling insert_session_trace with {len(self._current_trace.markov_blanket_message_history)} messages")
                 await self.db.insert_session_trace(self._current_trace)
+                _logger.info("[TRACE_DEBUG] insert_session_trace completed")
 
                 # Trigger post-save hooks
                 await self.hooks.trigger("after_save", session=self._current_trace)
+            else:
+                _logger.warning(f"[TRACE_DEBUG] Skipping save: should_save={should_save}, self.db={self.db is not None}")
 
             # Trigger session end hooks
             await self.hooks.trigger("session_end", session=self._current_trace)
@@ -370,7 +409,7 @@ class SessionTracer:
         self,
         session_id: str | None = None,
         metadata: dict[str, Any] | None = None,
-        save: bool = None,
+        save: bool | None = None,
     ):
         """Context manager for a session.
 
@@ -414,8 +453,16 @@ class SessionTracer:
         if limit:
             query += f" LIMIT {limit}"
 
-        df = await self.db.query_traces(query)
-        return df.to_dict("records")
+        # Ensure DB initialized before querying
+        if self.db is None:
+            await self.initialize()
+        df_or_records = await self.db.query_traces(query)  # type: ignore[union-attr]
+        try:
+            # If pandas DataFrame
+            return df_or_records.to_dict("records")  # type: ignore[call-arg, attr-defined]
+        except AttributeError:
+            # Already list of dicts
+            return df_or_records
 
     async def close(self):
         """Close database connections."""
@@ -427,7 +474,14 @@ class SessionTracer:
     # Reward recording helpers
     # -------------------------------
 
-    async def record_outcome_reward(self, *, total_reward: int, achievements_count: int, total_steps: int, reward_metadata: dict[str, Any] | None = None) -> int | None:
+    async def record_outcome_reward(
+        self,
+        *,
+        total_reward: int,
+        achievements_count: int,
+        total_steps: int,
+        reward_metadata: dict[str, Any] | None = None,
+    ) -> int | None:
         """Record an episode-level outcome reward for the current session."""
         if self._current_trace is None:
             raise RuntimeError("No active session")
@@ -454,7 +508,18 @@ class SessionTracer:
 
     # 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:
+    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")

synth_ai/tracing_v3/storage/base.py

@@ -4,8 +4,6 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import Any
 
-import pandas as pd
-
 from ..abstractions import SessionTrace
 
 
@@ -42,22 +40,25 @@ class TraceStorage(ABC):
         pass
 
     @abstractmethod
-    async def query_traces(self, query: str, params: dict[str, Any] = None) -> pd.DataFrame:
-        """Execute a query and return results as DataFrame.
+    async def query_traces(self, query: str, params: dict[str, Any] | None = None) -> Any:
+        """Execute a query and return results.
 
         Args:
             query: The SQL query to execute
             params: Optional query parameters
 
         Returns:
-            Query results as a DataFrame
+            Query results as a DataFrame-like object or list of dict records
         """
         pass
 
     @abstractmethod
     async def get_model_usage(
-        self, start_date: datetime = None, end_date: datetime = None, model_name: str = None
-    ) -> pd.DataFrame:
+        self,
+        start_date: datetime | None = None,
+        end_date: datetime | None = None,
+        model_name: str | None = None,
+    ) -> Any:
         """Get model usage statistics.
 
         Args:
@@ -66,7 +67,7 @@ class TraceStorage(ABC):
             model_name: Optional model name filter
 
         Returns:
-            Model usage statistics as a DataFrame
+            Model usage statistics as a DataFrame-like object or list of dict records
         """
         pass
 
@@ -87,13 +88,98 @@ class TraceStorage(ABC):
         """Close the storage connection."""
         pass
 
+    # Incremental helpers -------------------------------------------------
+
+    @abstractmethod
+    async def ensure_session(
+        self,
+        session_id: str,
+        *,
+        created_at: datetime | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """Ensure a session row exists for the given session id."""
+        pass
+
+    @abstractmethod
+    async def ensure_timestep(
+        self,
+        session_id: str,
+        *,
+        step_id: str,
+        step_index: int,
+        turn_number: int | None = None,
+        started_at: datetime | None = None,
+        completed_at: datetime | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> int:
+        """Ensure a timestep row exists and return its database id."""
+        pass
+
+    @abstractmethod
+    async def insert_event_row(
+        self,
+        session_id: str,
+        *,
+        timestep_db_id: int | None,
+        event: Any,
+        metadata_override: dict[str, Any] | None = None,
+    ) -> int:
+        """Insert an event and return its database id."""
+        pass
+
+    @abstractmethod
+    async def insert_message_row(
+        self,
+        session_id: str,
+        *,
+        timestep_db_id: int | None,
+        message_type: str,
+        content: Any,
+        event_time: float | None = None,
+        message_time: int | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> int:
+        """Insert a message row linked to a session/timestep."""
+        pass
+
+    @abstractmethod
+    async def insert_outcome_reward(
+        self,
+        session_id: str,
+        *,
+        total_reward: int,
+        achievements_count: int,
+        total_steps: int,
+        reward_metadata: dict | None = None,
+    ) -> int:
+        """Record an outcome reward for a session."""
+        pass
+
+    @abstractmethod
+    async def insert_event_reward(
+        self,
+        session_id: str,
+        *,
+        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:
+        """Record a reward tied to a specific event."""
+        pass
+
     # Optional experiment management methods
     async def create_experiment(
         self,
         experiment_id: str,
         name: str,
-        description: str = None,
-        configuration: dict[str, Any] = None,
+        description: str | None = None,
+        configuration: dict[str, Any] | None = None,
     ) -> str:
         """Create a new experiment."""
         raise NotImplementedError("Experiment management not supported by this backend")
@@ -103,14 +189,14 @@ class TraceStorage(ABC):
         raise NotImplementedError("Experiment management not supported by this backend")
 
     async def get_sessions_by_experiment(
-        self, experiment_id: str, limit: int = None
+        self, experiment_id: str, limit: int | None = None
     ) -> 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
+        self, traces: list[SessionTrace], batch_size: int | None = 1000
     ) -> list[str]:
         """Batch insert multiple session traces.
 

synth_ai/tracing_v3/storage/config.py

@@ -1,29 +1,34 @@
 """Storage configuration for tracing v3."""
 
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any
 
+from ..config import resolve_trace_db_auth_token, resolve_trace_db_settings
+
 
 class StorageBackend(str, Enum):
     """Supported storage backends."""
 
-    TURSO = "turso"
+    TURSO_NATIVE = "turso_native"
     SQLITE = "sqlite"
     POSTGRES = "postgres"  # Future support
 
 
+def _is_enabled(value: str | None) -> bool:
+    if value is None:
+        return False
+    return value.lower() in {"1", "true", "yes", "on"}
+
+
 @dataclass
 class StorageConfig:
     """Configuration for storage backend."""
 
-    backend: StorageBackend = StorageBackend.TURSO
     connection_string: str | None = None
-
-    # Turso-specific settings
-    turso_url: str = os.getenv("TURSO_DATABASE_URL", "sqlite+libsql://http://127.0.0.1:8080")
-    turso_auth_token: str = os.getenv("TURSO_AUTH_TOKEN", "")
+    backend: StorageBackend | None = None
+    turso_auth_token: str | None = field(default=None)
 
     # Common settings
     pool_size: int = int(os.getenv("STORAGE_POOL_SIZE", "8"))
@@ -34,23 +39,65 @@ class StorageConfig:
     enable_compression: bool = os.getenv("STORAGE_COMPRESSION", "false").lower() == "true"
     max_content_length: int = int(os.getenv("STORAGE_MAX_CONTENT_LENGTH", "1000000"))  # 1MB
 
+    def __post_init__(self):
+        # Allow legacy override while keeping compatibility with existing TURSO_NATIVE env flag
+        native_env = os.getenv("TURSO_NATIVE")
+        native_flag = _is_enabled(native_env) if native_env is not None else None
+        resolved_url: str | None = self.connection_string
+        resolved_token: str | None = self.turso_auth_token
+
+        if resolved_url is None:
+            resolved_url, inferred_token = resolve_trace_db_settings()
+            self.connection_string = resolved_url
+            resolved_token = inferred_token
+
+        if resolved_token is None:
+            resolved_token = resolve_trace_db_auth_token()
+
+        self.turso_auth_token = resolved_token or ""
+
+        if self.backend is None:
+            self.backend = self._infer_backend(self.connection_string or "")
+
+        if native_flag is False:
+            raise RuntimeError("TURSO_NATIVE=false is no longer supported; only Turso/libSQL backend is available.")
+
+        # Allow both TURSO_NATIVE and SQLITE backends (both use libsql.connect)
+        if self.backend not in (StorageBackend.TURSO_NATIVE, StorageBackend.SQLITE):
+            raise RuntimeError(f"Unsupported backend: {self.backend}. Only Turso/libSQL and SQLite are supported.")
+
+    @staticmethod
+    def _infer_backend(connection_string: str) -> StorageBackend:
+        """Infer backend type from the connection string."""
+        scheme = connection_string.split(":", 1)[0].lower()
+        
+        # Plain SQLite files: file://, /absolute/path, or no scheme
+        if (
+            scheme == "file"
+            or scheme.startswith("sqlite")
+            or connection_string.startswith("/")
+            or "://" not in connection_string
+        ):
+            return StorageBackend.SQLITE
+        
+        # Turso/sqld: libsql://, http://, https://
+        if scheme.startswith("libsql") or "libsql" in scheme or scheme in ("http", "https"):
+            return StorageBackend.TURSO_NATIVE
+        
+        raise RuntimeError(f"Unsupported tracing backend scheme: {scheme}")
+
     def get_connection_string(self) -> str:
         """Get the appropriate connection string for the backend."""
         if self.connection_string:
             return self.connection_string
 
-        if self.backend == StorageBackend.TURSO:
-            return self.turso_url
-        elif self.backend == StorageBackend.SQLITE:
-            return "sqlite+aiosqlite:///traces.db"
-        elif self.backend == StorageBackend.POSTGRES:
-            return os.getenv("POSTGRES_URL", "postgresql+asyncpg://localhost/traces")
-        else:
-            raise ValueError(f"Unknown backend: {self.backend}")
+        if self.backend == StorageBackend.TURSO_NATIVE:
+            return self.connection_string or ""
+        raise ValueError(f"Unsupported backend: {self.backend}")
 
     def get_backend_config(self) -> dict[str, Any]:
         """Get backend-specific configuration."""
-        if self.backend == StorageBackend.TURSO:
+        if self.backend == StorageBackend.TURSO_NATIVE:
             config = {}
             if self.turso_auth_token:
                 config["auth_token"] = self.turso_auth_token

synth_ai/tracing_v3/storage/factory.py

@@ -1,7 +1,6 @@
 """Factory for creating storage instances."""
 
-
-from ..turso.manager import AsyncSQLTraceManager
+from ..turso.native_manager import NativeLibsqlTraceManager
 from .base import TraceStorage
 from .config import StorageBackend, StorageConfig
 
@@ -23,13 +22,16 @@ def create_storage(config: StorageConfig | None = None) -> TraceStorage:
 
         config = STORAGE_CONFIG
 
-    if config.backend == StorageBackend.TURSO:
-        # Turso uses the AsyncSQLTraceManager
-        return AsyncSQLTraceManager(db_url=config.get_connection_string())
-    elif config.backend == StorageBackend.SQLITE:
-        # For pure SQLite, we can still use AsyncSQLTraceManager
-        # but with a file-based URL
-        return AsyncSQLTraceManager(db_url=config.get_connection_string())
+    connection_string = config.get_connection_string()
+
+    # Both TURSO_NATIVE and SQLITE use NativeLibsqlTraceManager
+    # because libsql.connect() handles both remote and local file databases
+    if config.backend in (StorageBackend.TURSO_NATIVE, StorageBackend.SQLITE):
+        backend_config = config.get_backend_config()
+        return NativeLibsqlTraceManager(
+            db_url=connection_string,
+            auth_token=backend_config.get("auth_token"),
+        )
     elif config.backend == StorageBackend.POSTGRES:
         # Future: PostgreSQL implementation
         raise NotImplementedError("PostgreSQL backend not yet implemented")

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,10 +18,10 @@ 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):
-            last_exception = None
+        async def wrapper(*args: Any, **kwargs: Any) -> T:
+            last_exception: Exception | None = None
             current_delay = delay
 
             for attempt in range(max_attempts):
@@ -35,7 +35,9 @@ def retry_async(max_attempts: int = 3, delay: float = 1.0, backoff: float = 2.0)
                     else:
                         raise
 
-            raise last_exception
+            if last_exception:
+                raise last_exception
+            raise RuntimeError("Retry logic failed without exception")
 
         return wrapper
 
@@ -169,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:
@@ -183,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: