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

synth_ai/lm/warmup.py

@@ -3,13 +3,14 @@ Model warmup utilities for Synth backend.
 Handles model preloading and warmup polling.
 """
 
-import httpx
 import asyncio
 import logging
 import sys
 import time
-from typing import Optional, Dict, Any
 from datetime import datetime, timedelta
+
+import httpx
+
 from .config import SynthConfig
 
 logger = logging.getLogger(__name__)
@@ -19,7 +20,7 @@ class WarmupStatus:
     """Track warmup status for models with TTL."""
 
     def __init__(self):
-        self._warmed_models: Dict[str, datetime] = {}
+        self._warmed_models: dict[str, datetime] = {}
         self._ttl = timedelta(minutes=10)  # Consider models warm for 10 minutes
 
     def is_warm(self, model_name: str) -> bool:
@@ -47,11 +48,11 @@ _warmup_status = WarmupStatus()
 
 async def warmup_synth_model(
     model_name: str,
-    config: Optional[SynthConfig] = None,
-    max_attempts: Optional[int] = None,
+    config: SynthConfig | None = None,
+    max_attempts: int | None = None,
     force: bool = False,
     verbose: bool = True,
-    gpu_preference: Optional[str] = None,
+    gpu_preference: str | None = None,
 ) -> bool:
     """
     Warm up a model on the Synth backend using fire-and-forget approach.
@@ -161,7 +162,7 @@ async def warmup_synth_model(
                 )
                 sys.stdout.flush()
                 await asyncio.sleep(1.0)
-            except Exception as e:
+            except Exception:
                 # Continue polling; update spinner line with error label
                 elapsed = int(time.time() - start_time)
                 wheel = spinner[spin_idx % len(spinner)]

synth_ai/rl/__init__.py

@@ -0,0 +1,30 @@
+from .contracts import (
+    RolloutEnvSpec,
+    RolloutPolicySpec,
+    RolloutRecordConfig,
+    RolloutSafetyConfig,
+    RolloutRequest,
+    RolloutStep,
+    RolloutTrajectory,
+    RolloutMetrics,
+    RolloutResponse,
+)
+from .env_keys import MAX_ENVIRONMENT_API_KEY_BYTES, encrypt_for_backend, setup_environment_api_key
+from .secrets import mint_environment_api_key
+
+__all__ = [
+    "RolloutEnvSpec",
+    "RolloutPolicySpec",
+    "RolloutRecordConfig",
+    "RolloutSafetyConfig",
+    "RolloutRequest",
+    "RolloutStep",
+    "RolloutTrajectory",
+    "RolloutMetrics",
+    "RolloutResponse",
+    "encrypt_for_backend",
+    "setup_environment_api_key",
+    "mint_environment_api_key",
+    "MAX_ENVIRONMENT_API_KEY_BYTES",
+]
+

synth_ai/rl/contracts.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+"""
+Compatibility layer: re-export Task App rollout contracts from synth_ai.task.contracts
+so existing imports continue to work while consolidating under synth_ai.task.
+"""
+
+from synth_ai.task.contracts import (
+    RolloutEnvSpec,
+    RolloutPolicySpec,
+    RolloutRecordConfig,
+    RolloutSafetyConfig,
+    RolloutRequest,
+    RolloutStep,
+    RolloutTrajectory,
+    RolloutMetrics,
+    RolloutResponse,
+)
+
+__all__ = [
+    "RolloutEnvSpec",
+    "RolloutPolicySpec",
+    "RolloutRecordConfig",
+    "RolloutSafetyConfig",
+    "RolloutRequest",
+    "RolloutStep",
+    "RolloutTrajectory",
+    "RolloutMetrics",
+    "RolloutResponse",
+]
+
+

