synth-ai 0.2.13.dev1__py3-none-any.whl → 0.2.14__py3-none-any.whl

synth_ai/tracing_v3/storage/base.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any
+from typing import Any, Optional
 
 from ..abstractions import SessionTrace
 
@@ -28,7 +28,7 @@ class TraceStorage(ABC):
         pass
 
     @abstractmethod
-    async def get_session_trace(self, session_id: str) -> dict[str, Any] | None:
+    async def get_session_trace(self, session_id: str) -> Optional[dict[str, Any]]:
         """Retrieve a session trace by ID.
 
         Args:
@@ -40,7 +40,7 @@ class TraceStorage(ABC):
         pass
 
     @abstractmethod
-    async def query_traces(self, query: str, params: dict[str, Any] | None = None) -> Any:
+    async def query_traces(self, query: str, params: Optional[dict[str, Any]] = None) -> Any:
         """Execute a query and return results.
 
         Args:
@@ -55,9 +55,9 @@ class TraceStorage(ABC):
     @abstractmethod
     async def get_model_usage(
         self,
-        start_date: datetime | None = None,
-        end_date: datetime | None = None,
-        model_name: str | None = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        model_name: Optional[str] = None,
     ) -> Any:
         """Get model usage statistics.
 
@@ -95,8 +95,8 @@ class TraceStorage(ABC):
         self,
         session_id: str,
         *,
-        created_at: datetime | None = None,
-        metadata: dict[str, Any] | None = None,
+        created_at: Optional[datetime] = None,
+        metadata: Optional[dict[str, Any]] = None,
     ) -> None:
         """Ensure a session row exists for the given session id."""
         pass
@@ -108,10 +108,10 @@ class TraceStorage(ABC):
         *,
         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,
+        turn_number: Optional[int] = None,
+        started_at: Optional[datetime] = None,
+        completed_at: Optional[datetime] = None,
+        metadata: Optional[dict[str, Any]] = None,
     ) -> int:
         """Ensure a timestep row exists and return its database id."""
         pass
@@ -121,9 +121,9 @@ class TraceStorage(ABC):
         self,
         session_id: str,
         *,
-        timestep_db_id: int | None,
+        timestep_db_id: Optional[int],
         event: Any,
-        metadata_override: dict[str, Any] | None = None,
+        metadata_override: Optional[dict[str, Any]] = None,
     ) -> int:
         """Insert an event and return its database id."""
         pass
@@ -133,12 +133,12 @@ class TraceStorage(ABC):
         self,
         session_id: str,
         *,
-        timestep_db_id: int | None,
+        timestep_db_id: Optional[int],
         message_type: str,
         content: Any,
-        event_time: float | None = None,
-        message_time: int | None = None,
-        metadata: dict[str, Any] | None = None,
+        event_time: Optional[float] = None,
+        message_time: Optional[int] = None,
+        metadata: Optional[dict[str, Any]] = None,
     ) -> int:
         """Insert a message row linked to a session/timestep."""
         pass
@@ -151,7 +151,7 @@ class TraceStorage(ABC):
         total_reward: int,
         achievements_count: int,
         total_steps: int,
-        reward_metadata: dict | None = None,
+        reward_metadata: Optional[dict] = None,
     ) -> int:
         """Record an outcome reward for a session."""
         pass
@@ -162,13 +162,13 @@ class TraceStorage(ABC):
         session_id: str,
         *,
         event_id: int,
-        message_id: int | None = None,
-        turn_number: int | None = None,
+        message_id: Optional[int] = None,
+        turn_number: Optional[int] = 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,
+        reward_type: Optional[str] = None,
+        key: Optional[str] = None,
+        annotation: Optional[dict[str, Any]] = None,
+        source: Optional[str] = None,
     ) -> int:
         """Record a reward tied to a specific event."""
         pass
@@ -178,8 +178,8 @@ class TraceStorage(ABC):
         self,
         experiment_id: str,
         name: str,
-        description: str | None = None,
-        configuration: dict[str, Any] | None = None,
+        description: Optional[str] = None,
+        configuration: Optional[dict[str, Any]] = None,
     ) -> str:
         """Create a new experiment."""
         raise NotImplementedError("Experiment management not supported by this backend")
@@ -189,14 +189,14 @@ class TraceStorage(ABC):
         raise NotImplementedError("Experiment management not supported by this backend")
 
     async def get_sessions_by_experiment(
-        self, experiment_id: str, limit: int | None = None
+        self, experiment_id: str, limit: Optional[int] = 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
+        self, traces: list[SessionTrace], batch_size: Optional[int] = 1000
     ) -> list[str]:
         """Batch insert multiple session traces.
 

