synth-ai 0.2.9.dev7__py3-none-any.whl → 0.2.10__py3-none-any.whl

synth_ai/tracing_v3/session_tracer.py

@@ -1,15 +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,
@@ -17,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
 
 
@@ -29,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.
 
@@ -41,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
 
@@ -58,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(
@@ -98,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=[],
@@ -144,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 {},
         )
@@ -189,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(
@@ -262,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,
@@ -280,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 {},
         )
@@ -318,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,
@@ -326,6 +336,22 @@ class SessionTracer:
             return message_id
         return None
 
+    @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.
 
@@ -342,7 +368,7 @@ 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)

synth_ai/tracing_v3/storage/base.py

@@ -88,6 +88,91 @@ 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,

synth_ai/tracing_v3/storage/config.py

@@ -9,16 +9,22 @@ from typing import Any
 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
+    backend: StorageBackend = StorageBackend.TURSO_NATIVE
     connection_string: str | None = None
 
     # Turso-specific settings
@@ -34,23 +40,30 @@ 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
+
+        if native_flag is False:
+            self.backend = StorageBackend.SQLITE
+
     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:
+        if self.backend == StorageBackend.TURSO_NATIVE:
             return self.turso_url
-        elif self.backend == StorageBackend.SQLITE:
+        if self.backend == StorageBackend.SQLITE:
             return "sqlite+aiosqlite:///traces.db"
-        elif self.backend == StorageBackend.POSTGRES:
+        if self.backend == StorageBackend.POSTGRES:
             return os.getenv("POSTGRES_URL", "postgresql+asyncpg://localhost/traces")
-        else:
-            raise ValueError(f"Unknown backend: {self.backend}")
+        raise ValueError(f"Unknown 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,6 +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
 
@@ -22,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())
+    connection_string = config.get_connection_string()
+
+    if config.backend == StorageBackend.TURSO_NATIVE:
+        backend_config = config.get_backend_config()
+        return NativeLibsqlTraceManager(
+            db_url=connection_string,
+            auth_token=backend_config.get("auth_token"),
+        )
     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())
+        return NativeLibsqlTraceManager(db_url=connection_string)
     elif config.backend == StorageBackend.POSTGRES:
         # Future: PostgreSQL implementation
         raise NotImplementedError("PostgreSQL backend not yet implemented")

synth_ai/tracing_v3/storage/utils.py

@@ -21,7 +21,7 @@ def retry_async(max_attempts: int = 3, delay: float = 1.0, backoff: float = 2.0)
     def decorator(func: Callable[..., T]) -> Callable[..., T]:
         @functools.wraps(func)
         async def wrapper(*args, **kwargs):
-            last_exception = None
+            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
 

synth_ai/tracing_v3/turso/daemon.py

@@ -13,7 +13,12 @@ from ..config import CONFIG
 class SqldDaemon:
     """Manages local sqld daemon lifecycle."""
 
-    def __init__(self, db_path: str = None, http_port: int = None, binary_path: str = None):
+    def __init__(
+        self,
+        db_path: str | None = None,
+        http_port: int | None = None,
+        binary_path: str | None = None,
+    ):
         """Initialize sqld daemon manager.
 
         Args:
@@ -120,7 +125,7 @@ class SqldDaemon:
 _daemon: SqldDaemon | None = None
 
 
-def start_sqld(db_path: str = None, port: int = None) -> SqldDaemon:
+def start_sqld(db_path: str | None = None, port: int | None = None) -> SqldDaemon:
     """Start a global sqld daemon instance."""
     global _daemon
     if _daemon and _daemon.is_running():

synth_ai/tracing_v3/turso/models.py

@@ -1,7 +1,7 @@
-from __future__ import annotations
-
 """SQLAlchemy declarative models for tracing v3."""
 
+from __future__ import annotations
+
 import json
 
 from sqlalchemy import (