synth_ai/rl/env_keys.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+"""Helpers for uploading RL environment credentials to the backend."""
+
+import base64
+import binascii
+import json
+from typing import Any, Dict
+import os
+
+import requests
+from nacl.public import PublicKey, SealedBox
+
+__all__ = ["encrypt_for_backend", "setup_environment_api_key", "MAX_ENVIRONMENT_API_KEY_BYTES"]
+
+MAX_ENVIRONMENT_API_KEY_BYTES = 8 * 1024
+_ALGORITHM = "libsodium.sealedbox.v1"
+
+
+def encrypt_for_backend(pubkey_b64: str, secret: str | bytes) -> str:
+    """Encrypt ``secret`` for storage by the backend using libsodium sealed boxes."""
+
+    if not isinstance(pubkey_b64, str) or not pubkey_b64.strip():
+        raise ValueError("public key must be a non-empty base64 string")
+
+    try:
+        key_bytes = base64.b64decode(pubkey_b64, validate=True)
+    except binascii.Error as exc:  # pragma: no cover - defensive guard
+        raise ValueError("public key must be base64-encoded") from exc
+
+    if len(key_bytes) != 32:
+        raise ValueError("public key must be 32 bytes for X25519")
+
+    if isinstance(secret, str):
+        secret_bytes = secret.encode("utf-8")
+    elif isinstance(secret, bytes):
+        secret_bytes = secret
+    else:  # pragma: no cover - type guard
+        raise TypeError("secret must be str or bytes")
+
+    if not secret_bytes:
+        raise ValueError("secret must not be empty")
+
+    box = SealedBox(PublicKey(key_bytes))
+    ciphertext = box.encrypt(secret_bytes)
+    return base64.b64encode(ciphertext).decode("ascii")
+
+
+def setup_environment_api_key(
+    backend_base: str,
+    synth_api_key: str,
+    token: str | None = None,
+    *,
+    timeout: float = 15.0,
+) -> Dict[str, Any]:
+    """Upload an ENVIRONMENT_API_KEY to the backend."""
+
+    backend = backend_base.rstrip("/")
+    if not backend:
+        raise ValueError("backend_base must be provided")
+    if not synth_api_key:
+        raise ValueError("synth_api_key must be provided")
+
+    # Require caller-provided plaintext. If not provided, read from ENVIRONMENT_API_KEY.
+    plaintext = token if token is not None else os.getenv("ENVIRONMENT_API_KEY", "").strip()
+    if not plaintext:
+        raise ValueError("ENVIRONMENT_API_KEY must be set (or pass token=...) to upload")
+    if not isinstance(plaintext, str):  # pragma: no cover - defensive guard
+        raise TypeError("token must be a string")
+
+    token_bytes = plaintext.encode("utf-8")
+    if not token_bytes:
+        raise ValueError("ENVIRONMENT_API_KEY token must not be empty")
+    if len(token_bytes) > MAX_ENVIRONMENT_API_KEY_BYTES:
+        raise ValueError("ENVIRONMENT_API_KEY token exceeds 8 KiB limit")
+
+    headers = {"Authorization": f"Bearer {synth_api_key}"}
+    pub_url = f"{backend}/api/v1/crypto/public-key"
+    response = requests.get(pub_url, headers=headers, timeout=timeout)
+    _raise_with_detail(response)
+
+    try:
+        doc = response.json()
+    except ValueError as exc:  # pragma: no cover - backend invariant
+        raise RuntimeError("backend returned invalid JSON for public key") from exc
+
+    if not isinstance(doc, dict):
+        raise RuntimeError("backend public key response must be an object")
+
+    pubkey = doc.get("public_key")
+    if not isinstance(pubkey, str) or not pubkey:
+        raise RuntimeError("backend response missing public_key")
+
+    # The backend currently returns a single algorithm identifier; keep a guard in
+    # case future versions change the value and we need to surface that to callers.
+    alg = doc.get("alg")
+    if alg is not None and alg != _ALGORITHM:
+        raise RuntimeError(f"unsupported sealed box algorithm: {alg}")
+
+    ciphertext_b64 = encrypt_for_backend(pubkey, token_bytes)
+
+    body = {"name": "ENVIRONMENT_API_KEY", "ciphertext_b64": ciphertext_b64}
+    post_url = f"{backend}/api/v1/env-keys"
+    response2 = requests.post(post_url, headers={**headers, "Content-Type": "application/json"}, json=body, timeout=timeout)
+    _raise_with_detail(response2)
+
+    try:
+        upload_doc = response2.json()
+    except ValueError:
+        upload_doc = {}
+
+    if not isinstance(upload_doc, dict):
+        upload_doc = {}
+
+    return {
+        "stored": True,
+        "id": upload_doc.get("id"),
+        "name": upload_doc.get("name"),
+        "updated_at": upload_doc.get("updated_at"),
+    }
+
+
+def _raise_with_detail(response: requests.Response) -> None:
+    try:
+        response.raise_for_status()
+    except requests.HTTPError as exc:
+        detail_snippet: str | None = None
+        try:
+            detail = response.json()
+            detail_snippet = json.dumps(detail, separators=(",", ":"))[:200]
+        except Exception:
+            body = response.text if response.text is not None else ""
+            detail_snippet = body[:200] if body else None
+        message = str(exc)
+        if detail_snippet:
+            message = f"{message} | body={detail_snippet}"
+        raise requests.HTTPError(message, request=exc.request, response=exc.response) from None

