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

synth_ai/task/server.py

@@ -0,0 +1,398 @@
+from __future__ import annotations
+
+"""FastAPI scaffolding for Task Apps (local dev + deployment)."""
+
+import asyncio
+import inspect
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence
+
+import httpx
+from fastapi import APIRouter, Depends, FastAPI, Query, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from starlette.middleware import Middleware
+
+from .auth import (
+    is_api_key_header_authorized,
+    normalize_environment_api_key,
+    require_api_key_dependency,
+)
+from .contracts import RolloutRequest, RolloutResponse, TaskInfo
+from .datasets import TaskDatasetRegistry
+from .errors import http_exception
+from .json import to_jsonable
+from .proxy import (
+    prepare_for_groq,
+    prepare_for_openai,
+    inject_system_hint,
+    synthesize_tool_call_if_missing,
+)
+from .rubrics import Rubric
+from .vendors import get_groq_key_or_503, get_openai_key_or_503, normalize_vendor_keys
+
+
+TasksetDescriptor = Callable[[], Mapping[str, Any] | Awaitable[Mapping[str, Any]]]
+InstanceProvider = Callable[[Sequence[int]], Iterable[TaskInfo] | Awaitable[Iterable[TaskInfo]]]
+RolloutExecutor = Callable[[RolloutRequest, Request], Any | Awaitable[Any]]
+
+
+@dataclass(slots=True)
+class RubricBundle:
+    """Optional rubrics advertised by the task app."""
+
+    outcome: Rubric | None = None
+    events: Rubric | None = None
+
+
+@dataclass(slots=True)
+class ProxyConfig:
+    """Configuration for optional vendor proxy endpoints."""
+
+    enable_openai: bool = False
+    enable_groq: bool = False
+    system_hint: str | None = None
+    openai_url: str = "https://api.openai.com/v1/chat/completions"
+    groq_url: str = "https://api.groq.com/openai/v1/chat/completions"
+
+
+@dataclass(slots=True)
+class TaskAppConfig:
+    """Declarative configuration describing a Task App."""
+
+    app_id: str
+    name: str
+    description: str
+    base_task_info: TaskInfo
+    describe_taskset: TasksetDescriptor
+    provide_task_instances: InstanceProvider
+    rollout: RolloutExecutor
+    dataset_registry: TaskDatasetRegistry | None = None
+    rubrics: RubricBundle = field(default_factory=RubricBundle)
+    proxy: ProxyConfig | None = None
+    routers: Sequence[APIRouter] = field(default_factory=tuple)
+    middleware: Sequence[Middleware] = field(default_factory=tuple)
+    app_state: Mapping[str, Any] = field(default_factory=dict)
+    require_api_key: bool = True
+    expose_debug_env: bool = True
+    cors_origins: Sequence[str] | None = None
+    startup_hooks: Sequence[Callable[[], None | Awaitable[None]]] = field(default_factory=tuple)
+    shutdown_hooks: Sequence[Callable[[], None | Awaitable[None]]] = field(default_factory=tuple)
+
+    def clone(self) -> "TaskAppConfig":
+        """Return a shallow copy safe to mutate when wiring the app."""
+
+        return TaskAppConfig(
+            app_id=self.app_id,
+            name=self.name,
+            description=self.description,
+            base_task_info=self.base_task_info,
+            describe_taskset=self.describe_taskset,
+            provide_task_instances=self.provide_task_instances,
+            rollout=self.rollout,
+            dataset_registry=self.dataset_registry,
+            rubrics=self.rubrics,
+            proxy=self.proxy,
+            routers=tuple(self.routers),
+            middleware=tuple(self.middleware),
+            app_state=dict(self.app_state),
+            require_api_key=self.require_api_key,
+            expose_debug_env=self.expose_debug_env,
+            cors_origins=tuple(self.cors_origins or ()),
+            startup_hooks=tuple(self.startup_hooks),
+            shutdown_hooks=tuple(self.shutdown_hooks),
+        )
+
+
+def _maybe_await(result: Any) -> Awaitable[Any]:
+    if inspect.isawaitable(result):
+        return asyncio.ensure_future(result)
+    loop = asyncio.get_event_loop()
+    future: asyncio.Future[Any] = loop.create_future()
+    future.set_result(result)
+    return future
+
+
+def _ensure_task_info(obj: Any) -> TaskInfo:
+    if isinstance(obj, TaskInfo):
+        return obj
+    if isinstance(obj, MutableMapping):
+        return TaskInfo.model_validate(obj)
+    raise TypeError(f"Task instance provider must yield TaskInfo-compatible objects (got {type(obj)!r})")
+
+
+def _normalise_seeds(values: Sequence[int]) -> list[int]:
+    seeds: list[int] = []
+    for value in values:
+        try:
+            seeds.append(int(value))
+        except Exception as exc:  # pragma: no cover - defensive
+            raise ValueError(f"Seed values must be convertible to int (got {value!r})") from exc
+    return seeds
+
+
+def _build_proxy_routes(
+    app: FastAPI, config: TaskAppConfig, auth_dependency: Callable[[Request], None]
+) -> None:
+    proxy = config.proxy
+    if not proxy:
+        return
+
+    async def _call_vendor(url: str, payload: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]:
+        async with httpx.AsyncClient(timeout=httpx.Timeout(600.0), follow_redirects=True) as client:
+            response = await client.post(url, json=payload, headers=headers)
+        data = (
+            response.json()
+            if response.headers.get("content-type", "").startswith("application/json")
+            else {"raw": response.text}
+        )
+        if response.status_code >= 400:
+            code = "vendor_error"
+            if url.startswith("https://api.openai.com"):
+                code = "openai_error"
+            elif "groq" in url:
+                code = "groq_error"
+            raise http_exception(
+                response.status_code,
+                code,
+                "Vendor proxy error",
+                extra={"status": response.status_code, "body": data},
+            )
+        return data
+
+    def _log_proxy(route: str, payload: dict[str, Any]) -> None:
+        try:
+            messages = payload.get("messages") if isinstance(payload, dict) else None
+            msg_count = len(messages) if isinstance(messages, list) else 0
+            tool_count = len(payload.get("tools") or []) if isinstance(payload, dict) else 0
+            model = payload.get("model") if isinstance(payload, dict) else None
+            print(f"[task:proxy:{route}] model={model} messages={msg_count} tools={tool_count}", flush=True)
+        except Exception:  # pragma: no cover - best effort logging
+            pass
+
+    system_hint = proxy.system_hint
+
+    if proxy.enable_openai:
+        @app.post("/proxy/v1/chat/completions", dependencies=[Depends(auth_dependency)])
+        async def proxy_openai(body: dict[str, Any], request: Request) -> Any:  # type: ignore[no-redef]
+            key = get_openai_key_or_503()
+            model = body.get("model") if isinstance(body.get("model"), str) else None
+            payload = prepare_for_openai(model, body)
+            payload = inject_system_hint(payload, system_hint or "")
+            _log_proxy("openai", payload)
+            data = await _call_vendor(proxy.openai_url, payload, {"Authorization": f"Bearer {key}"})
+            sanitized = synthesize_tool_call_if_missing(data)
+            return to_jsonable(sanitized)
+
+    if proxy.enable_groq:
+        @app.post("/proxy/groq/v1/chat/completions", dependencies=[Depends(auth_dependency)])
+        async def proxy_groq(body: dict[str, Any], request: Request) -> Any:  # type: ignore[no-redef]
+            key = get_groq_key_or_503()
+            model = body.get("model") if isinstance(body.get("model"), str) else None
+            payload = prepare_for_groq(model, body)
+            payload = inject_system_hint(payload, system_hint or "")
+            _log_proxy("groq", payload)
+            data = await _call_vendor(proxy.groq_url.rstrip("/"), payload, {"Authorization": f"Bearer {key}"})
+            sanitized = synthesize_tool_call_if_missing(data)
+            return to_jsonable(sanitized)
+
+
+def _auth_dependency_factory(config: TaskAppConfig) -> Callable[[Request], None]:
+    def _dependency(request: Request) -> None:
+        if not config.require_api_key:
+            return
+        require_api_key_dependency(request)
+
+    return _dependency
+
+
+def create_task_app(config: TaskAppConfig) -> FastAPI:
+    cfg = config.clone()
+    app = FastAPI(title=cfg.name, description=cfg.description)
+
+    for key, value in cfg.app_state.items():
+        setattr(app.state, key, value)
+
+    if cfg.cors_origins is not None:
+        app.add_middleware(
+            CORSMiddleware,
+            allow_origins=list(cfg.cors_origins) or ["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+    # Note: additional middleware from cfg.middleware is currently disabled to avoid typing ambiguity.
+    # for middleware in cfg.middleware:
+    #     try:
+    #         opts = getattr(middleware, "options", {})
+    #     except Exception:
+    #         opts = {}
+    #     app.add_middleware(middleware.cls, **(opts if isinstance(opts, dict) else {}))
+
+    for router in cfg.routers:
+        try:
+            app.include_router(router)
+        except Exception:
+            try:
+                inner = getattr(router, "router", None)
+                if inner is not None:
+                    app.include_router(inner)
+            except Exception:
+                raise
+
+    auth_dependency = _auth_dependency_factory(cfg)
+
+    def _call_hook(hook: Callable[..., Any]) -> Awaitable[Any]:
+        try:
+            params = inspect.signature(hook).parameters  # type: ignore[arg-type]
+        except (TypeError, ValueError):
+            params = {}
+        if params:
+            return _maybe_await(hook(app))  # type: ignore[misc]
+        return _maybe_await(hook())
+
+    @app.on_event("startup")
+    async def _startup() -> None:  # pragma: no cover - FastAPI lifecycle
+        normalize_environment_api_key()
+        normalize_vendor_keys()
+        for hook in cfg.startup_hooks:
+            await _call_hook(hook)
+
+    @app.on_event("shutdown")
+    async def _shutdown() -> None:  # pragma: no cover - FastAPI lifecycle
+        for hook in cfg.shutdown_hooks:
+            await _call_hook(hook)
+
+    @app.get("/")
+    async def root() -> Mapping[str, Any]:
+        return to_jsonable({"status": "ok", "service": cfg.app_id})
+
+    @app.head("/")
+    async def root_head() -> Mapping[str, Any]:
+        return to_jsonable({"status": "ok"})
+
+    @app.get("/health", dependencies=[Depends(auth_dependency)])
+    async def health(request: Request) -> Mapping[str, Any]:
+        # If we got here, auth_dependency already verified the key exactly matches
+        expected = normalize_environment_api_key()
+        return to_jsonable({"healthy": True, "auth": {"required": True, "expected_prefix": (expected[:6] + '...') if expected else '<unset>'}})
+
+    @app.get("/info", dependencies=[Depends(auth_dependency)])
+    async def info() -> Mapping[str, Any]:
+        dataset_meta = cfg.base_task_info.dataset
+        rubrics: dict[str, Any] | None = None
+        if cfg.rubrics.outcome or cfg.rubrics.events:
+            rubrics = {
+                "outcome": cfg.rubrics.outcome.model_dump() if cfg.rubrics.outcome else None,
+                "events": cfg.rubrics.events.model_dump() if cfg.rubrics.events else None,
+            }
+        payload = {
+            "service": {
+                "task": cfg.base_task_info.task,
+                "version": cfg.base_task_info.task.get("version"),
+            },
+            "dataset": dataset_meta,
+            "rubrics": rubrics,
+            "inference": cfg.base_task_info.inference,
+            "capabilities": cfg.base_task_info.capabilities,
+            "limits": cfg.base_task_info.limits,
+        }
+        return to_jsonable(payload)
+
+    @app.get("/task_info", dependencies=[Depends(auth_dependency)])
+    async def task_info(
+        seed: Sequence[int] | None = Query(default=None),
+        seeds: Sequence[int] | None = Query(default=None),
+    ) -> Any:
+        all_seeds: list[int] = []
+        if seed:
+            all_seeds.extend(_normalise_seeds(seed))
+        if seeds:
+            all_seeds.extend(_normalise_seeds(seeds))
+
+        if not all_seeds:
+            descriptor_result = await _maybe_await(cfg.describe_taskset())
+            return to_jsonable({"taskset": descriptor_result})
+
+        instances = await _maybe_await(cfg.provide_task_instances(all_seeds))
+        payload = [to_jsonable(_ensure_task_info(instance).model_dump()) for instance in instances]
+        return payload[0] if len(payload) == 1 else payload
+
+    @app.post("/rollout", dependencies=[Depends(auth_dependency)])
+    async def rollout_endpoint(rollout_request: RolloutRequest, request: Request) -> Any:
+        result = await _maybe_await(cfg.rollout(rollout_request, request))
+        if isinstance(result, RolloutResponse):
+            return to_jsonable(result.model_dump())
+        if isinstance(result, Mapping):
+            try:
+                validated = RolloutResponse.model_validate(result)
+            except Exception:
+                return to_jsonable(result)
+            return to_jsonable(validated.model_dump())
+        raise TypeError("Rollout executor must return RolloutResponse or mapping")
+
+    if cfg.expose_debug_env:
+        @app.get("/debug/env", dependencies=[Depends(auth_dependency)])
+        async def debug_env() -> Mapping[str, Any]:
+            def _mask(value: str | None) -> str:
+                if not value:
+                    return ""
+                return f"{value[:6]}…" if len(value) > 6 else value
+
+            return to_jsonable(
+                {
+                    "has_ENVIRONMENT_API_KEY": bool(os.getenv("ENVIRONMENT_API_KEY")),
+                    "OPENAI_API_KEY_prefix": _mask(os.getenv("OPENAI_API_KEY")),
+                    "GROQ_API_KEY_prefix": _mask(os.getenv("GROQ_API_KEY")),
+                }
+            )
+
+    _build_proxy_routes(app, cfg, auth_dependency)
+
+    return app
+
+
+def _load_env_files(env_files: Sequence[str]) -> list[str]:
+    loaded: list[str] = []
+    if not env_files:
+        return loaded
+    try:
+        import dotenv
+    except Exception:  # pragma: no cover - optional dep
+        return loaded
+    for path_str in env_files:
+        path = Path(path_str)
+        if not path.is_file():
+            continue
+        dotenv.load_dotenv(path, override=False)
+        loaded.append(str(path))
+    return loaded
+
+
+def run_task_app(
+    config_factory: Callable[[], TaskAppConfig],
+    *,
+    host: str = "0.0.0.0",
+    port: int = 8001,
+    reload: bool = False,
+    env_files: Sequence[str] = (),
+) -> None:
+    """Run the provided Task App configuration with uvicorn."""
+
+    loaded_files = _load_env_files(env_files)
+    if loaded_files:
+        print(f"[task:server] Loaded environment from: {', '.join(loaded_files)}", flush=True)
+
+    config = config_factory()
+    app = create_task_app(config)
+
+    try:
+        import uvicorn
+    except ImportError as exc:  # pragma: no cover - uvicorn optional
+        raise RuntimeError("uvicorn must be installed to run the task app locally") from exc
+
+    print(f"[task:server] Starting '{config.app_id}' on {host}:{port}", flush=True)
+    uvicorn.run(app, host=host, port=port, reload=reload)

synth_ai/task/tracing_utils.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+"""Utilities for wiring tracing_v3 into task apps."""
+
+import os
+import time
+from pathlib import Path
+from typing import Any, Callable
+
+
+def tracing_env_enabled(default: bool = False) -> bool:
+    """Return True when tracing is enabled for task apps via environment variable."""
+
+    raw = os.getenv("TASKAPP_TRACING_ENABLED")
+    if raw is None:
+        return default
+    raw = raw.strip().lower()
+    if raw in {"1", "true", "t", "yes", "y", "on"}:
+        return True
+    if raw in {"0", "false", "f", "no", "n", "off"}:
+        return False
+    return default
+
+
+def resolve_tracing_db_url() -> str | None:
+    """Resolve tracing database URL and prefer async drivers for SQLite."""
+
+    db_url = os.getenv("TURSO_LOCAL_DB_URL")
+    if db_url:
+        return db_url
+
+    sqld_path = os.getenv("SQLD_DB_PATH")
+    if sqld_path:
+        path = Path(sqld_path).expanduser()
+        if path.is_dir():
+            candidate = path / "dbs" / "default" / "data"
+            candidate.parent.mkdir(parents=True, exist_ok=True)
+            return f"sqlite+aiosqlite:///{candidate}"
+        else:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            return f"sqlite+aiosqlite:///{path}"
+
+    fallback_path = Path("traces/v3/synth_ai.db").expanduser()
+    fallback_path.parent.mkdir(parents=True, exist_ok=True)
+    return f"sqlite+aiosqlite:///{fallback_path}"
+
+
+def build_tracer_factory(make_tracer: Callable[..., Any], *, enabled: bool, db_url: str | None) -> Callable[[], Any] | None:
+    """Return a factory that instantiates a tracer when enabled, else None."""
+
+    if not enabled:
+        return None
+
+    def _factory() -> Any:
+        return make_tracer(db_url=db_url) if db_url else make_tracer()
+
+    return _factory
+
+
+def resolve_sft_output_dir() -> str | None:
+    """Resolve location for writing SFT records, creating directory if requested."""
+
+    raw = os.getenv("TASKAPP_SFT_OUTPUT_DIR") or os.getenv("SFT_OUTPUT_DIR")
+    if not raw:
+        return None
+    path = Path(raw).expanduser()
+    try:
+        path.mkdir(parents=True, exist_ok=True)
+    except Exception:
+        return None
+    return str(path)
+
+
+def unique_sft_path(base_dir: str, *, run_id: str) -> Path:
+    """Return a unique JSONL path for an SFT record batch."""
+
+    ts = int(time.time() * 1000)
+    name = f"{run_id}_{ts}.jsonl"
+    return Path(base_dir) / name

synth_ai/task/vendors.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+"""Vendor API key helpers shared by Task Apps."""
+
+import os
+from typing import Optional
+
+from .errors import http_exception
+
+_VENDOR_KEYS = {
+    "OPENAI_API_KEY": ("dev_openai_api_key", "DEV_OPENAI_API_KEY"),
+    "GROQ_API_KEY": ("dev_groq_api_key", "DEV_GROQ_API_KEY"),
+}
+
+
+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_single(key: str) -> Optional[str]:
+    direct = os.getenv(key)
+    if direct:
+        return direct
+    fallbacks = _VENDOR_KEYS.get(key, ())
+    for env in fallbacks:
+        candidate = os.getenv(env)
+        if candidate:
+            os.environ[key] = candidate
+            print(
+                f"[task:vendor] {key} set from {env} (prefix={_mask(candidate)})",
+                flush=True,
+            )
+            return candidate
+    return None
+
+
+def normalize_vendor_keys() -> dict[str, Optional[str]]:
+    """Normalise known vendor keys from dev fallbacks and return the mapping."""
+
+    resolved: dict[str, Optional[str]] = {}
+    for key in _VENDOR_KEYS:
+        resolved[key] = _normalize_single(key)
+    return resolved
+
+
+def get_openai_key_or_503() -> str:
+    key = _normalize_single("OPENAI_API_KEY")
+    if not key:
+        raise http_exception(503, "missing_openai_api_key", "OPENAI_API_KEY is not configured")
+    return key
+
+
+def get_groq_key_or_503() -> str:
+    key = _normalize_single("GROQ_API_KEY")
+    if not key:
+        raise http_exception(503, "missing_groq_api_key", "GROQ_API_KEY is not configured")
+    return key
+

synth_ai/tracing_v3/session_tracer.py

@@ -234,7 +234,7 @@ class SessionTracer:
             event_id = await self.db.insert_event_row(
                 self._current_trace.session_id,
                 timestep_db_id=timestep_db_id,
-                event=event,
+                event=event,  # type: ignore[arg-type]
             )
             # Auto-insert an event reward if EnvironmentEvent carries reward
             try:
@@ -323,7 +323,7 @@ class SessionTracer:
             return message_id
         return None
 
-    async def end_session(self, save: bool = None) -> SessionTrace:
+    async def end_session(self, save: bool | None = None) -> SessionTrace:
         """End the current session.
 
         Args:
@@ -370,7 +370,7 @@ class SessionTracer:
         self,
         session_id: str | None = None,
         metadata: dict[str, Any] | None = None,
-        save: bool = None,
+        save: bool | None = None,
     ):
         """Context manager for a session.
 
@@ -414,8 +414,16 @@ class SessionTracer:
         if limit:
             query += f" LIMIT {limit}"
 
-        df = await self.db.query_traces(query)
-        return df.to_dict("records")
+        # Ensure DB initialized before querying
+        if self.db is None:
+            await self.initialize()
+        df_or_records = await self.db.query_traces(query)  # type: ignore[union-attr]
+        try:
+            # If pandas DataFrame
+            return df_or_records.to_dict("records")  # type: ignore[call-arg, attr-defined]
+        except AttributeError:
+            # Already list of dicts
+            return df_or_records
 
     async def close(self):
         """Close database connections."""

synth_ai/tracing_v3/storage/base.py

@@ -4,8 +4,6 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import Any
 
-import pandas as pd
-
 from ..abstractions import SessionTrace
 
 
@@ -42,22 +40,22 @@ class TraceStorage(ABC):
         pass
 
     @abstractmethod
-    async def query_traces(self, query: str, params: dict[str, Any] = None) -> pd.DataFrame:
-        """Execute a query and return results as DataFrame.
+    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
+            Query results as a DataFrame-like object or list of dict records
         """
         pass
 
     @abstractmethod
     async def get_model_usage(
-        self, start_date: datetime = None, end_date: datetime = None, model_name: str = None
-    ) -> pd.DataFrame:
+        self, start_date: datetime | None = None, end_date: datetime | None = None, model_name: str | None = None
+    ) -> Any:
         """Get model usage statistics.
 
         Args:
@@ -66,7 +64,7 @@ class TraceStorage(ABC):
             model_name: Optional model name filter
 
         Returns:
-            Model usage statistics as a DataFrame
+            Model usage statistics as a DataFrame-like object or list of dict records
         """
         pass
 
@@ -92,8 +90,8 @@ class TraceStorage(ABC):
         self,
         experiment_id: str,
         name: str,
-        description: str = None,
-        configuration: dict[str, Any] = None,
+        description: str | None = None,
+        configuration: dict[str, Any] | None = None,
     ) -> str:
         """Create a new experiment."""
         raise NotImplementedError("Experiment management not supported by this backend")
@@ -103,14 +101,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
+        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 = 1000
+        self, traces: list[SessionTrace], batch_size: int | None = 1000
     ) -> list[str]:
         """Batch insert multiple session traces.