synth-ai 0.2.9.dev0__py3-none-any.whl → 0.2.23.dev3__py3-none-any.whl

synth_ai/task/contracts.py

@@ -1,18 +1,25 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Optional, Any, Dict, List, Literal
-from pydantic import BaseModel, Field
+from enum import Enum
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class RolloutMode(str, Enum):
+    """Mode controls how rollout infrastructure processes inference URLs."""
+    RL = "rl"
+    EVAL = "eval"
 
 
 @dataclass(frozen=True)
 class TaskAppEndpoints:
-    """Canonical Task App endpoint shapes used by RL trainers.
+    """Required Task App endpoints used by RL trainers and clients.
 
-    Task Apps run as lightweight HTTP services (often on Modal) that expose a
-    consistent set of endpoints for health, metadata, environment lifecycle,
-    rollouts, and optional proxy access to vendor models. The endpoint strings
-    defined here act as defaults and documentation for clients.
+    Task Apps run as lightweight HTTP services (often on Modal) that expose these
+    standard endpoints. Additional endpoints (proxies, debug routes) may be added
+    by individual task apps as needed.
     """
 
     root: str = "/"
@@ -20,43 +27,22 @@ class TaskAppEndpoints:
     info: str = "/info"
     task_info: str = "/task_info"
     rollout: str = "/rollout"
-    proxy_chat_completions: str = "/proxy/v1/chat/completions"
-    proxy_groq_chat_completions: str = "/proxy/groq/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
+    env_id: str | None = None
+    env_name: str | None = None
+    config: dict[str, Any] = Field(default_factory=dict)
+    seed: int | None = None
 
 
 class RolloutPolicySpec(BaseModel):
-    policy_id: Optional[str] = None
-    policy_name: Optional[str] = None
-    config: Dict[str, Any] = {}
+    policy_id: str | None = None
+    policy_name: str | None = None
+    config: dict[str, Any] = Field(default_factory=dict)
 
 
 class RolloutRecordConfig(BaseModel):
@@ -64,7 +50,7 @@ class RolloutRecordConfig(BaseModel):
     logprobs: bool = False
     value: bool = False
     return_trace: bool = False
-    trace_format: Literal["compact", "full"] = "compact"
+    trace_format: Literal["compact", "full", "structured"] = "compact"
 
 
 class RolloutSafetyConfig(BaseModel):
@@ -76,60 +62,175 @@ class RolloutRequest(BaseModel):
     run_id: str
     env: RolloutEnvSpec
     policy: RolloutPolicySpec
-    ops: List[Dict[str, Any]] | List[str]
+    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
+    training_session_id: str | None = None
+    synth_base_url: str | None = None
+    mode: RolloutMode  # Required: explicit RL vs EVAL mode
 
 
 class RolloutStep(BaseModel):
-    obs: Dict[str, Any]
-    tool_calls: List[Dict[str, Any]]
-    reward: Optional[float] = None
+    """Single step in a rollout trajectory.
+    
+    DEPRECATED: This is part of the legacy trajectory format. New code should
+    consume v3 traces (RolloutResponse.trace) instead. See monorepo/trace_single_source.txt
+    for migration plan.
+    """
+    obs: dict[str, Any]
+    tool_calls: list[dict[str, Any]]
+    reward: float | None = None
     done: bool = False
-    truncated: Optional[bool] = None
-    info: Optional[Dict[str, Any]] = None
+    truncated: bool | None = None
+    info: dict[str, Any] | None = None
 
 
 class RolloutTrajectory(BaseModel):
+    """Legacy trajectory format for rollout results.
+    
+    DEPRECATED: This format duplicates data already present in v3 traces and will
+    be removed once training code migrates to consuming RolloutResponse.trace.
+    
+    Current state:
+    - Task apps emit BOTH this format AND v3 traces (dual serialization)
+    - Training code (GSPO) reads from this format
+    - Eval/filter tools read from v3 traces
+    
+    Migration plan:
+    - Phase 1: Training code learns to read from v3 traces (with fallback to this)
+    - Phase 2: Make this field optional once training is migrated
+    - Phase 3: Remove this field entirely and delete this class
+    
+    See: monorepo/trace_single_source.txt for full migration plan and timeline.
+    
+    Why v3 traces are better:
+    - Single source of truth (no duplication/drift)
+    - Richer data: token IDs, logprobs, reasoning, timing, images
+    - Built-in audit trail and replay capability
+    - Standard schema across all Synth AI tooling
+    """
     env_id: str
     policy_id: str
-    steps: List[RolloutStep]
-    final: Optional[Dict[str, Any]] = None
+    steps: list[RolloutStep]
+    final: dict[str, Any] | None = None
     length: int
+    
+    # Required for trace correlation with inference mesh (optional initially for backward compat)
+    # See: monorepo/INFERENCE_URL_REQUIREMENT_PLAN.md and trace_creation_and_judgement.txt
+    inference_url: str
+    
+    decision_samples: list[dict[str, Any]] | None = None
 
 
 class RolloutMetrics(BaseModel):
-    episode_returns: List[float]
+    episode_returns: list[float]
     mean_return: float
     num_steps: int
     num_episodes: int = 0
-    outcome_score: Optional[float] = None
-    events_score: Optional[float] = None
-    details: Dict[str, Any] = Field(default_factory=dict)
+    outcome_score: float | None = None
+    events_score: float | None = None
+    details: dict[str, Any] = Field(default_factory=dict)
 
 
 class RolloutResponse(BaseModel):
+    """Response from a rollout execution.
+    
+    Contains both legacy trajectory format (for backward compatibility) and
+    modern v3 trace format (preferred going forward).
+    """
     run_id: str
-    trajectories: List[RolloutTrajectory]
-    branches: Dict[str, List[str]] = {}
+    
+    # DEPRECATED: Legacy format maintained for training code compatibility.
+    # Will be removed once training migrates to reading from `trace` field.
+    # See: monorepo/trace_single_source.txt for migration plan.
+    trajectories: list[RolloutTrajectory]
+    
+    branches: dict[str, list[str]] = Field(default_factory=dict)
     metrics: RolloutMetrics
     aborted: bool = False
     ops_executed: int = 0
-    trace: Dict[str, Any] | None = None
+    
+    # OPTIONAL: correlation ID for linking rollout to inference traces
+    # If not provided, trainer will infer it from trajectory.inference_url ?cid=... parameter
+    trace_correlation_id: str | None = None
+    
+    # PREFERRED: v3 trace format (SessionTrace). This is the single source of truth
+    # for rollout data and should be used by all new code. Contains richer data than
+    # trajectories including token IDs, logprobs, timing, and multimodal content.
+    trace: dict[str, Any] | None = None
+    pipeline_metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class _ExtraAllowModel(BaseModel):
+    """Base helper that preserves unknown keys while still exposing typed attributes."""
+
+    model_config = ConfigDict(extra="allow")
+
+
+class TaskDescriptor(_ExtraAllowModel):
+    """Human-readable task identifiers shown in UIs and logs."""
+
+    id: str
+    name: str
+    description: str | None = None
+    version: str | None = None
+
+
+class DatasetInfo(_ExtraAllowModel):
+    """Metadata about the prompt/task dataset powering the environment."""
+
+    id: str | None = None
+    name: str | None = None
+    version: str | None = None
+    splits: list[str] | None = None
+    default_split: str | None = None
+    description: str | None = None
+
+
+class RubricCriterion(_ExtraAllowModel):
+    id: str
+    description: str
+    weight: float | None = None
+
+
+class RubricSection(_ExtraAllowModel):
+    name: str
+    criteria: list[RubricCriterion] = Field(default_factory=list)
+
+
+class RubricInfo(_ExtraAllowModel):
+    """Outcome and event scoring definitions used by judges."""
+
+    outcome: RubricSection | None = None
+    events: RubricSection | None = None
+
+
+class InferenceInfo(_ExtraAllowModel):
+    """Recommended defaults for policy model routing."""
+
+    model: str | None = None
+    inference_url: str | None = None
+
+
+class LimitsInfo(_ExtraAllowModel):
+    """Operational limits the environment enforces."""
+
+    max_turns: int | None = None
+    max_response_tokens: int | None = None
+    timeout_seconds: int | None = None
 
 
-class TaskInfo(BaseModel):
+class TaskInfo(_ExtraAllowModel):
     """Static metadata describing the capabilities of a Task App task."""
 
-    task: Dict[str, Any]
-    environments: List[str]
-    action_space: Dict[str, Any]
-    observation: Dict[str, Any]
-    dataset: Dict[str, Any]
-    rubric: Dict[str, Any]
-    inference: Dict[str, Any]
-    capabilities: Dict[str, Any]
-    limits: Dict[str, Any]
+    task: TaskDescriptor
+    environment: str
+    dataset: DatasetInfo
+    rubric: RubricInfo
+    inference: InferenceInfo
+    limits: LimitsInfo
+    task_metadata: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Task-specific extras (e.g. prompt version info, documentation links).",
+    )

synth_ai/task/datasets.py

@@ -1,8 +1,9 @@
-from __future__ import annotations
-
 """Dataset registry and helpers shared by Task Apps."""
 
-from typing import Any, Callable, Dict, Hashable, Tuple
+from __future__ import annotations
+
+from collections.abc import Callable, Hashable
+from typing import Any
 
 from pydantic import BaseModel, Field, field_validator
 
@@ -34,10 +35,12 @@ class TaskDatasetRegistry:
     """Lightweight registry mapping dataset specs to loader callables."""
 
     def __init__(self) -> None:
-        self._entries: Dict[str, Tuple[TaskDatasetSpec, RegistryLoader, bool]] = {}
-        self._cache: Dict[Hashable, Any] = {}
+        self._entries: dict[str, tuple[TaskDatasetSpec, RegistryLoader, bool]] = {}
+        self._cache: dict[Hashable, Any] = {}
 
-    def register(self, spec: TaskDatasetSpec, loader: RegistryLoader, *, cache: bool = True) -> None:
+    def register(
+        self, spec: TaskDatasetSpec, loader: RegistryLoader, *, cache: bool = True
+    ) -> None:
         """Register a dataset loader and its metadata."""
 
         self._entries[spec.id] = (spec, loader, cache)

synth_ai/task/errors.py

@@ -1,14 +1,16 @@
-from __future__ import annotations
-
 """Error helpers used across Task App implementations."""
 
-from typing import Any, Dict, Optional
+from __future__ import annotations
+
+from typing import Any
 
 from .json import to_jsonable
 
 
-def error_payload(code: str, message: str, *, extra: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
-    payload: Dict[str, Any] = {"error": {"code": code, "message": message}}
+def error_payload(
+    code: str, message: str, *, extra: dict[str, Any] | None = None
+) -> dict[str, Any]:
+    payload: dict[str, Any] = {"error": {"code": code, "message": message}}
     if extra:
         payload["error"].update(extra)
     return payload
@@ -19,8 +21,8 @@ def http_exception(
     code: str,
     message: str,
     *,
-    extra: Optional[Dict[str, Any]] = None,
-    headers: Optional[Dict[str, str]] = None,
+    extra: dict[str, Any] | None = None,
+    headers: dict[str, str] | None = None,
 ):
     try:
         from fastapi import HTTPException  # type: ignore
@@ -36,8 +38,8 @@ def json_error_response(
     code: str,
     message: str,
     *,
-    extra: Optional[Dict[str, Any]] = None,
-    headers: Optional[Dict[str, str]] = None,
+    extra: dict[str, Any] | None = None,
+    headers: dict[str, str] | None = None,
 ):
     try:
         from fastapi.responses import JSONResponse  # type: ignore
@@ -46,4 +48,3 @@ def json_error_response(
 
     payload = error_payload(code, message, extra=extra)
     return JSONResponse(status_code=status_code, content=to_jsonable(payload), headers=headers)
-

synth_ai/task/health.py

@@ -1,10 +1,13 @@
+"""Helpers for probing Task App health endpoints."""
+
 from __future__ import annotations
 
-from typing import Any, Dict
+from typing import Any
+
 import aiohttp
 
 
-async def task_app_health(task_app_url: str) -> Dict[str, Any]:
+async def task_app_health(task_app_url: str) -> dict[str, Any]:
     """Probe a Task App base URL for basic reachability.
 
     Behavior:
@@ -12,17 +15,20 @@ async def task_app_health(task_app_url: str) -> Dict[str, Any]:
     - Fallback to GET if HEAD is unsupported
     - Returns {ok: bool, status?: int, error?: str}
     """
+
+    async def _try_request(session: aiohttp.ClientSession, method: str) -> dict[str, Any] | None:
+        request = getattr(session, method)
+        async with request(task_app_url, allow_redirects=True) as response:
+            if 200 <= response.status < 400:
+                return {"ok": True, "status": response.status}
+        return None
+
     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}
+            for method in ("head", "get"):
+                result = await _try_request(session, method)
+                if result is not None:
+                    return result
         return {"ok": False, "status": None}
     except Exception as e:
         return {"ok": False, "error": f"{type(e).__name__}: {e}"}
-
-

synth_ai/task/inference_api.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, List, Optional, Sequence
+
+
+class InferenceAPIClient:
+    """Async client that normalizes chat completion calls across providers."""
+
+    def __init__(self, *, provider: str, inference_url: Optional[str] = None, timeout: float = 30.0) -> None:
+        self.provider = (provider or "").lower()
+        self.inference_url = inference_url.rstrip("/") if inference_url else None
+        self.timeout = timeout
+
+        if self.provider == "openai":
+            try:
+                from openai import AsyncOpenAI
+            except Exception as exc:  # pragma: no cover - import guard
+                raise RuntimeError("openai package not installed") from exc
+
+            api_key = os.getenv("OPENAI_API_KEY")
+            if not api_key:
+                raise RuntimeError("OPENAI_API_KEY must be set for meta-model proposals")
+
+            kwargs: Dict[str, Any] = {"api_key": api_key}
+            if self.inference_url:
+                kwargs["base_url"] = self.inference_url
+            self._client = AsyncOpenAI(**kwargs)
+
+        elif self.provider == "groq":
+            try:
+                from groq import AsyncGroq
+            except Exception as exc:  # pragma: no cover - import guard
+                raise RuntimeError("groq package not installed") from exc
+
+            api_key = os.getenv("GROQ_API_KEY")
+            if not api_key:
+                raise RuntimeError("GROQ_API_KEY must be set for meta-model proposals")
+
+            kwargs: Dict[str, Any] = {"api_key": api_key}
+            if self.inference_url:
+                kwargs["base_url"] = self.inference_url
+            self._client = AsyncGroq(**kwargs)
+
+        else:
+            raise RuntimeError(f"Unsupported meta-model provider: {provider}")
+
+    async def chat_completion(
+        self,
+        *,
+        messages: Sequence[Dict[str, Any]],
+        model: str,
+        temperature: float,
+        max_tokens: int,
+    ) -> Dict[str, Any]:
+        if self.provider == "openai":
+            response = await self._client.chat.completions.create(  # type: ignore[attr-defined]
+                model=model,
+                messages=self._normalize_messages(messages),
+                temperature=temperature,
+                max_tokens=max_tokens,
+                timeout=self.timeout,
+            )
+            return self._model_dump(response)
+
+        if self.provider == "groq":
+            response = await self._client.chat.completions.create(  # type: ignore[attr-defined]
+                model=model,
+                messages=self._normalize_messages(messages),
+                temperature=temperature,
+                max_tokens=max_tokens,
+                timeout=self.timeout,
+            )
+            return self._model_dump(response)
+
+        raise RuntimeError(f"Unsupported meta-model provider: {self.provider}")
+
+    @staticmethod
+    def _normalize_messages(messages: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        return [dict(message) for message in messages]
+
+    @staticmethod
+    def _model_dump(obj: Any) -> Dict[str, Any]:
+        if hasattr(obj, "model_dump"):
+            data = obj.model_dump()
+        elif hasattr(obj, "dict"):
+            data = obj.dict()
+        else:
+            data = dict(obj)  # type: ignore[arg-type]
+
+        usage = getattr(obj, "usage", None)
+        if usage and "usage" not in data:
+            if hasattr(usage, "model_dump"):
+                data["usage"] = usage.model_dump()
+            elif isinstance(usage, dict):
+                data["usage"] = usage
+        return data
+
+
+__all__ = ["InferenceAPIClient"]
+

synth_ai/task/json.py

@@ -1,9 +1,9 @@
-from __future__ import annotations
-
 """Shared JSON sanitisation helpers for Task Apps."""
 
+from __future__ import annotations
+
 from collections.abc import Mapping, Sequence
-from dataclasses import is_dataclass, asdict
+from dataclasses import asdict, is_dataclass
 from enum import Enum
 from typing import Any
 
@@ -13,13 +13,19 @@ except Exception:  # pragma: no cover - handled at runtime
     _np = None  # type: ignore
 
 
-def _mask_numpy_array(arr: "_np.ndarray") -> str:
+def _mask_numpy_array(arr: Any) -> str:
     shape = getattr(arr, "shape", None)
     dtype = getattr(arr, "dtype", None)
     return f"<ndarray shape={shape} dtype={dtype}>"
 
 
-def to_jsonable(value: Any) -> Any:
+def to_jsonable(
+    value: Any,
+    *,
+    _visited: set[int] | None = None,
+    _depth: int = 0,
+    _max_depth: int = 32,
+) -> Any:
     """Convert `value` into structures compatible with JSON serialisation.
 
     - numpy scalars are converted to their Python counterparts