synth_ai/tracing_v3/storage/config.py

@@ -3,7 +3,7 @@
 import os
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any
+from typing import Any, Optional
 
 
 class StorageBackend(str, Enum):
@@ -14,7 +14,7 @@ class StorageBackend(str, Enum):
     POSTGRES = "postgres"  # Future support
 
 
-def _is_enabled(value: str | None) -> bool:
+def _is_enabled(value: Optional[str]) -> bool:
     if value is None:
         return False
     return value.lower() in {"1", "true", "yes", "on"}
@@ -25,7 +25,7 @@ class StorageConfig:
     """Configuration for storage backend."""
 
     backend: StorageBackend = StorageBackend.TURSO_NATIVE
-    connection_string: str | None = None
+    connection_string: Optional[str] = None
 
     # Turso-specific settings
     turso_url: str = os.getenv("TURSO_DATABASE_URL", "sqlite+libsql://http://127.0.0.1:8080")

synth_ai/tracing_v3/trace_utils.py

@@ -0,0 +1,317 @@
+from __future__ import annotations
+
+import json
+import sqlite3
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+Row = sqlite3.Row
+
+
+def connect(db_path: str | bytes | int) -> 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 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"],
+                "reward_metadata": _json_load(row["reward_metadata"]) 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",
+]
+
+

synth_ai/tracing_v3/turso/daemon.py

@@ -7,6 +7,7 @@ import time
 
 import requests
 from requests import RequestException
+from typing import Any, Optional
 
 from ..config import CONFIG
 
@@ -16,9 +17,9 @@ class SqldDaemon:
 
     def __init__(
         self,
-        db_path: str | None = None,
-        http_port: int | None = None,
-        binary_path: str | None = None,
+        db_path: Optional[str] = None,
+        http_port: Optional[int] = None,
+        binary_path: Optional[str] = None,
     ):
         """Initialize sqld daemon manager.
 
@@ -30,7 +31,7 @@ class SqldDaemon:
         self.db_path = db_path or CONFIG.sqld_db_path
         self.http_port = http_port or CONFIG.sqld_http_port
         self.binary_path = binary_path or self._find_binary()
-        self.process: subprocess.Popen | None = None
+        self.process: Optional[Any] = None
 
     def _find_binary(self) -> str:
         """Find sqld binary in PATH."""
@@ -123,10 +124,10 @@ class SqldDaemon:
 
 
 # Convenience functions
-_daemon: SqldDaemon | None = None
+_daemon: Optional[SqldDaemon] = None
 
 
-def start_sqld(db_path: str | None = None, port: int | None = None) -> SqldDaemon:
+def start_sqld(db_path: Optional[str] = None, port: Optional[int] = None) -> SqldDaemon:
     """Start a global sqld daemon instance."""
     global _daemon
     if _daemon and _daemon.is_running():
@@ -145,6 +146,6 @@ def stop_sqld():
         _daemon = None
 
 
-def get_daemon() -> SqldDaemon | None:
+def get_daemon() -> Optional[SqldDaemon]:
     """Get the global daemon instance."""
     return _daemon