synth-ai 0.2.14__py3-none-any.whl → 0.4.1__py3-none-any.whl

synth_ai/core/tracing_v3/storage/base.py

@@ -0,0 +1,211 @@
+"""Abstract base class for trace storage."""
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Any
+
+from ..abstractions import SessionTrace
+
+
+class TraceStorage(ABC):
+    """Abstract base class for trace storage implementations."""
+
+    @abstractmethod
+    async def initialize(self):
+        """Initialize the storage backend."""
+        pass
+
+    @abstractmethod
+    async def insert_session_trace(self, trace: SessionTrace) -> str:
+        """Insert a complete session trace.
+
+        Args:
+            trace: The session trace to insert
+
+        Returns:
+            The session ID
+        """
+        pass
+
+    @abstractmethod
+    async def get_session_trace(self, session_id: str) -> dict[str, Any] | None:
+        """Retrieve a session trace by ID.
+
+        Args:
+            session_id: The session ID to retrieve
+
+        Returns:
+            The session trace data or None if not found
+        """
+        pass
+
+    @abstractmethod
+    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-like object or list of dict records
+        """
+        pass
+
+    @abstractmethod
+    async def get_model_usage(
+        self,
+        start_date: datetime | None = None,
+        end_date: datetime | None = None,
+        model_name: str | None = None,
+    ) -> Any:
+        """Get model usage statistics.
+
+        Args:
+            start_date: Optional start date filter
+            end_date: Optional end date filter
+            model_name: Optional model name filter
+
+        Returns:
+            Model usage statistics as a DataFrame-like object or list of dict records
+        """
+        pass
+
+    @abstractmethod
+    async def delete_session(self, session_id: str) -> bool:
+        """Delete a session and all related data.
+
+        Args:
+            session_id: The session ID to delete
+
+        Returns:
+            True if deleted, False if not found
+        """
+        pass
+
+    @abstractmethod
+    async def close(self):
+        """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: float,
+        achievements_count: int,
+        total_steps: int,
+        reward_metadata: dict | None = None,
+        annotation: dict[str, Any] | 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 = None,
+        configuration: dict[str, Any] | None = None,
+    ) -> str:
+        """Create a new experiment."""
+        raise NotImplementedError("Experiment management not supported by this backend")
+
+    async def link_session_to_experiment(self, session_id: str, experiment_id: str):
+        """Link a session to an experiment."""
+        raise NotImplementedError("Experiment management not supported by this backend")
+
+    async def get_sessions_by_experiment(
+        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 | None = 1000
+    ) -> list[str]:
+        """Batch insert multiple session traces.
+
+        Default implementation calls insert_session_trace for each trace.
+        Subclasses can override for more efficient batch operations.
+        """
+        inserted_ids = []
+        for trace in traces:
+            session_id = await self.insert_session_trace(trace)
+            inserted_ids.append(session_id)
+        return inserted_ids

synth_ai/core/tracing_v3/storage/config.py