synth_ai/rl/secrets.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+"""Helpers for generating RL environment credentials."""
+
+import secrets
+
+__all__ = ["mint_environment_api_key"]
+
+
+def mint_environment_api_key() -> str:
+    """Mint a random ENVIRONMENT_API_KEY value.
+
+    The current format is 64 hexadecimal characters (256 bits of entropy), which
+    matches the shell helpers used by the RL examples. This keeps the token easy
+    to copy while remaining suitably strong for authentication.
+    """
+
+    # secrets.token_hex(32) → 32 random bytes rendered as 64 hex characters.
+    return secrets.token_hex(32)

synth_ai/scripts/verify_rewards.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""
+Verify reward persistence in a traces database.
+
+Usage:
+  uv run python -m synth_ai.scripts.verify_rewards --db /path/to/db.sqlite --min-reward 1
+"""
+
+import argparse
+import asyncio
+import os
+from typing import Dict
+
+from sqlalchemy import text
+
+from synth_ai.tracing_v3.turso.manager import AsyncSQLTraceManager
+
+
+async def verify(db_path: str, min_reward: int) -> int:
+    db_url = db_path
+    if not db_url.startswith("sqlite+aiosqlite:///"):
+        db_url = f"sqlite+aiosqlite:///{os.path.abspath(db_path)}"
+
+    mgr = AsyncSQLTraceManager(db_url=db_url)
+    await mgr.initialize()
+
+    try:
+        async with mgr.session() as session:
+            # Sessions with outcome_rewards
+            q_good = text(
+                """
+                SELECT session_id, MAX(total_reward) as total_reward
+                FROM outcome_rewards
+                GROUP BY session_id
+                """
+            )
+            res = await session.execute(q_good)
+            outcomes = {row[0]: int(row[1]) for row in res.fetchall()}
+
+            # Sessions without outcome_rewards
+            q_missing = text(
+                """
+                SELECT s.session_id
+                FROM session_traces s
+                LEFT JOIN outcome_rewards o ON s.session_id = o.session_id
+                WHERE o.session_id IS NULL
+                """
+            )
+            res2 = await session.execute(q_missing)
+            missing = [row[0] for row in res2.fetchall()]
+
+            # Aggregate event_rewards per session (informational)
+            q_event = text(
+                """
+                SELECT session_id, COALESCE(SUM(reward_value), 0.0) as sum_rewards
+                FROM event_rewards
+                GROUP BY session_id
+                """
+            )
+            res3 = await session.execute(q_event)
+            event_sums: Dict[str, float] = {row[0]: float(row[1]) for row in res3.fetchall()}
+
+        print(f"Sessions with outcome_rewards: {len(outcomes)}")
+        print(f"Sessions missing outcome_rewards: {len(missing)}")
+        if missing:
+            print("Missing session_ids:", ", ".join(missing[:10]) + (" ..." if len(missing) > 10 else ""))
+
+        # Threshold check
+        qualifying = {sid: r for sid, r in outcomes.items() if r >= min_reward}
+        print(f"Sessions with total_reward >= {min_reward}: {len(qualifying)}")
+
+        # Show a small comparison snapshot
+        sample = list(qualifying.items())[:5]
+        for sid, tot in sample:
+            er = event_sums.get(sid, 0.0)
+            print(f"  {sid}: outcome={tot}, sum(event_rewards)={er:.2f}")
+
+        # Exit non-zero if any sessions are missing outcome rewards
+        if missing:
+            return 2
+        if min_reward > 0 and not qualifying:
+            return 3
+        return 0
+    finally:
+        await mgr.close()
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description="Verify reward persistence in traces DB")
+    ap.add_argument("--db", required=True, help="Path to traces SQLite DB (aiosqlite)")
+    ap.add_argument("--min-reward", type=int, default=0, help="Minimum total_reward to consider qualifying")
+    args = ap.parse_args()
+
+    return asyncio.run(verify(args.db, args.min_reward))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+
+

synth_ai/task/__init__.py

@@ -0,0 +1,10 @@
+from .validators import validate_task_app_url
+from .health import task_app_health
+from .contracts import TaskAppContract, TaskAppEndpoints
+
+__all__ = [
+    "validate_task_app_url",
+    "task_app_health",
+    "TaskAppContract",
+    "TaskAppEndpoints",
+]

synth_ai/task/contracts.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Any, Dict, List
+from pydantic import BaseModel
+
+
+@dataclass(frozen=True)
+class TaskAppEndpoints:
+    """Canonical Task App endpoint shapes used by RL trainers.
+
+    The Task App is an HTTP service (often deployed on Modal) that exposes:
+    - Health: GET /health
+      • Requires header X-API-Key (when ENVIRONMENT_API_KEY is configured)
+      • Returns { healthy: true }
+    - Environment lifecycle:
+      • POST /env/{env_name}/initialize → { env_id, observation }
+      • POST /env/{env_name}/step      → { observation, reward, done, info }
+      • POST /env/{env_name}/terminate → { ok: true }
+    - Rollout (optional, unified schema):
+      • POST /rollout → { run_id, trajectories[], metrics, ... }
+    - Proxy (optional):
+      • POST /proxy/v1/chat/completions (for direct OpenAI calls from Task App)
+    """
+
+    health: str = "/health"
+    rollout: str = "/rollout"
+    proxy_chat_completions: str = "/proxy/v1/chat/completions"
+    env_initialize: str = "/env/{env_name}/initialize"
+    env_step: str = "/env/{env_name}/step"
+    env_terminate: str = "/env/{env_name}/terminate"
+
+
+@dataclass(frozen=True)
+class TaskAppContract:
+    """Requirements and expectations for a Task App used by RL trainers.
+
+    - Auth: ENVIRONMENT_API_KEY must be set in the Task App environment; requests include X-API-Key.
+    - Health: /health returns 200 and JSON; may verify X-API-Key header.
+    - Env API: initialize/step/terminate are present for the target env (e.g., CrafterClassic).
+    - Rollout API: optional; provides a single-call rollout for convenience/testing.
+    - Inference routing: policy config passes an inference_url (Synth backend or OpenAI proxy).
+    - URL: base must be reachable via HTTPS and should be under .modal.run in production.
+    """
+
+    base_url: str
+    env_name: Optional[str] = None
+    requires_api_key_header: bool = True
+
+
+# --- Unified rollout schema used by Task App services and SDK utilities ---
+
+class RolloutEnvSpec(BaseModel):
+    env_id: Optional[str] = None
+    env_name: Optional[str] = None
+    config: Dict[str, Any] = {}
+    seed: Optional[int] = None
+
+
+class RolloutPolicySpec(BaseModel):
+    policy_id: Optional[str] = None
+    policy_name: Optional[str] = None
+    config: Dict[str, Any] = {}
+
+
+class RolloutRecordConfig(BaseModel):
+    trajectories: bool = True
+    logprobs: bool = False
+    value: bool = False
+
+
+class RolloutSafetyConfig(BaseModel):
+    max_ops: int = 100000
+    max_time_s: float = 3600.0
+
+
+class RolloutRequest(BaseModel):
+    run_id: str
+    env: RolloutEnvSpec
+    policy: RolloutPolicySpec
+    ops: List[Dict[str, Any]] | List[str]
+    record: RolloutRecordConfig = RolloutRecordConfig()
+    on_done: str = "reset"
+    safety: RolloutSafetyConfig = RolloutSafetyConfig()
+    training_session_id: Optional[str] = None
+    synth_base_url: Optional[str] = None
+
+
+class RolloutStep(BaseModel):
+    obs: Dict[str, Any]
+    tool_calls: List[Dict[str, Any]]
+    reward: Optional[float] = None
+    done: bool = False
+    truncated: Optional[bool] = None
+    info: Optional[Dict[str, Any]] = None
+
+
+class RolloutTrajectory(BaseModel):
+    env_id: str
+    policy_id: str
+    steps: List[RolloutStep]
+    final: Optional[Dict[str, Any]] = None
+    length: int
+
+
+class RolloutMetrics(BaseModel):
+    episode_returns: List[float]
+    mean_return: float
+    num_steps: int
+    num_episodes: int = 0
+
+
+class RolloutResponse(BaseModel):
+    run_id: str
+    trajectories: List[RolloutTrajectory]
+    branches: Dict[str, List[str]] = {}
+    metrics: RolloutMetrics
+    aborted: bool = False
+    ops_executed: int = 0
+

synth_ai/task/health.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+import aiohttp
+
+
+async def task_app_health(task_app_url: str) -> Dict[str, Any]:
+    """Probe a Task App base URL for basic reachability.
+
+    Behavior:
+    - Try HEAD first (follows redirects)
+    - Fallback to GET if HEAD is unsupported
+    - Returns {ok: bool, status?: int, error?: str}
+    """
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.head(task_app_url, allow_redirects=True) as r:
+                if 200 <= r.status < 400:
+                    return {"ok": True, "status": r.status}
+        async with aiohttp.ClientSession() as session:
+            async with session.get(task_app_url, allow_redirects=True) as r2:
+                if 200 <= r2.status < 400:
+                    return {"ok": True, "status": r2.status}
+        return {"ok": False, "status": None}
+    except Exception as e:
+        return {"ok": False, "error": f"{type(e).__name__}: {e}"}
+
+

synth_ai/task/validators.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from urllib.parse import urlparse
+
+
+def validate_task_app_url(url: str, *, name: str = "TASK_APP_BASE_URL") -> None:
+    """Validate a Task App base URL (scheme + host present)."""
+
+    p = urlparse(url)
+    if p.scheme not in ("http", "https") or not p.netloc:
+        raise ValueError(f"Invalid {name}: malformed: {url}")
+

synth_ai/tracing/__init__.py

@@ -1,18 +1,30 @@
-import sys as _sys
 import importlib as _importlib
+import sys as _sys
 
-_pkg = _importlib.import_module('synth_ai.v0.tracing')
+_pkg = _importlib.import_module("synth_ai.v0.tracing")
 _sys.modules[__name__] = _pkg
 
 _SUBMODULES = [
-    'abstractions', 'base_client', 'client_manager', 'config', 'context',
-    'decorators', 'immediate_client', 'local', 'log_client_base', 'retry_queue',
-    'trackers', 'upload', 'utils'
+    "abstractions",
+    "base_client",
+    "client_manager",
+    "config",
+    "context",
+    "decorators",
+    "immediate_client",
+    "local",
+    "log_client_base",
+    "retry_queue",
+    "trackers",
+    "upload",
+    "utils",
 ]
 for _m in _SUBMODULES:
-    _sys.modules[f'{__name__}.{_m}'] = _importlib.import_module(f'synth_ai.v0.tracing.{_m}')
+    _sys.modules[f"{__name__}.{_m}"] = _importlib.import_module(f"synth_ai.v0.tracing.{_m}")
 
-_events_pkg = _importlib.import_module('synth_ai.v0.tracing.events')
-_sys.modules[f'{__name__}.events'] = _events_pkg
-for _m in ['manage', 'scope', 'store']:
-    _sys.modules[f'{__name__}.events.{_m}'] = _importlib.import_module(f'synth_ai.v0.tracing.events.{_m}')
+_events_pkg = _importlib.import_module("synth_ai.v0.tracing.events")
+_sys.modules[f"{__name__}.events"] = _events_pkg
+for _m in ["manage", "scope", "store"]:
+    _sys.modules[f"{__name__}.events.{_m}"] = _importlib.import_module(
+        f"synth_ai.v0.tracing.events.{_m}"
+    )

synth_ai/tracing_v1/__init__.py

@@ -1,31 +1,33 @@
-import sys as _sys
 import importlib as _importlib
+import sys as _sys
 
 # Forward top-level package
-_pkg = _importlib.import_module('synth_ai.v0.tracing_v1')
+_pkg = _importlib.import_module("synth_ai.v0.tracing_v1")
 _sys.modules[__name__] = _pkg
 
 # Explicitly forward submodules so `synth_ai.tracing_v1.X` works
 _SUBMODULES = [
-    'abstractions',
-    'base_client',
-    'client_manager',
-    'config',
-    'context',
-    'decorators',
-    'immediate_client',
-    'local',
-    'log_client_base',
-    'retry_queue',
-    'trackers',
-    'upload',
-    'utils',
+    "abstractions",
+    "base_client",
+    "client_manager",
+    "config",
+    "context",
+    "decorators",
+    "immediate_client",
+    "local",
+    "log_client_base",
+    "retry_queue",
+    "trackers",
+    "upload",
+    "utils",
 ]
 for _m in _SUBMODULES:
-    _sys.modules[f'{__name__}.{_m}'] = _importlib.import_module(f'synth_ai.v0.tracing_v1.{_m}')
+    _sys.modules[f"{__name__}.{_m}"] = _importlib.import_module(f"synth_ai.v0.tracing_v1.{_m}")
 
 # Forward events package and its submodules
-_events_pkg = _importlib.import_module('synth_ai.v0.tracing_v1.events')
-_sys.modules[f'{__name__}.events'] = _events_pkg
-for _m in ['manage', 'scope', 'store']:
-    _sys.modules[f'{__name__}.events.{_m}'] = _importlib.import_module(f'synth_ai.v0.tracing_v1.events.{_m}')
+_events_pkg = _importlib.import_module("synth_ai.v0.tracing_v1.events")
+_sys.modules[f"{__name__}.events"] = _events_pkg
+for _m in ["manage", "scope", "store"]:
+    _sys.modules[f"{__name__}.events.{_m}"] = _importlib.import_module(
+        f"synth_ai.v0.tracing_v1.events.{_m}"
+    )

synth_ai/tracing_v3/__init__.py

@@ -1,7 +1,7 @@
 """Tracing v3 - Turso/sqld based tracing implementation.
 
 This module provides a modern, async-first tracing system for capturing and storing
-detailed execution traces from AI systems. It's designed to handle high-throughput 
+detailed execution traces from AI systems. It's designed to handle high-throughput
 scenarios with proper async/await patterns throughout.
 
 Architecture Overview:
@@ -49,10 +49,10 @@ Key Features:
 Usage Example:
 -------------
     from synth_ai.tracing_v3 import SessionTracer
-    
+
     tracer = SessionTracer()
     await tracer.initialize()
-    
+
     async with tracer.session() as session_id:
         async with tracer.timestep("step1", turn_number=1):
             # Record events during execution
@@ -70,17 +70,17 @@ The system uses environment variables for configuration:
 See `config.py` for full configuration options.
 """
 
-from .session_tracer import SessionTracer
 from .abstractions import (
-    SessionTrace,
-    SessionTimeStep,
     BaseEvent,
-    RuntimeEvent,
     EnvironmentEvent,
+    RuntimeEvent,
     SessionEventMarkovBlanketMessage,
+    SessionTimeStep,
+    SessionTrace,
     TimeRecord,
 )
 from .config import TursoConfig
+from .session_tracer import SessionTracer
 
 __all__ = [
     "SessionTracer",