synth-ai 0.2.9.dev11__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/storage/utils.py

@@ -0,0 +1,206 @@
+"""Utility functions for storage layer."""
+
+import asyncio
+import functools
+import time
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar, cast
+
+T = TypeVar("T")
+
+
+def retry_async(max_attempts: int = 3, delay: float = 1.0, backoff: float = 2.0):
+    """Decorator for retrying async functions on failure.
+
+    Args:
+        max_attempts: Maximum number of attempts
+        delay: Initial delay between attempts in seconds
+        backoff: Backoff multiplier for each retry
+    """
+
+    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> T:
+            last_exception: Exception | None = None
+            current_delay = delay
+
+            for attempt in range(max_attempts):
+                try:
+                    return await func(*args, **kwargs)
+                except Exception as e:
+                    last_exception = e
+                    if attempt < max_attempts - 1:
+                        await asyncio.sleep(current_delay)
+                        current_delay *= backoff
+                    else:
+                        raise
+
+            if last_exception:
+                raise last_exception
+            raise RuntimeError("Retry logic failed without exception")
+
+        return wrapper
+
+    return decorator
+
+
+async def batch_process(
+    items: list[Any], processor: Callable, batch_size: int = 100, max_concurrent: int = 5
+) -> list[Any]:
+    """Process items in batches with concurrency control.
+
+    Args:
+        items: List of items to process
+        processor: Async function to process a batch
+        batch_size: Size of each batch
+        max_concurrent: Maximum concurrent batches
+
+    Returns:
+        List of results from all batches
+    """
+    semaphore = asyncio.Semaphore(max_concurrent)
+    results = []
+
+    async def process_batch(batch):
+        async with semaphore:
+            return await processor(batch)
+
+    # Create batches
+    batches = [items[i : i + batch_size] for i in range(0, len(items), batch_size)]
+
+    # Process all batches
+    batch_results = await asyncio.gather(*[process_batch(batch) for batch in batches])
+
+    # Flatten results
+    for result in batch_results:
+        if isinstance(result, list):
+            results.extend(result)
+        else:
+            results.append(result)
+
+    return results
+
+
+def sanitize_json(data: Any) -> Any:
+    """Sanitize data for JSON storage.
+
+    Converts non-JSON-serializable types to strings.
+    """
+    if isinstance(data, dict):
+        return {k: sanitize_json(v) for k, v in data.items()}
+    elif isinstance(data, list):
+        return [sanitize_json(item) for item in data]
+    elif isinstance(data, str | int | float | bool | type(None)):
+        return data
+    else:
+        return str(data)
+
+
+def estimate_size(obj: Any) -> int:
+    """Estimate the size of an object in bytes.
+
+    This is a rough estimate for storage planning.
+    """
+    if isinstance(obj, str):
+        return len(obj.encode("utf-8"))
+    elif isinstance(obj, int | float):
+        return 8
+    elif isinstance(obj, bool):
+        return 1
+    elif isinstance(obj, dict):
+        size = 0
+        for k, v in obj.items():
+            size += estimate_size(k) + estimate_size(v)
+        return size
+    elif isinstance(obj, list):
+        return sum(estimate_size(item) for item in obj)
+    else:
+        return len(str(obj).encode("utf-8"))
+
+
+class StorageMetrics:
+    """Track storage operation metrics."""
+
+    def __init__(self):
+        self.operations: dict[str, dict[str, Any]] = {}
+
+    def record_operation(
+        self, operation: str, duration: float, success: bool, size: int | None = None
+    ):
+        """Record a storage operation."""
+        if operation not in self.operations:
+            self.operations[operation] = {
+                "count": 0,
+                "success_count": 0,
+                "total_duration": 0.0,
+                "total_size": 0,
+            }
+
+        stats = self.operations[operation]
+        stats["count"] += 1
+        if success:
+            stats["success_count"] += 1
+        stats["total_duration"] += duration
+        if size:
+            stats["total_size"] += size
+
+    def get_stats(self, operation: str | None = None) -> dict[str, Any]:
+        """Get statistics for operations."""
+        if operation:
+            stats = self.operations.get(operation, {})
+            if stats:
+                return {
+                    "count": stats["count"],
+                    "success_rate": stats["success_count"] / stats["count"]
+                    if stats["count"] > 0
+                    else 0,
+                    "avg_duration": stats["total_duration"] / stats["count"]
+                    if stats["count"] > 0
+                    else 0,
+                    "total_size": stats["total_size"],
+                }
+            return {}
+        else:
+            return {op: self.get_stats(op) for op in self.operations}
+
+
+# Global metrics instance
+STORAGE_METRICS = StorageMetrics()
+
+
+def track_metrics(operation: str):
+    """Decorator to track storage operation metrics."""
+
+    def decorator(func: Callable[..., Awaitable[T]] | Callable[..., T]) -> Callable[..., Awaitable[T]] | Callable[..., T]:
+        @functools.wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> T:
+            start_time = time.time()
+            success = False
+            try:
+                async_func = cast(Callable[..., Awaitable[T]], func)
+                result = await async_func(*args, **kwargs)
+                success = True
+                return result
+            finally:
+                duration = time.time() - start_time
+                STORAGE_METRICS.record_operation(operation, duration, success)
+
+        @functools.wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> T:
+            start_time = time.time()
+            success = False
+            try:
+                sync_func = cast(Callable[..., T], func)
+                result = sync_func(*args, **kwargs)
+                success = True
+                return result
+            finally:
+                duration = time.time() - start_time
+                STORAGE_METRICS.record_operation(operation, duration, success)
+
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    return decorator