synth-ai 0.2.8.dev12__py3-none-any.whl → 0.2.9.dev0__py3-none-any.whl

synth_ai/task/auth.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+"""Authentication helpers shared by Task Apps."""
+
+import os
+from typing import Iterable, Optional, Any, Set
+
+from .errors import http_exception
+
+_API_KEY_ENV = "ENVIRONMENT_API_KEY"
+_DEV_API_KEY_ENVS = ("dev_environment_api_key", "DEV_ENVIRONMENT_API_KEY")
+_API_KEY_HEADER = "x-api-key"
+_API_KEYS_HEADER = "x-api-keys"
+_AUTH_HEADER = "authorization"
+_API_KEY_ALIASES_ENV = "ENVIRONMENT_API_KEY_ALIASES"  # comma-separated list of additional valid keys
+
+
+def _mask(value: str, *, prefix: int = 4) -> str:
+    if not value:
+        return "<empty>"
+    visible = value[:prefix]
+    return f"{visible}{'…' if len(value) > prefix else ''}"
+
+
+def normalize_environment_api_key() -> Optional[str]:
+    """Ensure `ENVIRONMENT_API_KEY` is populated from dev fallbacks.
+
+    Returns the resolved key (if any) so callers can branch on configuration.
+    """
+
+    key = os.getenv(_API_KEY_ENV)
+    if key:
+        return key
+    for env in _DEV_API_KEY_ENVS:
+        candidate = os.getenv(env)
+        if candidate:
+            os.environ[_API_KEY_ENV] = candidate
+            print(
+                f"[task:auth] {_API_KEY_ENV} set from {env} (prefix={_mask(candidate)})",
+                flush=True,
+            )
+            return candidate
+    return None
+
+
+def allowed_environment_api_keys() -> Set[str]:
+    """Return the set of valid environment API keys for this Task App.
+
+    Includes:
+    - The primary ENVIRONMENT_API_KEY (normalized from dev fallbacks if needed)
+    - Any comma-separated aliases from ENVIRONMENT_API_KEY_ALIASES
+    """
+    keys: set[str] = set()
+    primary = normalize_environment_api_key()
+    if primary:
+        keys.add(primary)
+    aliases = (os.getenv(_API_KEY_ALIASES_ENV) or "").strip()
+    if aliases:
+        for part in aliases.split(","):
+            trimmed = part.strip()
+            if trimmed:
+                keys.add(trimmed)
+    return keys
+
+
+def _header_values(request: Any, header: str) -> Iterable[str]:
+    header_lower = header.lower()
+    if request is None:
+        return []
+    headers = getattr(request, "headers", None)
+    if headers:
+        raw = headers.get(header) or headers.get(header_lower)
+        if raw is not None:
+            return [raw]
+    if isinstance(request, dict):
+        raw = request.get(header) or request.get(header_lower)
+        if raw is not None:
+            return [raw]
+    # Support passing explicit header dict via keyword arg on FastAPI route handlers
+    for attr in ("headers", "state"):
+        maybe = getattr(request, attr, None)
+        if isinstance(maybe, dict):
+            raw = maybe.get(header) or maybe.get(header_lower)
+            if raw is not None:
+                return [raw]
+    return []
+
+
+def _split_csv(values: Iterable[str]) -> list[str]:
+    seen: list[str] = []
+    for v in values:
+        if not isinstance(v, str):
+            continue
+        for part in v.split(","):
+            trimmed = part.strip()
+            if trimmed:
+                seen.append(trimmed)
+    return seen
+
+
+def is_api_key_header_authorized(request: Any) -> bool:
+    """Return True if any header-provided key matches any allowed environment key."""
+
+    allowed = allowed_environment_api_keys()
+    if not allowed:
+        return False
+    single = list(_header_values(request, _API_KEY_HEADER))
+    multi = list(_header_values(request, _API_KEYS_HEADER))
+    auths = list(_header_values(request, _AUTH_HEADER))
+    bearer: list[str] = []
+    for a in auths:
+        if isinstance(a, str) and a.lower().startswith("bearer "):
+            bearer.append(a.split(" ", 1)[1].strip())
+    candidates = _split_csv(single + multi + bearer)
+    return any(candidate in allowed for candidate in candidates)
+
+
+def require_api_key_dependency(request: Any) -> None:
+    """FastAPI dependency enforcing Task App authentication headers."""
+
+    allowed = allowed_environment_api_keys()
+    if not allowed:
+        raise http_exception(503, "missing_environment_api_key", "ENVIRONMENT_API_KEY is not configured")
+    # Build candidate list for verbose diagnostics
+    single = list(_header_values(request, _API_KEY_HEADER))
+    multi = list(_header_values(request, _API_KEYS_HEADER))
+    auths = list(_header_values(request, _AUTH_HEADER))
+    bearer: list[str] = []
+    for a in auths:
+        if isinstance(a, str) and a.lower().startswith("bearer "):
+            bearer.append(a.split(" ", 1)[1].strip())
+    candidates = _split_csv(single + multi + bearer)
+    if not any(candidate in allowed for candidate in candidates):
+        try:
+            print({
+                "task_auth_failed": True,
+                "allowed_first15": [k[:15] for k in allowed],
+                "allowed_count": len(allowed),
+                "got_first15": [c[:15] for c in candidates],
+                "got_lens": [len(c) for c in candidates],
+                "have_x_api_key": bool(single),
+                "have_x_api_keys": bool(multi),
+                "have_authorization": bool(auths),
+            }, flush=True)
+        except Exception:
+            pass
+        # Use 400 to make failures unmistakable during preflight
+        raise http_exception(400, "unauthorised", "API key missing or invalid", extra={
+            "allowed_first15": [k[:15] for k in allowed],
+            "allowed_count": len(allowed),
+            "got_first15": [c[:15] for c in candidates],
+            "got_lens": [len(c) for c in candidates],
+        })