@@ -0,0 +1,109 @@
+"""Storage configuration for tracing v3."""
+
+import os
+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_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."""
+
+    connection_string: str | None = None
+    backend: StorageBackend | None = None
+    turso_auth_token: str | None = field(default=None)
+
+    # Common settings
+    pool_size: int = int(os.getenv("STORAGE_POOL_SIZE", "8"))
+    echo_sql: bool = os.getenv("STORAGE_ECHO_SQL", "false").lower() == "true"
+    batch_size: int = int(os.getenv("STORAGE_BATCH_SIZE", "1000"))
+
+    # Performance settings
+    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_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_NATIVE:
+            config = {}
+            if self.turso_auth_token:
+                config["auth_token"] = self.turso_auth_token
+            return config
+        return {}
+
+
+# Global storage configuration
+STORAGE_CONFIG = StorageConfig()

synth_ai/core/tracing_v3/storage/factory.py

@@ -0,0 +1,39 @@
+"""Factory for creating storage instances."""
+
+from ..turso.native_manager import NativeLibsqlTraceManager
+from .base import TraceStorage
+from .config import StorageBackend, StorageConfig
+
+
+def create_storage(config: StorageConfig | None = None) -> TraceStorage:
+    """Create a storage instance based on configuration.
+
+    Args:
+        config: Storage configuration (uses default if not provided)
+
+    Returns:
+        A TraceStorage implementation
+
+    Raises:
+        ValueError: If the backend is not supported
+    """
+    if config is None:
+        from .config import STORAGE_CONFIG
+
+        config = STORAGE_CONFIG
+
+    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")
+    else:
+        raise ValueError(f"Unknown storage backend: {config.backend}")

synth_ai/core/tracing_v3/trace_utils.py

@@ -0,0 +1,326 @@
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+Row = sqlite3.Row
+
+
+def connect(db_path: str | bytes | os.PathLike[str] | os.PathLike[bytes]) -> sqlite3.Connection:
+    conn = sqlite3.connect(db_path)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def _json_load(value: Any) -> Any:
+    if value is None:
+        return None
+    if isinstance(value, dict | list):
+        return value
+    if isinstance(value, bytes | bytearray):
+        value = value.decode("utf-8", errors="ignore")
+    try:
+        return json.loads(value)
+    except Exception:
+        return value
+
+
+def _row_get(row: Row, key: str) -> Any:
+    try:
+        return row[key]
+    except Exception:
+        return None
+
+
+def fetch_crafter_sessions(
+    conn: sqlite3.Connection,
+    *,
+    limit: int,
+    metadata_filter: str | None = None,
+    session_ids: Sequence[str] | None = None,
+    min_event_count: int = 0,
+) -> list[str]:
+    if session_ids:
+        placeholders = ",".join("?" for _ in session_ids)
+        rows = conn.execute(
+            f"""
+            SELECT session_id
+            FROM session_traces
+            WHERE session_id IN ({placeholders})
+            ORDER BY created_at DESC
+            """,
+            tuple(session_ids),
+        ).fetchall()
+        return [row["session_id"] for row in rows]
+
+    params: list[Any] = []
+    where_clauses: list[str] = []
+    if metadata_filter:
+        where_clauses.append("session_traces.metadata LIKE ?")
+        params.append(f"%{metadata_filter}%")
+    where_sql = ""
+    if where_clauses:
+        where_sql = "WHERE " + " AND ".join(where_clauses)
+
+    having_sql = ""
+    if min_event_count > 0:
+        having_sql = "HAVING COUNT(events.id) >= ?"
+        params.append(min_event_count)
+
+    query = f"""
+        SELECT session_traces.session_id
+        FROM session_traces
+        LEFT JOIN events ON session_traces.session_id = events.session_id
+        {where_sql}
+        GROUP BY session_traces.session_id
+        {having_sql}
+        ORDER BY session_traces.created_at DESC
+        LIMIT ?
+    """
+    rows = conn.execute(query, (*params, limit)).fetchall()
+    return [row["session_id"] for row in rows]
+
+
+def load_session_trace(conn: sqlite3.Connection, session_id: str) -> dict[str, Any]:
+    session_row = conn.execute(
+        """
+        SELECT session_id, created_at, metadata
+        FROM session_traces
+        WHERE session_id = ?
+        """,
+        (session_id,),
+    ).fetchone()
+    if not session_row:
+        raise ValueError(f"Session {session_id} not found")
+
+    timesteps = 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,),
+    ).fetchall()
+
+    event_rows = conn.execute(
+        """
+        SELECT *
+        FROM events
+        WHERE session_id = ?
+        ORDER BY event_time ASC, id ASC
+        """,
+        (session_id,),
+    ).fetchall()
+
+    message_rows = conn.execute(
+        """
+        SELECT *
+        FROM messages
+        WHERE session_id = ?
+        ORDER BY event_time ASC, id ASC
+        """,
+        (session_id,),
+    ).fetchall()
+
+    event_rewards = conn.execute(
+        """
+        SELECT *
+        FROM event_rewards
+        WHERE session_id = ?
+        ORDER BY turn_number ASC, id ASC
+        """,
+        (session_id,),
+    ).fetchall()
+
+    outcome_rewards = conn.execute(
+        """
+        SELECT *
+        FROM outcome_rewards
+        WHERE session_id = ?
+        ORDER BY created_at ASC
+        """,
+        (session_id,),
+    ).fetchall()
+
+    metadata = _json_load(session_row["metadata"]) or {}
+    if isinstance(metadata, dict):
+        episode_id = metadata.get("episode_id")
+        if episode_id is not None and not isinstance(episode_id, str):
+            metadata["episode_id"] = str(episode_id)
+
+    events_payload = [
+        {
+            "id": row["id"],
+            "event_type": row["event_type"],
+            "system_instance_id": row["system_instance_id"],
+            "time_record": {
+                "event_time": row["event_time"],
+                "message_time": row["message_time"],
+                "created_at": row["created_at"],
+            },
+            "model_name": row["model_name"],
+            "provider": row["provider"],
+            "input_tokens": row["input_tokens"],
+            "output_tokens": row["output_tokens"],
+            "total_tokens": row["total_tokens"],
+            "cost_usd": row["cost_usd"],
+            "latency_ms": row["latency_ms"],
+            "span_id": row["span_id"],
+            "trace_id": row["trace_id"],
+            "call_records": _json_load(row["call_records"]) or [],
+            "reward": row["reward"],
+            "terminated": row["terminated"],
+            "truncated": row["truncated"],
+            "system_state_before": _json_load(row["system_state_before"]),
+            "system_state_after": _json_load(row["system_state_after"]),
+            "metadata": _json_load(row["metadata"]) or {},
+            "event_metadata": _json_load(row["event_metadata"]),
+        }
+        for row in event_rows
+    ]
+
+    messages_payload = [
+        {
+            "id": row["id"],
+            "message_type": row["message_type"],
+            "content": row["content"],
+            "time_record": {
+                "event_time": row["event_time"],
+                "message_time": row["message_time"],
+                "timestamp": row["timestamp"],
+            },
+            "metadata": _json_load(row["metadata"]) or {},
+        }
+        for row in message_rows
+    ]
+
+    trace: dict[str, Any] = {
+        "session_id": session_row["session_id"],
+        "created_at": session_row["created_at"],
+        "metadata": metadata,
+        "session_time_steps": [
+            {
+                "step_id": row["step_id"],
+                "step_index": row["step_index"],
+                "turn_number": row["turn_number"],
+                "started_at": row["started_at"],
+                "completed_at": row["completed_at"],
+                "metadata": _json_load(row["step_metadata"]) or {},
+            }
+            for row in timesteps
+        ],
+        "event_history": events_payload,
+        "events": events_payload,
+        "markov_blanket_message_history": messages_payload,
+        "messages": messages_payload,
+        "event_rewards": [
+            {
+                "id": row["id"],
+                "event_id": row["event_id"],
+                "turn_number": row["turn_number"],
+                "reward_value": row["reward_value"],
+                "reward_type": row["reward_type"],
+                "key": row["key"],
+                "annotation": _json_load(row["annotation"]) or {},
+                "source": row["source"],
+                "created_at": row["created_at"],
+            }
+            for row in event_rewards
+        ],
+        "outcome_rewards": [
+            {
+                "id": row["id"],
+                "total_reward": row["total_reward"],
+                "achievements_count": row["achievements_count"],
+                "total_steps": row["total_steps"],
+                "reward_metadata": _json_load(row["reward_metadata"]) or {},
+                "annotation": _json_load(_row_get(row, "annotation")) or {},
+                "created_at": row["created_at"],
+            }
+            for row in outcome_rewards
+        ],
+    }
+    return trace
+
+
+@dataclass
+class DeterministicMetrics:
+    session_id: str
+    unique_achievement_reward: float
+    achievement_reward: float
+    outcome_total_reward: float
+    unique_achievement_count: int
+    final_achievement_count: int
+
+
+def compute_deterministic_metrics(conn: sqlite3.Connection, session_id: str) -> DeterministicMetrics:
+    event_rows = conn.execute(
+        """
+        SELECT reward_type, reward_value, annotation
+        FROM event_rewards
+        WHERE session_id = ?
+        """,
+        (session_id,),
+    ).fetchall()
+
+    unique_total = 0.0
+    all_total = 0.0
+    unique_achievements: set[str] = set()
+
+    for row in event_rows:
+        reward_type = row["reward_type"]
+        value = float(row["reward_value"] or 0.0)
+        if reward_type == "unique_achievement_delta":
+            unique_total += value
+            annotation = _json_load(row["annotation"]) or {}
+            for name in annotation.get("new_unique") or []:
+                if isinstance(name, str):
+                    unique_achievements.add(name)
+        elif reward_type == "achievement_delta":
+            all_total += value
+
+    outcome_rows = conn.execute(
+        """
+        SELECT total_reward, reward_metadata
+        FROM outcome_rewards
+        WHERE session_id = ?
+        """,
+        (session_id,),
+    ).fetchall()
+
+    outcome_total = 0.0
+    final_achievements: set[str] = set()
+    for row in outcome_rows:
+        outcome_total += float(row["total_reward"] or 0.0)
+        metadata = _json_load(row["reward_metadata"]) or {}
+        for name in metadata.get("achievements") or []:
+            if isinstance(name, str):
+                final_achievements.add(name)
+
+    return DeterministicMetrics(
+        session_id=session_id,
+        unique_achievement_reward=unique_total,
+        achievement_reward=all_total,
+        outcome_total_reward=outcome_total,
+        unique_achievement_count=len(unique_achievements),
+        final_achievement_count=len(final_achievements),
+    )
+
+
+__all__ = [
+    "DeterministicMetrics",
+    "compute_deterministic_metrics",
+    "connect",
+    "fetch_crafter_sessions",
+    "load_session_trace",
+]