synth-ai 0.2.4.dev6__py3-none-any.whl → 0.2.4.dev8__py3-none-any.whl

synth_ai/tracing_v3/turso/manager.py

@@ -21,37 +21,50 @@ 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 import event
+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,
+)
+from .models import (
+    OutcomeReward as DBOutcomeReward,
+)
+from .models import (
+    EventReward as DBEventReward,
 )
 
 logger = logging.getLogger(__name__)
@@ -59,10 +72,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 +84,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.
         """
@@ -119,6 +132,18 @@ class AsyncSQLTraceManager:
                     connect_args=connect_args,
                     echo=CONFIG.echo_sql,
                 )
+                # Ensure PRAGMA foreign_keys=ON for every connection
+                try:
+                    @event.listens_for(self.engine.sync_engine, "connect")
+                    def _set_sqlite_pragma(dbapi_connection, connection_record):  # type: ignore[no-redef]
+                        try:
+                            cursor = dbapi_connection.cursor()
+                            cursor.execute("PRAGMA foreign_keys=ON")
+                            cursor.close()
+                        except Exception:
+                            pass
+                except Exception:
+                    pass
             else:
                 connect_args = CONFIG.get_connect_args()
                 engine_kwargs = CONFIG.get_engine_kwargs()
@@ -136,7 +161,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 +179,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 +208,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 +220,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 +239,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 +255,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 +295,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 +366,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 +403,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 +413,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 +445,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 +471,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 +501,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 +546,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.
         """
@@ -526,3 +557,202 @@ class AsyncSQLTraceManager:
             self.engine = None
             self.SessionLocal = None
             self._schema_ready = False