synth_ai/task/client.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+"""Async HTTP client for interacting with Task Apps."""
+
+import asyncio
+from typing import Any, Dict, Iterable, List, Optional
+import os
+
+import httpx
+from pydantic import BaseModel
+
+from .contracts import RolloutRequest, RolloutResponse, TaskInfo
+from .json import to_jsonable
+
+
+def _prepare_payload(payload: Any) -> Any:
+    if payload is None:
+        return None
+    if isinstance(payload, BaseModel):
+        return payload.model_dump(mode="json", by_alias=True)
+    return to_jsonable(payload)
+
+
+class TaskAppClient:
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str | None = None,
+        *,
+        timeout: float = 600.0,
+        retries: int = 3,
+    ) -> None:
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.timeout = timeout
+        self.retries = max(1, retries)
+        self._client: httpx.AsyncClient | None = None
+        self.env = _TaskAppEnvironmentClient(self)
+
+    async def __aenter__(self) -> "TaskAppClient":
+        await self._ensure_client()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        await self.aclose()
+
+    async def _ensure_client(self) -> httpx.AsyncClient:
+        if self._client is None:
+            self._client = httpx.AsyncClient(
+                base_url=self.base_url,
+                timeout=httpx.Timeout(self.timeout),
+                follow_redirects=True,
+            )
+        return self._client
+
+    def _headers(self) -> Dict[str, str]:
+        headers: Dict[str, str] = {}
+        # Primary key
+        primary = (self.api_key or "").strip()
+        if primary:
+            headers["X-API-Key"] = primary
+            # Also set Authorization for clients that read bearer tokens
+            headers.setdefault("Authorization", f"Bearer {primary}")
+        # Include ALL available environment keys via CSV in X-API-Keys
+        keys: list[str] = []
+        if primary:
+            keys.append(primary)
+        aliases = (os.getenv("ENVIRONMENT_API_KEY_ALIASES") or "").strip()
+        if aliases:
+            for part in aliases.split(","):
+                trimmed = part.strip()
+                if trimmed and trimmed not in keys:
+                    keys.append(trimmed)
+        if keys:
+            headers["X-API-Keys"] = ",".join(keys)
+        return headers
+
+    async def aclose(self) -> None:
+        if self._client is not None:
+            await self._client.aclose()
+            self._client = None
+
+    async def _request(
+        self,
+        method: str,
+        path: str,
+        *,
+        params: Optional[Dict[str, Any] | List[tuple[str, Any]]] = None,
+        json_payload: Any = None,
+    ) -> httpx.Response:
+        client = await self._ensure_client()
+        payload = _prepare_payload(json_payload)
+        headers = self._headers()
+        last_exc: Exception | None = None
+        for attempt in range(self.retries):
+            try:
+                response = await client.request(
+                    method,
+                    path,
+                    headers=headers,
+                    params=params,
+                    json=payload,
+                )
+                response.raise_for_status()
+                return response
+            except httpx.HTTPStatusError as exc:
+                if 500 <= exc.response.status_code < 600 and attempt + 1 < self.retries:
+                    await asyncio.sleep(0.1 * (attempt + 1))
+                    last_exc = exc
+                    continue
+                raise
+            except httpx.HTTPError as exc:
+                last_exc = exc
+                if attempt + 1 >= self.retries:
+                    raise
+                await asyncio.sleep(0.1 * (attempt + 1))
+        if last_exc:  # pragma: no cover - defensive
+            raise last_exc
+        raise RuntimeError("Unreachable code in TaskAppClient._request")
+
+    async def health(self) -> Dict[str, Any]:
+        response = await self._request("GET", "/health")
+        return response.json()
+
+    async def info(self) -> Dict[str, Any]:
+        response = await self._request("GET", "/info")
+        return response.json()
+
+    async def task_info(self, seeds: list[int] | None = None) -> TaskInfo | list[TaskInfo]:
+        params: Optional[List[tuple[str, Any]]] = None
+        if seeds:
+            params = [("seed", seed) for seed in seeds]
+        response = await self._request("GET", "/task_info", params=params)
+        data = response.json()
+        if isinstance(data, list):
+            return [TaskInfo.model_validate(item) for item in data]
+        return TaskInfo.model_validate(data)
+
+    async def rollout(self, request: RolloutRequest) -> RolloutResponse:
+        response = await self._request("POST", "/rollout", json_payload=request)
+        data = response.json()
+        return RolloutResponse.model_validate(data)
+
+
+class _TaskAppEnvironmentClient:
+    def __init__(self, client: TaskAppClient) -> None:
+        self._client = client
+
+    async def initialize(self, env_name: str, payload: Dict[str, Any]) -> Dict[str, Any]:
+        response = await self._client._request(
+            "POST", f"/env/{env_name}/initialize", json_payload=payload
+        )
+        return response.json()
+
+    async def step(self, env_name: str, payload: Dict[str, Any]) -> Dict[str, Any]:
+        response = await self._client._request(
+            "POST", f"/env/{env_name}/step", json_payload=payload
+        )
+        return response.json()
+
+    async def terminate(self, env_name: str, payload: Dict[str, Any] | None = None) -> Dict[str, Any]:
+        response = await self._client._request(
+            "POST", f"/env/{env_name}/terminate", json_payload=payload or {}
+        )
+        return response.json()

