synth-ai 0.2.9.dev11__py3-none-any.whl → 0.4.1__py3-none-any.whl

synth_ai/tracing_v3/turso/native_manager.py

@@ -1,1173 +0,0 @@
-"""LibSQL-native trace manager prototype.
-
-This module provides the Turso/libsql-backed trace storage implementation. It
-mirrors the public surface area of the historical SQLAlchemy manager while
-executing all operations directly via libsql.
-"""
-
-from __future__ import annotations
-
-import asyncio
-import json
-import logging
-import re
-from dataclasses import asdict, dataclass
-from datetime import UTC, datetime
-from typing import Any
-
-import libsql
-from sqlalchemy.engine import make_url
-
-try:  # pragma: no cover - exercised only when pandas present
-    import pandas as pd  # type: ignore
-except Exception:  # pragma: no cover
-    pd = None  # type: ignore[assignment]
-
-from ..abstractions import (
-    EnvironmentEvent,
-    LMCAISEvent,
-    RuntimeEvent,
-    SessionMessageContent,
-    SessionTrace,
-)
-from ..config import CONFIG
-from ..storage.base import TraceStorage
-from .models import analytics_views
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass(slots=True)
-class _ConnectionTarget:
-    """Resolved connection target for libsql."""
-
-    database: str
-    sync_url: str | None = None
-    auth_token: str | None = None
-
-
-def _resolve_connection_target(db_url: str | None, auth_token: str | None) -> _ConnectionTarget:
-    """Normalise the configured database URL."""
-    url = db_url or CONFIG.db_url
-
-    # Fast-path local SQLite URLs (`sqlite+aiosqlite:///path/to/db`)
-    if url.startswith("sqlite+aiosqlite:///"):
-        return _ConnectionTarget(database=url.replace("sqlite+aiosqlite:///", ""), auth_token=auth_token)
-
-    # SQLAlchemy-compatible libsql scheme (`sqlite+libsql://<endpoint or path>`)
-    if url.startswith("sqlite+libsql://"):
-        target = url.replace("sqlite+libsql://", "", 1)
-        return _ConnectionTarget(database=target, sync_url=target if target.startswith("libsql://") else None, auth_token=auth_token)
-
-    # Native libsql URLs (`libsql://...`).
-    if url.startswith("libsql://"):
-        return _ConnectionTarget(database=url, sync_url=url, auth_token=auth_token)
-
-    # Fallback to SQLAlchemy URL parsing for anything else we missed.
-    try:
-        parsed = make_url(url)
-        if parsed.drivername.startswith("sqlite"):
-            if parsed.database:
-                return _ConnectionTarget(database=parsed.database, auth_token=auth_token)
-        if parsed.drivername.startswith("libsql"):
-            database = parsed.render_as_string(hide_password=False)
-            return _ConnectionTarget(database=database, sync_url=database, auth_token=auth_token)
-    except Exception:  # pragma: no cover - defensive guardrail
-        logger.debug("Unable to parse db_url via SQLAlchemy", exc_info=True)
-
-    # As a last resort use the raw value (libsql.connect can handle absolute paths).
-    return _ConnectionTarget(database=url, auth_token=auth_token)
-
-
-def _json_dumps(value: Any) -> str | None:
-    """Serialise Python objects as JSON compatible with the existing schema."""
-
-    def _default(obj: Any):
-        if isinstance(obj, datetime):
-            return obj.isoformat()
-        return str(obj)
-
-    if value is None:
-        return None
-    return json.dumps(value, separators=(",", ":"), default=_default)
-
-
-def _maybe_datetime(value: Any) -> Any:
-    if value is None or isinstance(value, datetime):
-        return value
-    if isinstance(value, str):
-        try:
-            return datetime.fromisoformat(value)
-        except ValueError:
-            pass
-    return value
-
-
-def _load_json(value: Any) -> Any:
-    if value is None or isinstance(value, (dict, list)):
-        return value or {}
-    if isinstance(value, str):
-        try:
-            return json.loads(value)
-        except (TypeError, ValueError):
-            return {}
-    return value
-
-
-_TABLE_DEFINITIONS: tuple[str, ...] = (
-    """
-    CREATE TABLE IF NOT EXISTS experiments (
-        experiment_id VARCHAR PRIMARY KEY,
-        name VARCHAR NOT NULL,
-        description TEXT,
-        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
-        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
-        configuration TEXT,
-        metadata TEXT
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS systems (
-        system_id VARCHAR PRIMARY KEY,
-        name VARCHAR NOT NULL,
-        system_type VARCHAR,
-        description TEXT,
-        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
-        metadata TEXT
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS system_versions (
-        version_id VARCHAR PRIMARY KEY,
-        system_id VARCHAR NOT NULL,
-        version_number VARCHAR NOT NULL,
-        commit_hash VARCHAR,
-        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
-        configuration TEXT,
-        metadata TEXT,
-        FOREIGN KEY(system_id) REFERENCES systems(system_id),
-        UNIQUE(system_id, version_number)
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS experimental_systems (
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        experiment_id VARCHAR NOT NULL,
-        system_id VARCHAR NOT NULL,
-        version_id VARCHAR NOT NULL,
-        FOREIGN KEY(experiment_id) REFERENCES experiments(experiment_id),
-        FOREIGN KEY(system_id) REFERENCES systems(system_id),
-        FOREIGN KEY(version_id) REFERENCES system_versions(version_id)
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS session_traces (
-        session_id VARCHAR PRIMARY KEY,
-        created_at DATETIME NOT NULL,
-        num_timesteps INTEGER NOT NULL,
-        num_events INTEGER NOT NULL,
-        num_messages INTEGER NOT NULL,
-        metadata TEXT,
-        experiment_id VARCHAR,
-        embedding VECTOR,
-        FOREIGN KEY(experiment_id) REFERENCES experiments(experiment_id)
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS session_timesteps (
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        session_id VARCHAR NOT NULL,
-        step_id VARCHAR NOT NULL,
-        step_index INTEGER NOT NULL,
-        turn_number INTEGER,
-        started_at DATETIME,
-        completed_at DATETIME,
-        num_events INTEGER,
-        num_messages INTEGER,
-        step_metadata TEXT,
-        UNIQUE(session_id, step_id),
-        FOREIGN KEY(session_id) REFERENCES session_traces(session_id)
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS events (
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        session_id VARCHAR NOT NULL,
-        timestep_id INTEGER,
-        event_type VARCHAR NOT NULL,
-        system_instance_id VARCHAR,
-        event_time FLOAT,
-        message_time INTEGER,
-        created_at DATETIME,
-        model_name VARCHAR,
-        provider VARCHAR,
-        input_tokens INTEGER,
-        output_tokens INTEGER,
-        total_tokens INTEGER,
-        cost_usd INTEGER,
-        latency_ms INTEGER,
-        span_id VARCHAR,
-        trace_id VARCHAR,
-        call_records TEXT,
-        reward FLOAT,
-        terminated BOOLEAN,
-        truncated BOOLEAN,
-        system_state_before TEXT,
-        system_state_after TEXT,
-        metadata TEXT,
-        event_metadata TEXT,
-        embedding VECTOR,
-        CHECK (event_type IN ('cais', 'environment', 'runtime')),
-        FOREIGN KEY(session_id) REFERENCES session_traces(session_id),
-        FOREIGN KEY(timestep_id) REFERENCES session_timesteps(id)
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS messages (
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        session_id VARCHAR NOT NULL,
-        timestep_id INTEGER,
-        message_type VARCHAR NOT NULL,
-        content TEXT NOT NULL,
-        timestamp DATETIME,
-        event_time FLOAT,
-        message_time INTEGER,
-        metadata TEXT,
-        embedding VECTOR,
-        CHECK (message_type IN ('user', 'assistant', 'system', 'tool_use', 'tool_result')),
-        FOREIGN KEY(session_id) REFERENCES session_traces(session_id),
-        FOREIGN KEY(timestep_id) REFERENCES session_timesteps(id)
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS outcome_rewards (
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        session_id VARCHAR NOT NULL,
-        total_reward INTEGER NOT NULL,
-        achievements_count INTEGER NOT NULL,
-        total_steps INTEGER NOT NULL,
-        created_at DATETIME NOT NULL,
-        reward_metadata TEXT,
-        FOREIGN KEY(session_id) REFERENCES session_traces(session_id)
-    )
-    """,
-    """
-    CREATE TABLE IF NOT EXISTS event_rewards (
-        id INTEGER PRIMARY KEY AUTOINCREMENT,
-        event_id INTEGER NOT NULL,
-        session_id VARCHAR NOT NULL,
-        message_id INTEGER,
-        turn_number INTEGER,
-        reward_value FLOAT NOT NULL,
-        reward_type VARCHAR,
-        "key" VARCHAR,
-        annotation TEXT,
-        source VARCHAR,
-        created_at DATETIME NOT NULL,
-        FOREIGN KEY(event_id) REFERENCES events(id),
-        FOREIGN KEY(session_id) REFERENCES session_traces(session_id),
-        FOREIGN KEY(message_id) REFERENCES messages(id)
-    )
-    """
-)
-
-
-_INDEX_DEFINITIONS: tuple[str, ...] = (
-    "CREATE INDEX IF NOT EXISTS idx_session_created ON session_traces (created_at)",
-    "CREATE INDEX IF NOT EXISTS idx_session_experiment ON session_traces (experiment_id)",
-    "CREATE INDEX IF NOT EXISTS idx_timestep_session_step ON session_timesteps (session_id, step_id)",
-    "CREATE INDEX IF NOT EXISTS idx_timestep_turn ON session_timesteps (turn_number)",
-    "CREATE INDEX IF NOT EXISTS idx_event_session_step ON events (session_id, timestep_id)",
-    "CREATE INDEX IF NOT EXISTS idx_event_type ON events (event_type)",
-    "CREATE INDEX IF NOT EXISTS idx_event_created ON events (created_at)",
-    "CREATE INDEX IF NOT EXISTS idx_event_model ON events (model_name)",
-    "CREATE INDEX IF NOT EXISTS idx_event_trace ON events (trace_id)",
-    "CREATE INDEX IF NOT EXISTS idx_message_session_step ON messages (session_id, timestep_id)",
-    "CREATE INDEX IF NOT EXISTS idx_message_type ON messages (message_type)",
-    "CREATE INDEX IF NOT EXISTS idx_message_timestamp ON messages (timestamp)",
-    "CREATE INDEX IF NOT EXISTS idx_experiment_created ON experiments (created_at)",
-    "CREATE INDEX IF NOT EXISTS idx_experiment_name ON experiments (name)",
-    "CREATE INDEX IF NOT EXISTS idx_system_name ON systems (name)",
-    "CREATE INDEX IF NOT EXISTS idx_system_type ON systems (system_type)",
-    "CREATE UNIQUE INDEX IF NOT EXISTS uq_system_version ON system_versions (system_id, version_number)",
-    "CREATE INDEX IF NOT EXISTS idx_version_system ON system_versions (system_id)",
-    "CREATE INDEX IF NOT EXISTS idx_version_created ON system_versions (created_at)",
-    "CREATE UNIQUE INDEX IF NOT EXISTS uq_experiment_system ON experimental_systems (experiment_id, system_id)",
-    "CREATE INDEX IF NOT EXISTS idx_experimental_system ON experimental_systems (experiment_id, system_id)",
-    "CREATE INDEX IF NOT EXISTS idx_outcome_rewards_session ON outcome_rewards (session_id)",
-    "CREATE INDEX IF NOT EXISTS idx_outcome_rewards_total ON outcome_rewards (total_reward)",
-    "CREATE INDEX IF NOT EXISTS idx_event_rewards_session ON event_rewards (session_id)",
-    "CREATE INDEX IF NOT EXISTS idx_event_rewards_event ON event_rewards (event_id)",
-    "CREATE INDEX IF NOT EXISTS idx_event_rewards_type ON event_rewards (reward_type)",
-    'CREATE INDEX IF NOT EXISTS idx_event_rewards_key ON event_rewards ("key")',
-)
-
-
-class NativeLibsqlTraceManager(TraceStorage):
-    """Libsql-backed trace manager."""
-
-    def __init__(
-        self,
-        db_url: str | None = None,
-        *,
-        auth_token: str | None = None,
-    ):
-        self._config_auth_token = auth_token
-        self._target = _resolve_connection_target(db_url, auth_token)
-        self._conn: libsql.Connection | None = None
-        self._conn_lock = asyncio.Lock()
-        self._op_lock = asyncio.Lock()
-        self._initialized = False
-
-    def _open_connection(self) -> libsql.Connection:
-        """Open a libsql connection for the resolved target."""
-        kwargs: dict[str, Any] = {}
-        if self._target.sync_url and self._target.sync_url.startswith("libsql://"):
-            kwargs["sync_url"] = self._target.sync_url
-        if self._target.auth_token:
-            kwargs["auth_token"] = self._target.auth_token
-        # Disable automatic background sync; ReplicaSync drives this explicitly.
-        kwargs.setdefault("sync_interval", 0)
-        logger.debug("Opening libsql connection to %s", self._target.database)
-        return libsql.connect(self._target.database, **kwargs)
-
-    async def initialize(self):
-        """Initialise the backend."""
-        async with self._conn_lock:
-            if self._initialized:
-                return
-
-            # Establish a libsql connection for future native operations.
-            self._conn = self._open_connection()
-            self._ensure_schema()
-            self._initialized = True
-
-    async def close(self):
-        """Close the libsql connection."""
-        async with self._conn_lock:
-            if self._conn:
-                logger.debug("Closing libsql connection to %s", self._target.database)
-                self._conn.close()
-                self._conn = None
-            self._initialized = False
-
-    # ------------------------------------------------------------------
-    # Delegated operations (to be swapped with native libsql versions).
-    # ------------------------------------------------------------------
-
-    async def insert_session_trace(self, trace: SessionTrace) -> str:
-        await self.initialize()
-
-        if await self._session_exists(trace.session_id):
-            async with self._op_lock:
-                conn = self._conn
-                assert conn is not None
-                conn.execute(
-                    "UPDATE session_traces SET metadata = ? WHERE session_id = ?",
-                    (_json_dumps(trace.metadata or {}), trace.session_id),
-                )
-                conn.commit()
-            return trace.session_id
-
-        created_at = trace.created_at or datetime.now(UTC)
-
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-            conn.execute(
-                """
-                INSERT INTO session_traces (
-                    session_id,
-                    created_at,
-                    num_timesteps,
-                    num_events,
-                    num_messages,
-                    metadata
-                )
-                VALUES (?, ?, 0, 0, 0, ?)
-                """,
-                (
-                    trace.session_id,
-                    created_at.isoformat(),
-                    _json_dumps(trace.metadata or {}),
-                ),
-            )
-            conn.commit()
-
-        step_id_map: dict[str, int] = {}
-
-        for step in trace.session_time_steps:
-            step_db_id = await self.ensure_timestep(
-                trace.session_id,
-                step_id=step.step_id,
-                step_index=step.step_index,
-                turn_number=step.turn_number,
-                started_at=step.timestamp,
-                completed_at=step.completed_at,
-                metadata=step.step_metadata or {},
-            )
-            step_id_map[step.step_id] = step_db_id
-
-        for event in trace.event_history:
-            step_ref = None
-            metadata = event.metadata or {}
-            if isinstance(metadata, dict):
-                step_ref = metadata.get("step_id")
-            timestep_db_id = step_id_map.get(step_ref) if step_ref else None
-            await self.insert_event_row(
-                trace.session_id,
-                timestep_db_id=timestep_db_id,
-                event=event,
-                metadata_override=event.metadata or {},
-            )
-
-        for msg in trace.markov_blanket_message_history:
-            metadata = dict(getattr(msg, "metadata", {}) or {})
-            step_ref = metadata.get("step_id")
-            content_value = msg.content
-            if isinstance(msg.content, SessionMessageContent):
-                if msg.content.json_payload:
-                    metadata.setdefault("json_payload", msg.content.json_payload)
-                    content_value = msg.content.json_payload
-                else:
-                    content_value = msg.content.as_text()
-                    if msg.content.text:
-                        metadata.setdefault("text", msg.content.text)
-            elif not isinstance(content_value, str):
-                try:
-                    content_value = json.dumps(content_value, ensure_ascii=False)
-                except (TypeError, ValueError):
-                    content_value = str(content_value)
-
-            await self.insert_message_row(
-                trace.session_id,
-                timestep_db_id=step_id_map.get(step_ref) if step_ref else None,
-                message_type=msg.message_type,
-                content=content_value,
-                event_time=msg.time_record.event_time,
-                message_time=msg.time_record.message_time,
-                metadata=metadata,
-            )
-
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-            conn.execute(
-                "UPDATE session_traces SET num_timesteps = ?, num_events = ?, num_messages = ?, metadata = ? WHERE session_id = ?",
-                (
-                    len(trace.session_time_steps),
-                    len(trace.event_history),
-                    len(trace.markov_blanket_message_history),
-                    _json_dumps(trace.metadata or {}),
-                    trace.session_id,
-                ),
-            )
-            conn.commit()
-
-        return trace.session_id
-
-    async def get_session_trace(self, session_id: str) -> dict[str, Any] | None:
-        await self.initialize()
-
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-
-            session_cursor = conn.execute(
-                """
-                SELECT session_id,
-                       created_at,
-                       num_timesteps,
-                       num_events,
-                       num_messages,
-                       metadata
-                FROM session_traces
-                WHERE session_id = ?
-                """,
-                (session_id,),
-            )
-            session_row = session_cursor.fetchone()
-            session_cursor.close()
-
-            if not session_row:
-                return None
-
-            session_columns = ["session_id", "created_at", "num_timesteps", "num_events", "num_messages", "metadata"]
-            session_data = dict(zip(session_columns, session_row))
-
-            timestep_cursor = conn.execute(
-                """
-                SELECT step_id,
-                       step_index,
-                       turn_number,
-                       started_at,
-                       completed_at,
-                       step_metadata
-                FROM session_timesteps
-                WHERE session_id = ?
-                ORDER BY step_index ASC
-                """,
-                (session_id,),
-            )
-            timestep_rows = timestep_cursor.fetchall()
-            timestep_cursor.close()
-
-        return {
-            "session_id": session_data["session_id"],
-            "created_at": _maybe_datetime(session_data["created_at"]),
-            "num_timesteps": session_data["num_timesteps"],
-            "num_events": session_data["num_events"],
-            "num_messages": session_data["num_messages"],
-            "metadata": _load_json(session_data["metadata"]),
-            "timesteps": [
-                {
-                    "step_id": row[0],
-                    "step_index": row[1],
-                    "turn_number": row[2],
-                    "started_at": _maybe_datetime(row[3]),
-                    "completed_at": _maybe_datetime(row[4]),
-                    "metadata": _load_json(row[5]),
-                }
-                for row in timestep_rows
-            ],
-        }
-
-    async def _session_exists(self, session_id: str) -> bool:
-        await self.initialize()
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-            cursor = conn.execute(
-                "SELECT 1 FROM session_traces WHERE session_id = ?", (session_id,)
-            )
-            row = cursor.fetchone()
-            cursor.close()
-            return row is not None
-
-    @staticmethod
-    def _normalise_params(params: dict[str, Any] | None) -> dict[str, Any]:
-        if not params:
-            return {}
-        normalised: dict[str, Any] = {}
-        for key, value in params.items():
-            if isinstance(value, datetime):
-                normalised[key] = value.isoformat()
-            else:
-                normalised[key] = value
-        return normalised
-
-    @staticmethod
-    def _prepare_query_params(query: str, params: dict[str, Any] | list[Any] | tuple[Any, ...]) -> tuple[str, tuple[Any, ...]]:
-        if isinstance(params, dict):
-            keys: list[str] = []
-
-            def _replace(match: re.Match[str]) -> str:
-                key = match.group(1)
-                keys.append(key)
-                return "?"
-
-            new_query = re.sub(r":([a-zA-Z_][a-zA-Z0-9_]*)", _replace, query)
-            if not keys:
-                raise ValueError("No named parameters found in query for provided mapping")
-            values = tuple(params[key] for key in keys)
-            return new_query, values
-        if isinstance(params, (list, tuple)):
-            return query, tuple(params)
-        raise TypeError("Unsupported parameter type for query execution")
-
-    def _ensure_schema(self) -> None:
-        if not self._conn:
-            raise RuntimeError("Connection not initialised")
-
-        for ddl in _TABLE_DEFINITIONS:
-            self._conn.execute(ddl)
-        for ddl in _INDEX_DEFINITIONS:
-            self._conn.execute(ddl)
-        for view_sql in analytics_views.values():
-            self._conn.execute(view_sql)
-        self._conn.commit()
-
-    async def query_traces(self, query: str, params: dict[str, Any] | None = None) -> Any:
-        await self.initialize()
-
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-            normalised = self._normalise_params(params)
-            if normalised:
-                prepared_query, prepared_params = self._prepare_query_params(query, normalised)
-                cursor = conn.execute(prepared_query, prepared_params)
-            else:
-                cursor = conn.execute(query)
-            try:
-                description = cursor.description or []
-                columns = [col[0] for col in description]
-                rows = cursor.fetchall()
-            finally:
-                cursor.close()
-
-        if not rows:
-            if pd is not None:
-                return pd.DataFrame(columns=[col for col in columns])
-            return []
-
-        records = [dict(zip(columns, row)) for row in rows]
-        if pd is not None:
-            return pd.DataFrame(records)
-        return records
-
-    async def get_model_usage(
-        self,
-        start_date=None,
-        end_date=None,
-        model_name=None,
-    ) -> Any:
-        query = """
-            SELECT * FROM model_usage_stats
-            WHERE 1=1
-        """
-        params: dict[str, Any] = {}
-        if start_date:
-            params["start_date"] = start_date
-            query += " AND last_used >= :start_date"
-        if end_date:
-            params["end_date"] = end_date
-            query += " AND first_used <= :end_date"
-        if model_name:
-            params["model_name"] = model_name
-            query += " AND model_name = :model_name"
-        query += " ORDER BY usage_count DESC"
-        return await self.query_traces(query, params)
-
-    async def delete_session(self, session_id: str) -> bool:
-        await self.initialize()
-
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-
-            cursor = conn.execute(
-                "SELECT 1 FROM session_traces WHERE session_id = ?", (session_id,)
-            )
-            exists = cursor.fetchone() is not None
-            cursor.close()
-            if not exists:
-                return False
-
-            conn.execute("DELETE FROM event_rewards WHERE session_id = ?", (session_id,))
-            conn.execute("DELETE FROM outcome_rewards WHERE session_id = ?", (session_id,))
-            conn.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
-            conn.execute("DELETE FROM events WHERE session_id = ?", (session_id,))
-            conn.execute("DELETE FROM session_timesteps WHERE session_id = ?", (session_id,))
-            conn.execute("DELETE FROM session_traces WHERE session_id = ?", (session_id,))
-            conn.commit()
-            return True
-
-    # Experiment helpers -------------------------------------------------
-    async def create_experiment(
-        self,
-        experiment_id: str,
-        name: str,
-        description: str | None = None,
-        configuration: dict[str, Any] | None = None,
-    ) -> str:
-        await self.initialize()
-
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-            conn.execute(
-                """
-                INSERT INTO experiments (experiment_id, name, description, configuration)
-                VALUES (?, ?, ?, ?)
-                ON CONFLICT(experiment_id) DO UPDATE SET
-                    name = excluded.name,
-                    description = excluded.description,
-                    configuration = excluded.configuration
-                """,
-                (
-                    experiment_id,
-                    name,
-                    description,
-                    _json_dumps(configuration or {}),
-                ),
-            )
-            conn.commit()
-        return experiment_id
-
-    async def link_session_to_experiment(self, session_id: str, experiment_id: str):
-        await self.initialize()
-
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-            conn.execute(
-                "UPDATE session_traces SET experiment_id = ? WHERE session_id = ?",
-                (experiment_id, session_id),
-            )
-            conn.commit()
-
-    async def get_sessions_by_experiment(
-        self, experiment_id: str, limit: int | None = None
-    ) -> list[dict[str, Any]]:
-        await self.initialize()
-
-        sql = """
-            SELECT session_id,
-                   created_at,
-                   num_timesteps,
-                   num_events,
-                   num_messages,
-                   metadata
-            FROM session_traces
-            WHERE experiment_id = ?
-            ORDER BY created_at DESC
-        """
-        params: list[Any] = [experiment_id]
-        if limit is not None:
-            sql += " LIMIT ?"
-            params.append(limit)
-
-        async with self._op_lock:
-            conn = self._conn
-            assert conn is not None
-            cursor = conn.execute(sql, params)
-            rows = cursor.fetchall()
-            cursor.close()
-
-        return [
-            {
-                "session_id": row[0],
-                "created_at": _maybe_datetime(row[1]),
-                "num_timesteps": row[2],
-                "num_events": row[3],
-                "num_messages": row[4],
-                "metadata": _load_json(row[5]),
-            }
-            for row in rows
-        ]
-
-    async def batch_insert_sessions(
-        self, traces: list[SessionTrace], batch_size: int | None = None
-    ) -> list[str]:
-        batch_size = batch_size or CONFIG.batch_size
-        inserted: list[str] = []
-
-        for i in range(0, len(traces), batch_size):
-            chunk = traces[i : i + batch_size]
-            for trace in chunk:
-                session_id = await self.insert_session_trace(trace)
-                inserted.append(session_id)
-        return inserted
-
-    # Incremental helpers -----------------------------------------------
-    async def ensure_session(
-        self,
-        session_id: str,
-        *,
-        created_at=None,
-        metadata=None,
-    ) -> None:
-        await self.initialize()
-
-        created_at_val = (created_at or datetime.now(UTC)).isoformat()
-        metadata_json = _json_dumps(metadata or {})
-
-        async with self._op_lock:
-            conn = self._conn
-
-            assert conn is not None
-            conn.execute(
-                """
-                INSERT INTO session_traces (
-                    session_id, created_at, num_timesteps, num_events, num_messages, metadata
-                )
-                VALUES (?, ?, 0, 0, 0, ?)
-                ON CONFLICT(session_id) DO NOTHING
-                """,
-                (session_id, created_at_val, metadata_json),
-            )
-            conn.commit()
-
-    async def ensure_timestep(
-        self,
-        session_id: str,
-        *,
-        step_id: str,
-        step_index: int,
-        turn_number: int | None = None,
-        started_at=None,
-        completed_at=None,
-        metadata=None,
-    ) -> int:
-        await self.initialize()
-
-        started_at_val = (started_at or datetime.now(UTC)).isoformat()
-        completed_at_val = completed_at.isoformat() if completed_at else None
-        metadata_json = _json_dumps(metadata or {})
-
-        async with self._op_lock:
-            conn = self._conn
-
-            assert conn is not None
-            cur = conn.execute(
-                """
-                SELECT id FROM session_timesteps
-                WHERE session_id = ? AND step_id = ?
-                """,
-                (session_id, step_id),
-            )
-            row = cur.fetchone()
-            if row:
-                return int(row[0])
-
-            cur = conn.execute(
-                """
-                INSERT INTO session_timesteps (
-                    session_id,
-                    step_id,
-                    step_index,
-                    turn_number,
-                    started_at,
-                    completed_at,
-                    num_events,
-                    num_messages,
-                    step_metadata
-                )
-                VALUES (?, ?, ?, ?, ?, ?, 0, 0, ?)
-                """,
-                (
-                    session_id,
-                    step_id,
-                    step_index,
-                    turn_number,
-                    started_at_val,
-                    completed_at_val,
-                    metadata_json,
-                ),
-            )
-            timestep_id = int(cur.lastrowid)
-            conn.execute(
-                """
-                UPDATE session_traces
-                SET num_timesteps = num_timesteps + 1
-                WHERE session_id = ?
-                """,
-                (session_id,),
-            )
-            conn.commit()
-            return timestep_id
-
-    async def insert_event_row(
-        self,
-        session_id: str,
-        *,
-        timestep_db_id: int | None,
-        event: Any,
-        metadata_override: dict[str, Any] | None = None,
-    ) -> int:
-        await self.initialize()
-
-        if not isinstance(event, (EnvironmentEvent, LMCAISEvent, RuntimeEvent)):
-            raise TypeError(f"Unsupported event type for native manager: {type(event)!r}")
-
-        metadata_json = metadata_override or event.metadata or {}
-        event_extra_metadata = getattr(event, "event_metadata", None)
-        system_state_before = getattr(event, "system_state_before", None)
-        system_state_after = getattr(event, "system_state_after", None)
-
-        payload: 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,
-            "metadata": metadata_json,
-            "event_metadata": event_extra_metadata,
-            "system_state_before": system_state_before,
-            "system_state_after": system_state_after,
-        }
-
-        if isinstance(event, LMCAISEvent):
-            call_records = None
-            if getattr(event, "call_records", None):
-                call_records = [asdict(record) for record in event.call_records]
-            payload.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": int(event.cost_usd * 100) if event.cost_usd is not None else None,
-                    "latency_ms": event.latency_ms,
-                    "span_id": event.span_id,
-                    "trace_id": event.trace_id,
-                    "call_records": call_records,
-                }
-            )
-        elif isinstance(event, EnvironmentEvent):
-            payload.update(
-                {
-                    "event_type": "environment",
-                    "reward": event.reward,
-                    "terminated": event.terminated,
-                    "truncated": event.truncated,
-                }
-            )
-        elif isinstance(event, RuntimeEvent):
-            payload.update(
-                {
-                    "event_type": "runtime",
-                    "metadata": {**(event.metadata or {}), "actions": event.actions},
-                }
-            )
-
-        async with self._op_lock:
-            conn = self._conn
-
-            assert conn is not None
-            cur = conn.execute(
-                """
-                INSERT INTO events (
-                    session_id,
-                    timestep_id,
-                    event_type,
-                    system_instance_id,
-                    event_time,
-                    message_time,
-                    model_name,
-                    provider,
-                    input_tokens,
-                    output_tokens,
-                    total_tokens,
-                    cost_usd,
-                    latency_ms,
-                    span_id,
-                    trace_id,
-                    call_records,
-                    reward,
-                    terminated,
-                    truncated,
-                    system_state_before,
-                    system_state_after,
-                    metadata,
-                    event_metadata
-                )
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-                """,
-                (
-                    payload["session_id"],
-                    payload["timestep_id"],
-                    payload.get("event_type"),
-                    payload["system_instance_id"],
-                    payload["event_time"],
-                    payload["message_time"],
-                    payload.get("model_name"),
-                    payload.get("provider"),
-                    payload.get("input_tokens"),
-                    payload.get("output_tokens"),
-                    payload.get("total_tokens"),
-                    payload.get("cost_usd"),
-                    payload.get("latency_ms"),
-                    payload.get("span_id"),
-                    payload.get("trace_id"),
-                    _json_dumps(payload.get("call_records")),
-                    payload.get("reward"),
-                    payload.get("terminated"),
-                    payload.get("truncated"),
-                    _json_dumps(payload.get("system_state_before")),
-                    _json_dumps(payload.get("system_state_after")),
-                    _json_dumps(payload.get("metadata")),
-                    _json_dumps(payload.get("event_metadata")),
-                ),
-            )
-            event_id = int(cur.lastrowid)
-            conn.execute(
-                """
-                UPDATE session_traces
-                SET num_events = num_events + 1
-                WHERE session_id = ?
-                """,
-                (session_id,),
-            )
-            if timestep_db_id is not None:
-                conn.execute(
-                    """
-                    UPDATE session_timesteps
-                    SET num_events = num_events + 1
-                    WHERE id = ?
-                    """,
-                    (timestep_db_id,),
-                )
-            conn.commit()
-            return event_id
-
-    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:
-        await self.initialize()
-
-        metadata_payload = dict(metadata or {})
-        if isinstance(content, SessionMessageContent):
-            if content.json_payload:
-                metadata_payload.setdefault("json_payload", content.json_payload)
-                content_value = content.json_payload
-            else:
-                content_value = content.as_text()
-                if content.text:
-                    metadata_payload.setdefault("text", content.text)
-        else:
-            content_value = content
-            if not isinstance(content_value, str):
-                try:
-                    content_value = json.dumps(content_value, ensure_ascii=False)
-                except (TypeError, ValueError):
-                    content_value = str(content_value)
-
-        async with self._op_lock:
-            conn = self._conn
-
-            assert conn is not None
-            cur = conn.execute(
-                """
-                INSERT INTO messages (
-                    session_id,
-                    timestep_id,
-                    message_type,
-                    content,
-                    event_time,
-                    message_time,
-                    metadata
-                )
-                VALUES (?, ?, ?, ?, ?, ?, ?)
-                """,
-                (
-                    session_id,
-                    timestep_db_id,
-                    message_type,
-                    content_value,
-                    event_time,
-                    message_time,
-                    _json_dumps(metadata_payload),
-                ),
-            )
-            message_id = int(cur.lastrowid)
-            conn.execute(
-                """
-                UPDATE session_traces
-                SET num_messages = num_messages + 1
-                WHERE session_id = ?
-                """,
-                (session_id,),
-            )
-            if timestep_db_id is not None:
-                conn.execute(
-                    """
-                    UPDATE session_timesteps
-                    SET num_messages = num_messages + 1
-                    WHERE id = ?
-                    """,
-                    (timestep_db_id,),
-                )
-            conn.commit()
-            return message_id
-
-    async def insert_outcome_reward(
-        self,
-        session_id: str,
-        *,
-        total_reward: int,
-        achievements_count: int,
-        total_steps: int,
-        reward_metadata: dict | None = None,
-    ) -> int:
-        await self.initialize()
-
-        async with self._op_lock:
-            conn = self._conn
-
-            assert conn is not None
-            cur = conn.execute(
-                """
-                INSERT INTO outcome_rewards (
-                    session_id,
-                    total_reward,
-                    achievements_count,
-                    total_steps,
-                    created_at,
-                    reward_metadata
-                )
-                VALUES (?, ?, ?, ?, ?, ?)
-                """,
-                (
-                    session_id,
-                    total_reward,
-                    achievements_count,
-                    total_steps,
-                    datetime.now(UTC).isoformat(),
-                    _json_dumps(reward_metadata),
-                ),
-            )
-            conn.commit()
-            return int(cur.lastrowid)
-
-    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:
-        await self.initialize()
-
-        async with self._op_lock:
-            conn = self._conn
-
-            assert conn is not None
-            cur = conn.execute(
-                """
-                INSERT INTO event_rewards (
-                    event_id,
-                    session_id,
-                    message_id,
-                    turn_number,
-                    reward_value,
-                    reward_type,
-                    key,
-                    annotation,
-                    source,
-                    created_at
-                )
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-                """,
-                (
-                    event_id,
-                    session_id,
-                    message_id,
-                    turn_number,
-                    reward_value,
-                    reward_type,
-                    key,
-                    _json_dumps(annotation),
-                    source,
-                    datetime.now(UTC).isoformat(),
-                ),
-            )
-            conn.commit()
-            return int(cur.lastrowid)