+
+    # -------------------------------
+    # Incremental insert helpers
+    # -------------------------------
+
+    async def ensure_session(self, session_id: str, *, created_at: datetime | None = None, metadata: dict[str, Any] | None = None):
+        """Ensure a DB session row exists for session_id."""
+        async with self.session() as sess:
+            result = await sess.execute(select(DBSessionTrace).where(DBSessionTrace.session_id == session_id))
+            existing = result.scalar_one_or_none()
+            if existing:
+                return
+            row = DBSessionTrace(
+                session_id=session_id,
+                created_at=created_at or datetime.utcnow(),
+                num_timesteps=0,
+                num_events=0,
+                num_messages=0,
+                session_metadata=metadata or {},
+            )
+            sess.add(row)
+            await sess.commit()
+
+    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; return its DB id."""
+        async with self.session() as sess:
+            result = await sess.execute(
+                select(DBSessionTimestep).where(DBSessionTimestep.session_id == session_id, DBSessionTimestep.step_id == step_id)
+            )
+            row = result.scalar_one_or_none()
+            if row:
+                return row.id
+            row = DBSessionTimestep(
+                session_id=session_id,
+                step_id=step_id,
+                step_index=step_index,
+                turn_number=turn_number,
+                started_at=started_at or datetime.utcnow(),
+                completed_at=completed_at,
+                num_events=0,
+                num_messages=0,
+                step_metadata=metadata or {},
+            )
+            sess.add(row)
+            await sess.flush()
+            # increment session num_timesteps
+            await sess.execute(
+                update(DBSessionTrace)
+                .where(DBSessionTrace.session_id == session_id)
+                .values(num_timesteps=DBSessionTrace.num_timesteps + 1)
+            )
+            await sess.commit()
+            return row.id
+
+    async def insert_message_row(self, session_id: str, *, timestep_db_id: int | None, message_type: str, content: str, event_time: float | None = None, message_time: int | None = None, metadata: dict[str, Any] | None = None) -> int:
+        """Insert a message and return its id."""
+        async with self.session() as sess:
+            db_msg = DBMessage(
+                session_id=session_id,
+                timestep_id=timestep_db_id,
+                message_type=message_type,
+                content=content,
+                event_time=event_time,
+                message_time=message_time,
+                message_metadata=metadata or {},
+            )
+            sess.add(db_msg)
+            await sess.flush()
+            # increment session num_messages
+            await sess.execute(
+                update(DBSessionTrace)
+                .where(DBSessionTrace.session_id == session_id)
+                .values(num_messages=DBSessionTrace.num_messages + 1)
+            )
+            await sess.commit()
+            return db_msg.id
+
+    async def insert_event_row(self, session_id: str, *, timestep_db_id: int | None, event: EnvironmentEvent | LMCAISEvent | RuntimeEvent, metadata_override: dict[str, Any] | None = None) -> int:
+        """Insert an event and return its id."""
+        def to_cents(cost: float | None) -> int | None:
+            return int(cost * 100) if cost is not None else None
+
+        event_data: dict[str, Any] = {
+            "session_id": session_id,
+            "timestep_id": timestep_db_id,
+            "system_instance_id": event.system_instance_id,
+            "event_time": event.time_record.event_time,
+            "message_time": event.time_record.message_time,
+            "event_metadata_json": metadata_override or event.metadata or {},
+            "event_extra_metadata": getattr(event, "event_metadata", None),
+        }
+        if isinstance(event, LMCAISEvent):
+            call_records_data = None
+            if getattr(event, "call_records", None):
+                from dataclasses import asdict
+
+                call_records_data = [asdict(record) for record in event.call_records]
+            event_data.update({
+                "event_type": "cais",
+                "model_name": event.model_name,
+                "provider": event.provider,
+                "input_tokens": event.input_tokens,
+                "output_tokens": event.output_tokens,
+                "total_tokens": event.total_tokens,
+                "cost_usd": to_cents(event.cost_usd),
+                "latency_ms": event.latency_ms,
+                "span_id": event.span_id,
+                "trace_id": event.trace_id,
+                "system_state_before": event.system_state_before,
+                "system_state_after": event.system_state_after,
+                "call_records": call_records_data,
+            })
+        elif isinstance(event, EnvironmentEvent):
+            event_data.update({
+                "event_type": "environment",
+                "reward": event.reward,
+                "terminated": event.terminated,
+                "truncated": event.truncated,
+                "system_state_before": event.system_state_before,
+                "system_state_after": event.system_state_after,
+            })
+        elif isinstance(event, RuntimeEvent):
+            event_data.update({
+                "event_type": "runtime",
+                "event_metadata_json": {**(event.metadata or {}), "actions": event.actions},
+            })
+        else:
+            event_data["event_type"] = event.__class__.__name__.lower()
+
+        async with self.session() as sess:
+            db_event = DBEvent(**event_data)
+            sess.add(db_event)
+            await sess.flush()
+            # increment session num_events
+            await sess.execute(
+                update(DBSessionTrace)
+                .where(DBSessionTrace.session_id == session_id)
+                .values(num_events=DBSessionTrace.num_events + 1)
+            )
+            await sess.commit()
+            return db_event.id
+
+    # -------------------------------
+    # Reward helpers
+    # -------------------------------
+
+    async def insert_outcome_reward(self, session_id: str, *, total_reward: int, achievements_count: int, total_steps: int) -> int:
+        async with self.session() as sess:
+            row = DBOutcomeReward(
+                session_id=session_id,
+                total_reward=total_reward,
+                achievements_count=achievements_count,
+                total_steps=total_steps,
+            )
+            sess.add(row)
+            await sess.flush()
+            await sess.commit()
+            return row.id
+
+    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:
+        async with self.session() as sess:
+            row = DBEventReward(
+                event_id=event_id,
+                session_id=session_id,
+                message_id=message_id,
+                turn_number=turn_number,
+                reward_value=reward_value,
+                reward_type=reward_type,
+                key=key,
+                annotation=annotation or {},
+                source=source,
+            )
+            sess.add(row)
+            await sess.flush()
+            await sess.commit()
+            return row.id
+
+    async def get_outcome_rewards(self) -> list[dict[str, Any]]:
+        async with self.session() as sess:
+            result = await sess.execute(select(DBOutcomeReward))
+            rows = result.scalars().all()
+            return [
+                {
+                    "id": r.id,
+                    "session_id": r.session_id,
+                    "total_reward": r.total_reward,
+                    "achievements_count": r.achievements_count,
+                    "total_steps": r.total_steps,
+                    "created_at": r.created_at,
+                }
+                for r in rows
+            ]
+
+    async def get_outcome_rewards_by_min_reward(self, min_reward: int) -> list[str]:
+        async with self.session() as sess:
+            result = await sess.execute(
+                select(DBOutcomeReward.session_id).where(DBOutcomeReward.total_reward >= min_reward)
+            )
+            return [row[0] for row in result.all()]

synth_ai/tracing_v3/turso/models.py

@@ -1,28 +1,25 @@
 """SQLAlchemy declarative models for tracing v3."""
 
-from datetime import datetime
-from typing import Optional
+import json
+
 from sqlalchemy import (
+    Boolean,
+    CheckConstraint,
     Column,
-    Integer,
-    String,
     DateTime,
-    JSON,
+    Float,
     ForeignKey,
     Index,
+    Integer,
+    String,
     Text,
-    Float,
-    Boolean,
-    UniqueConstraint,
-    CheckConstraint,
     TypeDecorator,
+    UniqueConstraint,
 )
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import relationship
 from sqlalchemy.sql import func
 from sqlalchemy.types import UserDefinedType
-import json
-
 
 Base = declarative_base()
 
@@ -81,10 +78,11 @@ class JSONText(TypeDecorator):
 
 class SessionTrace(Base):
     """Database model for session traces.