synth_ai/task/contracts.py

@@ -1,31 +1,27 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Optional, Any, Dict, List
-from pydantic import BaseModel
+from typing import Optional, Any, Dict, List, Literal
+from pydantic import BaseModel, Field
 
 
 @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)
+    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.
     """
 
+    root: str = "/"
     health: str = "/health"
+    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"
@@ -67,6 +63,8 @@ class RolloutRecordConfig(BaseModel):
     trajectories: bool = True
     logprobs: bool = False
     value: bool = False
+    return_trace: bool = False
+    trace_format: Literal["compact", "full"] = "compact"
 
 
 class RolloutSafetyConfig(BaseModel):
@@ -108,6 +106,9 @@ class RolloutMetrics(BaseModel):
     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)
 
 
 class RolloutResponse(BaseModel):
@@ -117,4 +118,18 @@ class RolloutResponse(BaseModel):
     metrics: RolloutMetrics
     aborted: bool = False
     ops_executed: int = 0
+    trace: Dict[str, Any] | None = None
 
+
+class TaskInfo(BaseModel):
+    """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]

synth_ai/task/datasets.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+"""Dataset registry and helpers shared by Task Apps."""
+
+from typing import Any, Callable, Dict, Hashable, Tuple
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class TaskDatasetSpec(BaseModel):
+    """Declarative metadata describing a dataset that a Task App exposes."""
+
+    id: str
+    name: str
+    version: str | None = None
+    splits: list[str] = Field(default_factory=list)
+    default_split: str | None = None
+    cardinality: int | None = None
+    description: str | None = None
+
+    @field_validator("default_split")
+    @classmethod
+    def _validate_default_split(cls, value: str | None, info):
+        values = info.data if hasattr(info, "data") else {}  # type: ignore[attr-defined]
+        if value and value not in (values.get("splits") or []):
+            raise ValueError("default_split must be one of splits when provided")
+        return value
+
+
+RegistryLoader = Callable[[TaskDatasetSpec], Any]
+
+
+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] = {}
+
+    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)
+
+    def describe(self, dataset_id: str) -> TaskDatasetSpec:
+        if dataset_id not in self._entries:
+            raise KeyError(f"Dataset not registered: {dataset_id}")
+        spec, _, _ = self._entries[dataset_id]
+        return spec
+
+    def list(self) -> list[TaskDatasetSpec]:
+        return [entry[0] for entry in self._entries.values()]
+
+    def get(self, spec: TaskDatasetSpec | str) -> Any:
+        """Return dataset materialisation (with optional caching)."""
+
+        if isinstance(spec, str):
+            if spec not in self._entries:
+                raise KeyError(f"Dataset not registered: {spec}")
+            base_spec, loader, cache_enabled = self._entries[spec]
+            effective_spec = base_spec
+        else:
+            if spec.id not in self._entries:
+                raise KeyError(f"Dataset not registered: {spec.id}")
+            base_spec, loader, cache_enabled = self._entries[spec.id]
+            effective_spec = base_spec.model_copy(update=spec.model_dump(exclude_unset=True))
+
+        cache_key: Hashable = (
+            effective_spec.id,
+            effective_spec.version,
+            effective_spec.default_split,
+        )
+        if cache_enabled:
+            if cache_key not in self._cache:
+                self._cache[cache_key] = loader(effective_spec)
+            return self._cache[cache_key]
+        return loader(effective_spec)
+
+    @staticmethod
+    def ensure_split(spec: TaskDatasetSpec, split: str | None) -> str:
+        """Validate that `split` exists on the spec; return a concrete split."""
+
+        if not spec.splits:
+            return split or spec.default_split or "default"
+        if split is None:
+            if spec.default_split:
+                return spec.default_split
+            raise ValueError(f"split must be provided for dataset {spec.id}")
+        if split not in spec.splits:
+            raise ValueError(f"Unknown split '{split}' for dataset {spec.id}")
+        return split
+
+    @staticmethod
+    def normalise_seed(seed: Any, *, cardinality: int | None = None) -> int:
+        """Normalise arbitrary seed input into a bounded non-negative integer."""
+
+        try:
+            value = int(seed)
+        except Exception as exc:  # pragma: no cover - defensive
+            raise ValueError(f"Seed must be convertible to int (got {seed!r})") from exc
+        if value < 0:
+            value = abs(value)
+        if cardinality and cardinality > 0:
+            value = value % cardinality
+        return value