@@ -29,25 +35,33 @@ def to_jsonable(value: Any) -> Any:
     - non-serialisable objects fall back to `repr`
     """
 
-    if value is None or isinstance(value, (str, bool, int, float)):
+    if _visited is None:
+        _visited = set()
+
+    if _depth > _max_depth:
+        return f"<max_depth type={type(value).__name__}>"
+
+    if value is None or isinstance(value, str | bool | int | float):
         return value
 
     # numpy scalars / arrays
     if _np is not None:
-        if isinstance(value, (_np.integer,)):
+        if isinstance(value, _np.integer):
             return int(value)
-        if isinstance(value, (_np.floating,)):
+        if isinstance(value, _np.floating):
             return float(value)
-        if isinstance(value, (_np.bool_,)):
+        if isinstance(value, _np.bool_):
             return bool(value)
-        if isinstance(value, (_np.ndarray,)):
+        if isinstance(value, _np.ndarray):
             return _mask_numpy_array(value)
 
     if isinstance(value, Enum):
-        return to_jsonable(value.value)
+        return to_jsonable(value.value, _visited=_visited, _depth=_depth + 1, _max_depth=_max_depth)
 
     if is_dataclass(value):
-        return to_jsonable(asdict(value))
+        return to_jsonable(
+            asdict(value), _visited=_visited, _depth=_depth + 1, _max_depth=_max_depth
+        )
 
     # pydantic BaseModel / attrs objects
     for attr in ("model_dump", "dict", "to_dict", "to_json"):
@@ -56,22 +70,42 @@ def to_jsonable(value: Any) -> Any:
                 dumped = getattr(value, attr)()  # type: ignore[misc]
             except TypeError:
                 dumped = getattr(value, attr)(exclude_none=False)  # pragma: no cover
-            return to_jsonable(dumped)
+            return to_jsonable(
+                dumped, _visited=_visited, _depth=_depth + 1, _max_depth=_max_depth
+            )
 
-    if isinstance(value, Mapping):
-        return {str(k): to_jsonable(v) for k, v in value.items()}
-
-    if isinstance(value, (set, tuple)):
-        return [to_jsonable(v) for v in value]
-
-    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
-        return [to_jsonable(v) for v in value]
+    obj_id = id(value)
+    if obj_id in _visited:
+        return f"<circular type={type(value).__name__}>"
 
-    if isinstance(value, (bytes, bytearray)):
+    if isinstance(value, Mapping):
+        _visited.add(obj_id)
+        return {
+            str(k): to_jsonable(v, _visited=_visited, _depth=_depth + 1, _max_depth=_max_depth)
+            for k, v in value.items()
+        }
+
+    if isinstance(value, set | tuple):
+        _visited.add(obj_id)
+        return [
+            to_jsonable(v, _visited=_visited, _depth=_depth + 1, _max_depth=_max_depth)
+            for v in value
+        ]
+
+    if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray):
+        _visited.add(obj_id)
+        return [
+            to_jsonable(v, _visited=_visited, _depth=_depth + 1, _max_depth=_max_depth)
+            for v in value
+        ]
+
+    if isinstance(value, bytes | bytearray):
         return f"<bytes len={len(value)}>"
 
     if hasattr(value, "__dict__"):
-        return to_jsonable(vars(value))
+        _visited.add(obj_id)
+        return to_jsonable(
+            vars(value), _visited=_visited, _depth=_depth + 1, _max_depth=_max_depth
+        )
 
     return repr(value)
-