-    
+
     Stores high-level information about tracing sessions including
     metadata, statistics, and relationships to timesteps and events.
     """
+
     __tablename__ = "session_traces"
 
     session_id = Column(String, primary_key=True)
@@ -114,10 +112,11 @@ class SessionTrace(Base):
 
 class SessionTimestep(Base):
     """Database model for session timesteps.
-    
+
     Represents individual steps within a tracing session, with timing
     information and relationships to events and messages.
     """
+
     __tablename__ = "session_timesteps"
 
     id = Column(Integer, primary_key=True, autoincrement=True)
@@ -145,11 +144,12 @@ class SessionTimestep(Base):
 
 class Event(Base):
     """Database model for events.
-    
+
     Stores all types of events (LM CAIS, environment, runtime) with
     type-specific fields and common metadata. Supports vector embeddings
     for similarity search.
     """
+
     __tablename__ = "events"
 
     id = Column(Integer, primary_key=True, autoincrement=True)
@@ -209,10 +209,11 @@ class Event(Base):
 
 class Message(Base):
     """Database model for messages.
-    
+
     Stores conversational messages between users, assistants, and systems
     with support for embeddings and rich metadata.
     """
+
     __tablename__ = "messages"
 
     id = Column(Integer, primary_key=True, autoincrement=True)
@@ -247,10 +248,11 @@ class Message(Base):
 
 class Experiment(Base):
     """Database model for experiments.
-    
+
     Groups related sessions and systems for experimental evaluation
     and comparison. Supports rich configuration and metadata.
     """
+
     __tablename__ = "experiments"
 
     experiment_id = Column(String, primary_key=True)
@@ -277,10 +279,11 @@ class Experiment(Base):
 
 class System(Base):
     """Database model for systems.
-    
+
     Represents agents, environments, or runtime systems that participate
     in tracing sessions. Supports versioning and type classification.
     """
+
     __tablename__ = "systems"
 
     system_id = Column(String, primary_key=True)
@@ -302,10 +305,11 @@ class System(Base):
 
 class SystemVersion(Base):
     """Database model for system versions.
-    
+
     Tracks different versions of systems with commit hashes,
     configuration changes, and relationships to experiments.
     """
+
     __tablename__ = "system_versions"
 
     version_id = Column(String, primary_key=True)
@@ -329,10 +333,11 @@ class SystemVersion(Base):
 
 class ExperimentalSystem(Base):
     """Database model for experiment-system relationships.
-    
+
     Junction table linking experiments with specific system versions,
     allowing tracking of which systems participated in which experiments.
     """
+
     __tablename__ = "experimental_systems"
 
     id = Column(Integer, primary_key=True, autoincrement=True)
@@ -403,3 +408,56 @@ analytics_views = {
         GROUP BY e.experiment_id
     """,
 }
+
+
+# Reward persistence tables
+
+
+class OutcomeReward(Base):
+    """Episode-level rewards/outcomes per session.
+
+    Stores per-episode summary including total_reward (e.g., unique achievements),
+    achievements_count, and total_steps. Used for filtering episodes by outcome.
+    """
+
+    __tablename__ = "outcome_rewards"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    session_id = Column(String, ForeignKey("session_traces.session_id"), nullable=False)
+    total_reward = Column(Integer, nullable=False)
+    achievements_count = Column(Integer, nullable=False, default=0)
+    total_steps = Column(Integer, nullable=False, default=0)
+    created_at = Column(DateTime, default=func.current_timestamp(), nullable=False)
+
+    __table_args__ = (
+        Index("idx_outcome_rewards_session", "session_id"),
+        Index("idx_outcome_rewards_total", "total_reward"),
+    )
+
+
+class EventReward(Base):
+    """First-class event-level rewards with annotations.
+
+    Links to an event and session. `message_id` is optional.
+    """
+
+    __tablename__ = "event_rewards"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    event_id = Column(Integer, ForeignKey("events.id"), nullable=False)
+    session_id = Column(String, ForeignKey("session_traces.session_id"), nullable=False)
+    message_id = Column(Integer, ForeignKey("messages.id"), nullable=True)
+    turn_number = Column(Integer, nullable=True)
+    reward_value = Column(Float, nullable=False, default=0.0)
+    reward_type = Column(String, nullable=True)  # shaped | sparse | achievement | penalty | evaluator | human
+    key = Column(String, nullable=True)  # e.g., achievement name
+    annotation = Column(JSONText)  # free-form JSON
+    source = Column(String, nullable=True)  # environment | runner | evaluator | human
+    created_at = Column(DateTime, default=func.current_timestamp(), nullable=False)
+
+    __table_args__ = (
+        Index("idx_event_rewards_session", "session_id"),
+        Index("idx_event_rewards_event", "event_id"),
+        Index("idx_event_rewards_type", "reward_type"),
+        Index("idx_event_rewards_key", "key"),
+    )