synth_ai/task/errors.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+"""Error helpers used across Task App implementations."""
+
+from typing import Any, Dict, Optional
+
+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}}
+    if extra:
+        payload["error"].update(extra)
+    return payload
+
+
+def http_exception(
+    status_code: int,
+    code: str,
+    message: str,
+    *,
+    extra: Optional[Dict[str, Any]] = None,
+    headers: Optional[Dict[str, str]] = None,
+):
+    try:
+        from fastapi import HTTPException  # type: ignore
+    except Exception as exc:  # pragma: no cover - FastAPI not installed
+        raise RuntimeError("fastapi must be installed to raise HTTPException") from exc
+
+    payload = error_payload(code, message, extra=extra)
+    return HTTPException(status_code=status_code, detail=to_jsonable(payload), headers=headers)
+
+
+def json_error_response(
+    status_code: int,
+    code: str,
+    message: str,
+    *,
+    extra: Optional[Dict[str, Any]] = None,
+    headers: Optional[Dict[str, str]] = None,
+):
+    try:
+        from fastapi.responses import JSONResponse  # type: ignore
+    except Exception as exc:  # pragma: no cover - FastAPI not installed
+        raise RuntimeError("fastapi must be installed to build JSONResponse") from exc
+
+    payload = error_payload(code, message, extra=extra)
+    return JSONResponse(status_code=status_code, content=to_jsonable(payload), headers=headers)
+

synth_ai/task/json.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+"""Shared JSON sanitisation helpers for Task Apps."""
+
+from collections.abc import Mapping, Sequence
+from dataclasses import is_dataclass, asdict
+from enum import Enum
+from typing import Any
+
+try:  # numpy is optional at runtime; degrade gracefully if absent
+    import numpy as _np  # type: ignore
+except Exception:  # pragma: no cover - handled at runtime
+    _np = None  # type: ignore
+
+
+def _mask_numpy_array(arr: "_np.ndarray") -> str:
+    shape = getattr(arr, "shape", None)
+    dtype = getattr(arr, "dtype", None)
+    return f"<ndarray shape={shape} dtype={dtype}>"
+
+
+def to_jsonable(value: Any) -> Any:
+    """Convert `value` into structures compatible with JSON serialisation.
+
+    - numpy scalars are converted to their Python counterparts
+    - numpy arrays are represented by a compact descriptor string
+    - dataclasses, Enums, and pydantic models are unwrapped recursively
+    - sets and tuples are converted to lists
+    - non-serialisable objects fall back to `repr`
+    """
+
+    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,)):
+            return int(value)
+        if isinstance(value, (_np.floating,)):
+            return float(value)
+        if isinstance(value, (_np.bool_,)):
+            return bool(value)
+        if isinstance(value, (_np.ndarray,)):
+            return _mask_numpy_array(value)
+
+    if isinstance(value, Enum):
+        return to_jsonable(value.value)
+
+    if is_dataclass(value):
+        return to_jsonable(asdict(value))
+
+    # pydantic BaseModel / attrs objects
+    for attr in ("model_dump", "dict", "to_dict", "to_json"):
+        if hasattr(value, attr) and callable(getattr(value, attr, None)):
+            try:
+                dumped = getattr(value, attr)()  # type: ignore[misc]
+            except TypeError:
+                dumped = getattr(value, attr)(exclude_none=False)  # pragma: no cover
+            return to_jsonable(dumped)
+
+    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]
+
+    if isinstance(value, (bytes, bytearray)):
+        return f"<bytes len={len(value)}>"
+
+    if hasattr(value, "__dict__"):
+        return to_jsonable(vars(value))
+
+    return repr(value)
+