From 53578e5c534c7e71032a9bfa57b88ce73133c80a Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Mon, 27 Apr 2026 19:25:20 +0000 Subject: [PATCH 01/24] feat: enhance leaderboard with submission metadata including agent details and artifact SHA --- control_plane/owner_api/dashboard/queries.py | 54 ++++++++++++++++++++ control_plane/owner_api/dashboard/schemas.py | 9 ++++ 2 files changed, 63 insertions(+) diff --git a/control_plane/owner_api/dashboard/queries.py b/control_plane/owner_api/dashboard/queries.py index 80a4952..0102e0e 100644 --- a/control_plane/owner_api/dashboard/queries.py +++ b/control_plane/owner_api/dashboard/queries.py @@ -289,12 +289,29 @@ def _collect_single_run_rows( ) .group_by(TaskMinerResult.miner_hotkey) ).all() + # Resolve each live miner's currently-active deployment so the + # leaderboard can surface the agent name + artifact sha for the + # submission that is actually being benchmarked right now. + live_hotkeys = [r.miner_hotkey for r in rows if (r.task_count or 0) > 0] + submission_by_hk: dict[str, str] = {} + if live_hotkeys: + for dep in session.execute( + select(ManagedDeployment) + .where( + ManagedDeployment.family_id == family_id, + ManagedDeployment.miner_hotkey.in_(live_hotkeys), + ManagedDeployment.status != "retired", + ) + .order_by(ManagedDeployment.created_at.desc()) + ).scalars(): + submission_by_hk.setdefault(dep.miner_hotkey, dep.submission_id) return [ { "miner_hotkey": r.miner_hotkey, "raw_score": float(r.raw_score or 0.0), "normalized_score": None, "is_running": True, + "submission_id": submission_by_hk.get(r.miner_hotkey), } for r in rows if (r.task_count or 0) > 0 @@ -312,11 +329,39 @@ def _collect_single_run_rows( "raw_score": float(rec.raw_score), "normalized_score": float(rec.normalized_score), "is_running": False, + "submission_id": rec.submission_id, } for rec in records ] +def _resolve_submission_metadata( + session: Session, submission_ids: list[str] +) -> dict[str, dict[str, Any]]: + """Batch-load agent name/version + archive sha256 by submission id. + + Surfaced on the leaderboard so miners can spot-check that the artifact + the subnet is running for them matches their local sha (reproducibility + + tamper detection). + """ + out: dict[str, dict[str, Any]] = {} + if not submission_ids: + return out + for sub in session.execute( + select(ManagedMinerSubmission).where( + ManagedMinerSubmission.id.in_(submission_ids) + ) + ).scalars(): + manifest = sub.manifest_json or {} + agent = manifest.get("agent") or {} + out[sub.id] = { + "agent_name": agent.get("name") if isinstance(agent, dict) else None, + "agent_version": agent.get("version") if isinstance(agent, dict) else None, + "artifact_sha256": sub.archive_sha256, + } + return out + + def _rank_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: sorted_rows = sorted(rows, key=lambda r: (-(r["raw_score"] or 0.0), r["miner_hotkey"])) for idx, row in enumerate(sorted_rows, start=1): @@ -432,15 +477,24 @@ def fetch_leaderboard( ).all() ) + submission_ids = [ + row["submission_id"] for row in ranked if row.get("submission_id") + ] + submission_meta = _resolve_submission_metadata(session, submission_ids) + entries: list[LeaderboardEntry] = [] for row in ranked: hk = row["miner_hotkey"] prev = prev_ranks.get(hk) + meta = submission_meta.get(row.get("submission_id") or "", {}) entries.append( LeaderboardEntry( rank=row["rank"], hotkey=hk, hotkey_short=shorten_hotkey(hk), + agent_name=meta.get("agent_name"), + agent_version=meta.get("agent_version"), + artifact_sha256=meta.get("artifact_sha256"), raw_score=row["raw_score"], normalized_score=row.get("normalized_score"), is_serving_winner=(hk == winner_hk), diff --git a/control_plane/owner_api/dashboard/schemas.py b/control_plane/owner_api/dashboard/schemas.py index 479e0f4..7c52243 100644 --- a/control_plane/owner_api/dashboard/schemas.py +++ b/control_plane/owner_api/dashboard/schemas.py @@ -75,6 +75,15 @@ class LeaderboardEntry(BaseModel): rank: int hotkey: str hotkey_short: str + # Agent identity surfaced from the submission's manifest.yaml. Lets + # miners search by their own project name instead of memorising a + # hotkey, and lets observers tell agents apart at a glance. + agent_name: str | None = None + agent_version: str | None = None + # SHA-256 of the submission tarball as the subnet stored it. Miners + # can `sha256sum` their local archive and compare to confirm the + # subnet is running an unmodified copy of what they uploaded. + artifact_sha256: str | None = None raw_score: float normalized_score: float | None is_serving_winner: bool From b26b37fe1084707380833d963dfd6fe8c2a436ef Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:48:23 +0000 Subject: [PATCH 02/24] feat(orchestrator): chat-stream passthrough; consumer-api becomes thin SSE facade --- orchestration/consumer_api/chat.py | 174 +++--------- orchestration/consumer_api/main.py | 6 + orchestration/orchestrator/chat_stream.py | 316 ++++++++++++++++++++++ orchestration/orchestrator/main.py | 59 +++- shared/common/migrations.py | 53 ++++ shared/common/models.py | 7 + tests/services/test_consumer_api.py | 110 ++++---- tests/services/test_orchestrator.py | 144 ++++++++++ 8 files changed, 687 insertions(+), 182 deletions(-) create mode 100644 orchestration/orchestrator/chat_stream.py diff --git a/orchestration/consumer_api/chat.py b/orchestration/consumer_api/chat.py index e52233c..5f9b905 100644 --- a/orchestration/consumer_api/chat.py +++ b/orchestration/consumer_api/chat.py @@ -5,7 +5,6 @@ import logging import os import time -import uuid from collections.abc import AsyncIterator from typing import Any @@ -18,14 +17,9 @@ _logger = logging.getLogger(__name__) _orchestrator_circuit_breaker = CircuitBreaker(failure_threshold=5, recovery_timeout=30.0) -OWNER_API_URL = os.getenv("OWNER_API_URL", "http://owner-api:8000") -INTERNAL_SERVICE_TOKEN = os.getenv("EIREL_INTERNAL_SERVICE_TOKEN", "") -# Default family for the consumer chat surface. The streaming endpoint -# resolves a serving miner for this family from owner-api and proxies its -# /v1/agent/infer/stream NDJSON back to the client as SSE. -_CHAT_FAMILY_ID = os.getenv("EIREL_CONSUMER_CHAT_FAMILY", "general_chat") -# Total wall-clock budget for a streaming chat. Should clear the slowest -# acceptable miner completion (thinking mode = 600s today). +# Total wall-clock budget for a streaming chat. The orchestrator's +# upstream budget is the same; we set an envelope here so the consumer +# connection times out around the same time the orchestrator gives up. _CHAT_STREAM_TIMEOUT_SECONDS = float( os.getenv("EIREL_CONSUMER_CHAT_STREAM_TIMEOUT_SECONDS", "660") ) @@ -127,46 +121,12 @@ async def _orchestrator_post() -> httpx.Response: return response.status_code, response.json() -async def _resolve_serving_miner(family_id: str) -> dict[str, Any] | None: - """Fetch the current winner miner for a family from owner-api. - - Resolution order: - 1. `EIREL_CONSUMER_CHAT_MINER_OVERRIDE_ENDPOINT` env (test/debug - override — bypasses owner-api entirely). - 2. `/v1/internal/serving/{family_id}` (production: a published - serving release). - 3. `/v1/internal/managed-deployments/active/{family_id}` (fallback: - any healthy active managed deployment, useful before the first - serving release of a run is published). - """ - override = os.getenv("EIREL_CONSUMER_CHAT_MINER_OVERRIDE_ENDPOINT", "").strip() - if override: - return {"endpoint": override, "hotkey": "override", "family_id": family_id} - - headers: dict[str, str] = {} - if INTERNAL_SERVICE_TOKEN: - headers["Authorization"] = f"Bearer {INTERNAL_SERVICE_TOKEN}" - - async with httpx.AsyncClient(timeout=10.0) as client: - for path in ( - f"/v1/internal/serving/{family_id}", - f"/v1/internal/managed-deployments/active/{family_id}", - ): - try: - resp = await client.get(f"{OWNER_API_URL}{path}", headers=headers) - if resp.status_code == 404: - continue - resp.raise_for_status() - data = resp.json() - if data.get("endpoint"): - return data - except Exception as exc: # noqa: BLE001 - _logger.error( - "failed to resolve miner via %s for %s: %s", - path, family_id, exc, - ) - continue - return None +# consumer-chat-api is now a thin SSE facade. Resolving the +# serving deployment, building the slim invocation body, and proxying +# NDJSON from the family pod all live in the orchestrator now. This +# module only translates between consumer-chat-api's SSE surface and +# the orchestrator's NDJSON streaming endpoint. +ORCHESTRATOR_URL = os.getenv("ORCHESTRATOR_URL", "http://orchestrator:8050") def _sse_event(event: str, data: dict[str, Any]) -> bytes: @@ -188,100 +148,56 @@ async def stream_chat_request( user_id: str = "anonymous", session_id: str | None = None, context_history: list[dict[str, Any]] | None = None, + mode: str = "instant", + web_search: bool = False, ) -> AsyncIterator[bytes]: - """Stream a chat response as SSE. - - 1. Resolve the current serving miner for the chat family from owner-api. - 2. POST to that miner's `/v1/agent/infer/stream` (NDJSON, eirel SDK ≥ 0.2.3). - 3. Re-emit each NDJSON chunk as an SSE event with the chunk's `event` - name (`delta`/`citation`/`tool_call`/`done`). Falls back to the - unary endpoint on 404 and emits the whole answer as a single - `delta` followed by `done` so the client UX is identical. + """Stream a chat response as SSE — thin facade in front of the orchestrator. - Errors emit a final `error` SSE event then close the stream. - """ - task_id = session_id or f"chat-{uuid.uuid4().hex[:12]}" - yield _sse_event("started", {"task_id": task_id, "family_id": _CHAT_FAMILY_ID}) + Forwards to the orchestrator's ``/v1/orchestrate/chat/stream`` and + re-emits each NDJSON ``StreamChunk`` as an SSE event with the + chunk's ``event`` name (``started`` / ``delta`` / ``citation`` / + ``tool_call`` / ``done``). The orchestrator owns: - miner = await _resolve_serving_miner(_CHAT_FAMILY_ID) - if miner is None: - yield _sse_event("error", { - "message": f"no serving miner available for family {_CHAT_FAMILY_ID}", - }) - return + * session state (mode + web_search persistence), + * routing to the right family (today: passthrough to + ``general_chat``; later: DAG composition across families), + * miner resolution + NDJSON proxy. + A network failure to the orchestrator emits a terminal ``error`` + SSE event then closes the stream. + """ body = { - "task_id": task_id, - "family_id": _CHAT_FAMILY_ID, - "primary_goal": prompt, - "subtask": prompt, - "inputs": {}, + "prompt": prompt, + "user_id": user_id, + "session_id": session_id, "context_history": context_history or [], - "metadata": {"user_id": user_id, "session_id": session_id}, + "mode": mode, + "web_search": bool(web_search), } - headers = {"Content-Type": "application/json"} - endpoint = miner["endpoint"].rstrip("/") - stream_url = f"{endpoint}/v1/agent/infer/stream" - unary_url = f"{endpoint}/v1/agent/infer" start = time.monotonic() - used_stream = True final_status = "completed" + forward_url = f"{ORCHESTRATOR_URL.rstrip('/')}/v1/orchestrate/chat/stream" try: async with httpx.AsyncClient(timeout=_CHAT_STREAM_TIMEOUT_SECONDS) as client: - try: - async with client.stream( - "POST", stream_url, json=body, headers=headers, - ) as resp: - if resp.status_code == 404: - used_stream = False - else: - resp.raise_for_status() - async for line in resp.aiter_lines(): - line = line.strip() - if not line: - continue - try: - chunk = json.loads(line) - except json.JSONDecodeError: - _logger.warning( - "malformed NDJSON from miner: %r", line[:120], - ) - continue - event = chunk.get("event") or "delta" - if event == "done": - final_status = chunk.get("status") or "completed" - yield _sse_event(event, chunk) - except httpx.HTTPStatusError as exc: - if exc.response.status_code != 404: - raise - used_stream = False - - if not used_stream: - # Older miners on eirel SDK < 0.2.3 — assemble a single - # delta + done from the unary response so the client sees - # the same SSE shape. - resp = await client.post(unary_url, json=body, headers=headers) + async with client.stream("POST", forward_url, json=body) as resp: resp.raise_for_status() - payload = resp.json() if resp.content else {} - output = payload.get("output") or {} - text = "" - for key in ("answer", "response", "text", "content", "message"): - value = output.get(key) - if isinstance(value, str) and value: - text = value - break - if text: - yield _sse_event("delta", {"event": "delta", "text": text}) - yield _sse_event("done", { - "event": "done", - "output": output, - "citations": payload.get("citations") or [], - "tool_calls": payload.get("tool_calls") or [], - "status": payload.get("status") or "completed", - }) - final_status = payload.get("status") or "completed" + async for line in resp.aiter_lines(): + line = line.strip() + if not line: + continue + try: + chunk = json.loads(line) + except json.JSONDecodeError: + _logger.warning( + "malformed NDJSON from orchestrator: %r", line[:120], + ) + continue + event = chunk.get("event") or "delta" + if event == "done": + final_status = chunk.get("status") or "completed" + yield _sse_event(event, chunk) except Exception as exc: # noqa: BLE001 _logger.exception("stream_chat_request failed: %s", exc) yield _sse_event("error", {"message": str(exc)}) diff --git a/orchestration/consumer_api/main.py b/orchestration/consumer_api/main.py index 71e331a..09d9617 100644 --- a/orchestration/consumer_api/main.py +++ b/orchestration/consumer_api/main.py @@ -31,6 +31,10 @@ class ChatRequest(BaseModel): user_id: str = Field(default="anonymous", max_length=64) session_id: str | None = Field(default=None, max_length=128) context_history: list[dict] = Field(default_factory=list, max_length=50) + # Per-session toggles. Persisted on the session row in the + # orchestrator (consumer-chat-api forwards them as-is). + mode: str = Field(default="instant", pattern="^(instant|thinking)$") + web_search: bool = Field(default=False) @asynccontextmanager @@ -134,6 +138,8 @@ async def chat_stream( user_id=payload.user_id, session_id=payload.session_id, context_history=payload.context_history, + mode=payload.mode, + web_search=payload.web_search, ), media_type="text/event-stream", headers={ diff --git a/orchestration/orchestrator/chat_stream.py b/orchestration/orchestrator/chat_stream.py new file mode 100644 index 0000000..3fce9b9 --- /dev/null +++ b/orchestration/orchestrator/chat_stream.py @@ -0,0 +1,316 @@ +"""Streaming chat passthrough for the orchestrator. + +The orchestrator owns conversation state (per the multi-family +architecture); for now we only route to ``general_chat``, but the +session bookkeeping and the wire shape are already family-agnostic so +adding a DAG composition path later is local to this module. + +Two responsibilities: + + * **Session state** — load/upsert the per-session ``mode`` and + ``web_search`` toggles on ``ConsumerSessionState``. The toggles + persist across turns so the consumer-chat-api (and the user's + browser tab) don't have to re-assert them on every prompt. + + * **Miner proxy** — resolve the current serving deployment for the + target family from owner-api, build the slim 0.3.0 invocation body, + and stream the miner's NDJSON back to the caller line-by-line. + Falls back to the unary endpoint with a synthetic ``delta`` + + ``done`` if the miner pod doesn't expose ``/v1/agent/infer/stream`` + (eirel SDK < 0.2.3). +""" + +from __future__ import annotations + +import json +import logging +import os +import time +import uuid +from collections.abc import AsyncIterator +from typing import Any + +import httpx + +from shared.common.database import Database +from shared.common.models import ConsumerSessionState, utcnow + +_logger = logging.getLogger(__name__) + +OWNER_API_URL = os.getenv("OWNER_API_URL", "http://owner-api:8000") +INTERNAL_SERVICE_TOKEN = os.getenv("EIREL_INTERNAL_SERVICE_TOKEN", "") +# Total wall-clock budget for a streaming chat. Should clear the slowest +# acceptable miner completion (thinking mode = 600s today). +_CHAT_STREAM_TIMEOUT_SECONDS = float( + os.getenv("EIREL_ORCHESTRATOR_CHAT_STREAM_TIMEOUT_SECONDS", "660") +) + + +class ChatSessionStore: + """Thin facade over ``ConsumerSessionState`` for the chat path. + + All writes go through this so the orchestrator owns the single + write-path to the session row. Consumer-chat-api will eventually be + a stateless facade — no DB writes from there. + """ + + def __init__(self, db: Database) -> None: + self._db = db + + def resolve_toggles( + self, + *, + session_id: str | None, + user_id: str, + mode_override: str | None, + web_search_override: bool | None, + ) -> tuple[str, str, bool]: + """Load (and upsert) mode/web_search for the session. + + Resolution rules: + * If ``session_id`` is None, generate a new one and persist + with the override values (or defaults). + * If the session exists, override values win when present; + otherwise reuse the stored row values. + * If the session is missing, create it with override-or-default. + + Returns ``(session_id, mode, web_search)``. + """ + sid = session_id or str(uuid.uuid4()) + with self._db.sessionmaker() as s: + row = s.get(ConsumerSessionState, sid) + if row is None: + row = ConsumerSessionState( + session_id=sid, + user_id=user_id, + status="active", + messages_json=[], + mode=mode_override or "instant", + web_search=bool(web_search_override) if web_search_override is not None else False, + ) + s.add(row) + else: + if mode_override is not None: + row.mode = mode_override + if web_search_override is not None: + row.web_search = bool(web_search_override) + row.updated_at = utcnow() + mode = row.mode + web_search = bool(row.web_search) + s.commit() + return sid, mode, web_search + + +async def _resolve_serving_endpoint(family_id: str) -> dict[str, Any] | None: + """Look up the currently serving deployment for *family_id* in owner-api. + + Resolution order: + 1. ``EIREL_ORCHESTRATOR_MINER_OVERRIDE_ENDPOINT`` env override — + test-only, bypasses owner-api entirely. + 2. ``/v1/internal/serving/{family_id}`` — production winner pod + that won the most recent epoch for this family. + 3. ``/v1/internal/managed-deployments/active/{family_id}`` — + fallback when no serving release has been published yet + (e.g. the first run after a fresh subnet bring-up). Useful so + the chat surface still works in early-cluster states. + """ + override = os.getenv("EIREL_ORCHESTRATOR_MINER_OVERRIDE_ENDPOINT", "").strip() + if override: + return {"endpoint": override, "hotkey": "override", "family_id": family_id} + + headers: dict[str, str] = {} + if INTERNAL_SERVICE_TOKEN: + headers["Authorization"] = f"Bearer {INTERNAL_SERVICE_TOKEN}" + + paths = [ + f"/v1/internal/serving/{family_id}", + f"/v1/internal/managed-deployments/active/{family_id}", + ] + async with httpx.AsyncClient(timeout=10.0) as client: + for path in paths: + try: + resp = await client.get(f"{OWNER_API_URL}{path}", headers=headers) + if resp.status_code == 404: + continue + resp.raise_for_status() + data = resp.json() + if isinstance(data, dict) and data.get("endpoint"): + return data + except Exception as exc: + _logger.warning( + "orchestrator: serving lookup %s for %s failed: %s", + path, family_id, exc, + ) + continue + return None + + +def _ndjson(chunk: dict[str, Any]) -> bytes: + return (json.dumps(chunk, separators=(",", ":")) + "\n").encode("utf-8") + + +def _build_invocation_body( + *, + prompt: str, + mode: str, + web_search: bool, + history: list[dict[str, Any]], + turn_id: str, + family_id: str, +) -> dict[str, Any]: + """Construct the slim 0.3.0 family-agent body, with legacy fallbacks. + + Mirrors validator's ``_build_body`` so the orchestrator and the + eval pipeline send identical wire shapes — the family agent only + learns one contract. + """ + body: dict[str, Any] = { + # Slim 0.3.0 contract. + "turn_id": turn_id, + "prompt": prompt, + "mode": mode, + "web_search": web_search, + "history": history, + } + # Legacy mirror suppressed when ``EIREL_VALIDATOR_SLIM_ONLY=1`` — + # same kill switch the validator uses, so a slim-only test exercises + # both paths consistently. Drops in 0.4.0 either way. + if os.getenv("EIREL_VALIDATOR_SLIM_ONLY", "0") not in {"1", "true", "yes"}: + body.update({ + "task_id": turn_id, + "family_id": family_id, + "primary_goal": prompt, + "subtask": prompt, + "inputs": {"mode": mode, "web_search": web_search}, + "context_history": history, + }) + return body + + +async def stream_family_chat( + *, + store: ChatSessionStore, + family_id: str, + prompt: str, + user_id: str, + session_id: str | None, + context_history: list[dict[str, Any]], + mode_override: str | None, + web_search_override: bool | None, +) -> AsyncIterator[bytes]: + """Yield NDJSON StreamChunks for one chat turn. + + Sequence: + 1. Resolve session toggles (and create the session row if new). + 2. Resolve serving family deployment endpoint via owner-api. + 3. Stream the miner's ``/v1/agent/infer/stream``, line-by-line. + 4. On 404 streaming, fall back to the unary endpoint and emit a + synthetic single ``delta`` + ``done`` so the wire contract is + identical regardless of miner SDK version. + 5. Always end with a ``done`` chunk — error or success. + """ + sid, mode, web_search = store.resolve_toggles( + session_id=session_id, + user_id=user_id, + mode_override=mode_override, + web_search_override=web_search_override, + ) + turn_id = f"chat-{uuid.uuid4().hex[:12]}" + + # The very first chunk is informational so the consumer-chat-api can + # echo session_id + family back to the browser before content arrives. + yield _ndjson({ + "event": "started", + "metadata": { + "session_id": sid, + "family_id": family_id, + "mode": mode, + "web_search": web_search, + "turn_id": turn_id, + }, + }) + + miner = await _resolve_serving_endpoint(family_id) + if miner is None: + yield _ndjson({ + "event": "done", + "status": "failed", + "error": f"no serving deployment available for family {family_id}", + }) + return + + history = [ + {"role": h.get("role"), "content": h.get("content")} + for h in (context_history or []) + if isinstance(h, dict) and h.get("role") in ("user", "assistant") + ] + body = _build_invocation_body( + prompt=prompt, + mode=mode, + web_search=web_search, + history=history, + turn_id=turn_id, + family_id=family_id, + ) + + endpoint = miner["endpoint"].rstrip("/") + stream_url = f"{endpoint}/v1/agent/infer/stream" + unary_url = f"{endpoint}/v1/agent/infer" + + started_at = time.monotonic() + used_stream = True + try: + async with httpx.AsyncClient(timeout=_CHAT_STREAM_TIMEOUT_SECONDS) as client: + async with client.stream("POST", stream_url, json=body) as resp: + if resp.status_code == 404: + used_stream = False + else: + resp.raise_for_status() + async for line in resp.aiter_lines(): + line = line.strip() + if not line: + continue + yield (line + "\n").encode("utf-8") + return + except httpx.HTTPError as exc: + _logger.warning( + "orchestrator: stream call to %s failed: %s", + stream_url, exc, + ) + used_stream = False + + if not used_stream: + try: + async with httpx.AsyncClient(timeout=_CHAT_STREAM_TIMEOUT_SECONDS) as client: + resp = await client.post(unary_url, json=body) + resp.raise_for_status() + payload = resp.json() if resp.content else {} + except Exception as exc: + yield _ndjson({"event": "done", "status": "failed", "error": str(exc)}) + return + + text = "" + if isinstance(payload, dict): + out = payload.get("output") or {} + if isinstance(out, dict): + for key in ("answer", "response", "text", "content", "message"): + val = out.get(key) + if isinstance(val, str) and val: + text = val + break + if text: + yield _ndjson({"event": "delta", "text": text}) + yield _ndjson({ + "event": "done", + "status": payload.get("status") if isinstance(payload, dict) else "completed", + "output": payload.get("output") if isinstance(payload, dict) else {}, + "citations": payload.get("citations") if isinstance(payload, dict) else [], + "metadata": { + **(payload.get("metadata", {}) if isinstance(payload, dict) else {}), + "fallback": "unary", + "elapsed_seconds": round(time.monotonic() - started_at, 3), + }, + }) + + +__all__ = ["ChatSessionStore", "stream_family_chat"] diff --git a/orchestration/orchestrator/main.py b/orchestration/orchestrator/main.py index f21938b..bf0e623 100644 --- a/orchestration/orchestrator/main.py +++ b/orchestration/orchestrator/main.py @@ -15,12 +15,18 @@ import uvicorn from fastapi import FastAPI, Request -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse from starlette.responses import Response from pydantic import BaseModel, Field +from shared.common.config import get_settings +from shared.common.database import Database from shared.common.request_context import RequestIdMiddleware from shared.common.tracing import init_tracing, get_tracer +from orchestration.orchestrator.chat_stream import ( + ChatSessionStore, + stream_family_chat, +) from orchestration.orchestrator.orchestrator import Orchestrator init_tracing("orchestrator") @@ -41,6 +47,10 @@ async def lifespan(app: FastAPI): logger = logging.getLogger(__name__) logger.info("orchestrator starting up") app.state.orchestrator = Orchestrator() + settings = get_settings() + db = Database(settings.database_url) + db.create_all() + app.state.session_store = ChatSessionStore(db) yield logger.info("orchestrator shutting down") @@ -80,6 +90,53 @@ async def list_tools(request: Request) -> dict[str, Any]: } +class ChatStreamRequest(BaseModel): + """Streaming chat request from consumer-chat-api. + + Mirrors the slim family-agent contract — the orchestrator routes + this through (today: 1:1 to ``general_chat``) and proxies the + miner's NDJSON back to the caller. Per-session toggles + (``mode`` / ``web_search``) are persisted on the session row by + the orchestrator so the caller doesn't need to re-assert them on + every turn. + """ + + prompt: str + user_id: str = "anonymous" + session_id: str | None = None + context_history: list[dict[str, Any]] = Field(default_factory=list) + # Per-turn overrides — when present, persist on the session row. + # When absent, the orchestrator uses whatever's stored on the row. + mode: str | None = Field(default=None, pattern="^(instant|thinking)$") + web_search: bool | None = None + + +@app.post("/v1/orchestrate/chat/stream") +async def chat_stream(payload: ChatStreamRequest, request: Request): + """Streaming chat — single-family passthrough today, DAG-composed later. + + Today the orchestrator only routes to ``general_chat``. Once more + families come online, this entrypoint will fan out / synthesize + based on ``select_route()``. The wire shape stays stable: NDJSON + StreamChunks back to the caller (consumer-chat-api re-emits as SSE). + """ + session_store: ChatSessionStore = request.app.state.session_store + return StreamingResponse( + stream_family_chat( + store=session_store, + family_id="general_chat", + prompt=payload.prompt, + user_id=payload.user_id, + session_id=payload.session_id, + context_history=payload.context_history, + mode_override=payload.mode, + web_search_override=payload.web_search, + ), + media_type="application/x-ndjson", + headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"}, + ) + + @app.post("/v1/orchestrate") async def orchestrate(payload: OrchestratorRequest, request: Request): """Main orchestration endpoint. diff --git a/shared/common/migrations.py b/shared/common/migrations.py index 5277cc7..ffb2a26 100644 --- a/shared/common/migrations.py +++ b/shared/common/migrations.py @@ -374,6 +374,23 @@ def _migration_add_snapshot_unique_constraint(engine: Engine) -> None: ), apply=lambda engine: _migration_drop_miner_first_token_seconds(engine), ), + Migration( + version="add_session_mode_web_search", + description=( + "Persist per-session user toggles (mode, web_search) on " + "consumer_sessions so the orchestrator can apply them on every " + "turn without the client re-asserting them. Owned by the " + "orchestrator now that consumer-chat-api is a thin facade." + ), + apply=lambda engine: _migration_add_session_mode_web_search(engine), + ), + Migration( + version="add_proxy_and_judge_cost_to_task_miner_results", + description=( + "Add proxy_cost_usd + judge_cost_usd to task_miner_results" + ), + apply=lambda engine: _migration_add_proxy_and_judge_cost(engine), + ), ) @@ -418,6 +435,42 @@ def _migration_drop_miner_first_token_seconds(engine: Engine) -> None: )) +def _migration_add_session_mode_web_search(engine: Engine) -> None: + inspector = inspect(engine) + if "consumer_sessions" not in set(inspector.get_table_names()): + return + cols = {c["name"] for c in inspector.get_columns("consumer_sessions")} + with engine.begin() as conn: + if "mode" not in cols: + conn.execute(text( + "ALTER TABLE consumer_sessions " + "ADD COLUMN mode VARCHAR(16) NOT NULL DEFAULT 'instant'" + )) + if "web_search" not in cols: + conn.execute(text( + "ALTER TABLE consumer_sessions " + "ADD COLUMN web_search BOOLEAN NOT NULL DEFAULT FALSE" + )) + + +def _migration_add_proxy_and_judge_cost(engine: Engine) -> None: + inspector = inspect(engine) + if "task_miner_results" not in set(inspector.get_table_names()): + return + cols = {c["name"] for c in inspector.get_columns("task_miner_results")} + with engine.begin() as conn: + if "proxy_cost_usd" not in cols: + conn.execute(text( + "ALTER TABLE task_miner_results " + "ADD COLUMN proxy_cost_usd FLOAT NOT NULL DEFAULT 0.0" + )) + if "judge_cost_usd" not in cols: + conn.execute(text( + "ALTER TABLE task_miner_results " + "ADD COLUMN judge_cost_usd FLOAT NOT NULL DEFAULT 0.0" + )) + + def _drop_registered_neurons_is_active(engine: Engine) -> None: inspector = inspect(engine) if "registered_neurons" not in set(inspector.get_table_names()): diff --git a/shared/common/models.py b/shared/common/models.py index abc8b33..937c7ac 100644 --- a/shared/common/models.py +++ b/shared/common/models.py @@ -480,6 +480,13 @@ class ConsumerSessionState(Base): status: Mapped[str] = mapped_column(String(32), nullable=False, default="active") latest_task_id: Mapped[str | None] = mapped_column(String(36), nullable=True, index=True) messages_json: Mapped[list[dict[str, Any]]] = mapped_column(JSON, nullable=False, default=list) + # Per-session user toggles, owned by the orchestrator and applied to + # every turn dispatched to a family agent. Persisted so the + # frontend can render the current state on reconnect, and so the + # toggles survive page reloads. Defaults match the family-side + # defaults (instant + no search). + mode: Mapped[str] = mapped_column(String(16), nullable=False, default="instant") + web_search: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=False), default=utcnow, nullable=False ) diff --git a/tests/services/test_consumer_api.py b/tests/services/test_consumer_api.py index c3dd517..d5ea92b 100644 --- a/tests/services/test_consumer_api.py +++ b/tests/services/test_consumer_api.py @@ -182,24 +182,24 @@ def _parse_sse(body: bytes) -> list[tuple[str, dict]]: async def test_chat_stream_proxies_ndjson_as_sse(consumer_env, monkeypatch): - """Happy path: serving miner returns NDJSON stream → consumer-api re-emits as SSE.""" + """consumer-chat-api forwards to the orchestrator and + re-emits NDJSON StreamChunks as SSE events. The orchestrator now owns + miner resolution and the family-agent body shape; this test only + exercises the SSE translation layer. + """ import json as _json - async def _fake_resolve(family_id): - return {"endpoint": "http://miner.local", "hotkey": "hk"} - - monkeypatch.setattr( - "orchestration.consumer_api.chat._resolve_serving_miner", _fake_resolve, - ) - ndjson_body = ( - _json.dumps({"event": "delta", "text": "Hello "}) + "\n" + _json.dumps({ + "event": "started", + "metadata": {"session_id": "s-1", "family_id": "general_chat"}, + }) + "\n" + + _json.dumps({"event": "delta", "text": "Hello "}) + "\n" + _json.dumps({"event": "delta", "text": "world"}) + "\n" + _json.dumps({ "event": "done", "output": {"answer": "Hello world"}, "citations": [], - "tool_calls": [], "status": "completed", }) + "\n" ).encode("utf-8") @@ -245,38 +245,22 @@ def _patched(*args, **kwargs): deltas = [e[1].get("text", "") for e in events if e[0] == "delta"] assert "".join(deltas) == "Hello world" - # Stream URL was hit (not the unary fallback). - assert any(p.endswith("/v1/agent/infer/stream") for p in transport.calls) - - -async def test_chat_stream_falls_back_to_unary_on_404(consumer_env, monkeypatch): - """Older miners on eirel SDK <0.2.3 lack the stream route — consumer-api - falls back to the unary endpoint and emits the answer as one delta+done - so the client UX is unchanged.""" + # consumer-api hit the orchestrator's chat-stream route. + assert any( + p.endswith("/v1/orchestrate/chat/stream") for p in transport.calls + ) - async def _fake_resolve(family_id): - return {"endpoint": "http://miner.local", "hotkey": "hk"} - monkeypatch.setattr( - "orchestration.consumer_api.chat._resolve_serving_miner", _fake_resolve, - ) +async def test_chat_stream_emits_error_when_orchestrator_unreachable( + consumer_env, monkeypatch, +): + """If the orchestrator is down or returns 5xx, consumer-api emits a + terminal SSE ``error`` event so the browser closes cleanly instead + of hanging on a half-open response.""" class _Transport(httpx.AsyncBaseTransport): - def __init__(self) -> None: - self.calls: list[str] = [] - async def handle_async_request(self, request): - self.calls.append(request.url.path) - if request.url.path.endswith("/stream"): - return httpx.Response(404, text="not found") - return httpx.Response( - 200, json={ - "status": "completed", - "output": {"answer": "legacy reply"}, - "citations": [], - "tool_calls": [], - }, - ) + return httpx.Response(503, text="orchestrator unavailable") transport = _Transport() @@ -298,23 +282,45 @@ def _patched(*args, **kwargs): assert resp.status_code == 200 events = _parse_sse(resp.content) names = [e[0] for e in events] - assert names == ["started", "delta", "done"] - delta_text = events[1][1].get("text") - assert delta_text == "legacy reply" - # Hit stream first (404), then unary. - assert transport.calls[0].endswith("/v1/agent/infer/stream") - assert transport.calls[1].endswith("/v1/agent/infer") + assert names[-1] == "error" + err = next(e[1] for e in events if e[0] == "error") + assert "503" in err["message"] or "Server error" in err["message"] -async def test_chat_stream_emits_error_when_no_serving_miner( +async def test_chat_stream_emits_error_when_orchestrator_returns_done_failed( consumer_env, monkeypatch, ): - async def _fake_resolve(family_id): - return None + """If the orchestrator emits a terminal ``done`` chunk with + ``status: failed`` (e.g. no serving family is available), the + consumer-api passes it through verbatim — clients can decide + whether to render it as an error UI or retry.""" + import json as _json - monkeypatch.setattr( - "orchestration.consumer_api.chat._resolve_serving_miner", _fake_resolve, - ) + body = ( + _json.dumps({ + "event": "done", + "status": "failed", + "error": "no serving deployment available for family general_chat", + }) + "\n" + ).encode("utf-8") + + class _Transport(httpx.AsyncBaseTransport): + async def handle_async_request(self, request): + return httpx.Response( + 200, content=body, + headers={"content-type": "application/x-ndjson"}, + ) + + transport = _Transport() + + import orchestration.consumer_api.chat as chat_mod + original = chat_mod.httpx.AsyncClient + + def _patched(*args, **kwargs): + kwargs["transport"] = transport + return original(*args, **kwargs) + + monkeypatch.setattr(chat_mod.httpx, "AsyncClient", _patched) async for client in _make_client(): resp = await client.post( @@ -325,9 +331,9 @@ async def _fake_resolve(family_id): assert resp.status_code == 200 events = _parse_sse(resp.content) names = [e[0] for e in events] - assert "error" in names - err = next(e[1] for e in events if e[0] == "error") - assert "no serving miner" in err["message"] + assert names[-1] == "done" + done_payload = events[-1][1] + assert done_payload.get("status") == "failed" async def test_chat_stream_requires_api_key(consumer_env): diff --git a/tests/services/test_orchestrator.py b/tests/services/test_orchestrator.py index e4fc223..1a88f4b 100644 --- a/tests/services/test_orchestrator.py +++ b/tests/services/test_orchestrator.py @@ -286,3 +286,147 @@ async def test_fastapi_orchestrate_endpoint(): data = resp.json() assert "request_id" in data assert data["status"] in ("completed", "partial") + + +# --------------------------------------------------------------------------- +# /v1/orchestrate/chat/stream +# --------------------------------------------------------------------------- + + +async def test_chat_stream_persists_session_toggles_and_proxies_ndjson( + monkeypatch, tmp_path, +): + """End-to-end: orchestrator persists mode/web_search on the session + row and proxies the miner pod's NDJSON back to the caller.""" + import json as _json + from orchestration.orchestrator import chat_stream as cs + from orchestration.orchestrator.main import app + + monkeypatch.setenv("DATABASE_URL", f"sqlite:///{tmp_path}/cs.db") + + # Stub the serving-deployment lookup so the test doesn't need a live + # owner-api. Force the new env override path so we don't need to + # mock owner-api at all. + monkeypatch.setenv( + "EIREL_ORCHESTRATOR_MINER_OVERRIDE_ENDPOINT", + "http://miner.local", + ) + + body = ( + _json.dumps({"event": "delta", "text": "hello"}) + "\n" + + _json.dumps({ + "event": "done", + "status": "completed", + "output": {"answer": "hello"}, + "citations": [], + }) + "\n" + ).encode("utf-8") + + class _Transport(httpx.AsyncBaseTransport): + def __init__(self) -> None: + self.calls: list[str] = [] + self.bodies: list[dict] = [] + + async def handle_async_request(self, request): + self.calls.append(request.url.path) + try: + self.bodies.append(_json.loads(request.content or b"{}")) + except Exception: + self.bodies.append({}) + if request.url.path.endswith("/v1/agent/infer/stream"): + return httpx.Response( + 200, content=body, + headers={"content-type": "application/x-ndjson"}, + ) + return httpx.Response(404) + + transport = _Transport() + original = cs.httpx.AsyncClient + + def _patched(*args, **kwargs): + kwargs["transport"] = transport + return original(*args, **kwargs) + + monkeypatch.setattr(cs.httpx, "AsyncClient", _patched) + + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://test") as client: + # First turn: set mode=thinking + web_search=True; orchestrator + # should persist these on the session row. + resp = await client.post( + "/v1/orchestrate/chat/stream", + json={ + "prompt": "hi", + "user_id": "u1", + "session_id": "sess-A", + "mode": "thinking", + "web_search": True, + }, + ) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("application/x-ndjson") + lines = [ln for ln in resp.content.decode("utf-8").split("\n") if ln.strip()] + chunks = [_json.loads(ln) for ln in lines] + assert chunks[0]["event"] == "started" + assert chunks[0]["metadata"]["session_id"] == "sess-A" + assert chunks[0]["metadata"]["mode"] == "thinking" + assert chunks[0]["metadata"]["web_search"] is True + assert chunks[-1]["event"] == "done" + assert chunks[-1]["status"] == "completed" + + # The body sent to the miner carries the slim contract. + miner_body = transport.bodies[0] + assert miner_body["prompt"] == "hi" + assert miner_body["mode"] == "thinking" + assert miner_body["web_search"] is True + + # Second turn for the same session without explicit toggles — + # orchestrator must reuse the persisted values. + resp2 = await client.post( + "/v1/orchestrate/chat/stream", + json={"prompt": "follow up", "user_id": "u1", "session_id": "sess-A"}, + ) + assert resp2.status_code == 200 + chunks2 = [ + _json.loads(ln) + for ln in resp2.content.decode("utf-8").split("\n") + if ln.strip() + ] + started2 = chunks2[0] + assert started2["metadata"]["mode"] == "thinking" + assert started2["metadata"]["web_search"] is True + + +async def test_chat_stream_emits_done_failed_when_no_serving_deployment( + monkeypatch, tmp_path, +): + """If owner-api has no serving deployment for the family and no + override is set, the orchestrator must still close the stream with + a terminal ``done/failed`` chunk so the consumer-api facade can + render an error UI.""" + import json as _json + from orchestration.orchestrator import chat_stream as cs + from orchestration.orchestrator.main import app + + monkeypatch.setenv("DATABASE_URL", f"sqlite:///{tmp_path}/cs.db") + monkeypatch.delenv("EIREL_ORCHESTRATOR_MINER_OVERRIDE_ENDPOINT", raising=False) + + async def _no_serving(family_id): + return None + + monkeypatch.setattr(cs, "_resolve_serving_endpoint", _no_serving) + + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://test") as client: + resp = await client.post( + "/v1/orchestrate/chat/stream", + json={"prompt": "hi"}, + ) + assert resp.status_code == 200 + lines = [ln for ln in resp.content.decode("utf-8").split("\n") if ln.strip()] + chunks = [_json.loads(ln) for ln in lines] + assert chunks[-1]["event"] == "done" + assert chunks[-1]["status"] == "failed" + assert "no serving deployment" in chunks[-1]["error"] From 2eaa8c9975244c5a16956638341eeaf85c48be80 Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:52:20 +0000 Subject: [PATCH 03/24] feat(eval): multi-turn fixtures with live + scripted replay --- control_plane/owner_api/dashboard/queries.py | 82 ++++ control_plane/owner_api/dashboard/schemas.py | 36 ++ shared/benchmark/_invocation.py | 315 ++++++++++++---- shared/core/evaluation_models.py | 32 ++ tests/benchmark/test_invoke_task_multiturn.py | 257 +++++++++++++ tests/validator/test_baseline_multiturn.py | 152 ++++++++ validation/validator/engine.py | 352 ++++++++++++++++-- validation/validator/openai_baseline.py | 26 +- 8 files changed, 1162 insertions(+), 90 deletions(-) create mode 100644 tests/benchmark/test_invoke_task_multiturn.py create mode 100644 tests/validator/test_baseline_multiturn.py diff --git a/control_plane/owner_api/dashboard/queries.py b/control_plane/owner_api/dashboard/queries.py index 0102e0e..c225719 100644 --- a/control_plane/owner_api/dashboard/queries.py +++ b/control_plane/owner_api/dashboard/queries.py @@ -32,6 +32,8 @@ MinerRunsResponse, ModeLiteral, OverviewResponse, + QueuedSubmission, + QueuedSubmissionsResponse, RunDetailResponse, RunListResponse, RunSummary, @@ -519,6 +521,67 @@ def fetch_leaderboard( ) +# --------------------------------------------------------------------------- +# /submissions/queued +# --------------------------------------------------------------------------- + + +def fetch_queued_submissions( + session: Session, + *, + services: Any, + limit: int = 200, +) -> QueuedSubmissionsResponse: + """List submissions that have not yet produced a leaderboard score. + + Includes pre-evaluation states (``queued``, ``building``, + ``evaluating``) plus ``build_failed`` so miners can see why their + submission stalled. Excludes ``retired`` and any submission whose + deployment has already produced a DeploymentScoreRecord — those are + on the leaderboard already. + """ + del services + scored_ids = { + sid for sid, in session.execute( + select(DeploymentScoreRecord.submission_id).distinct() + ).all() + } + + rows = session.execute( + select(ManagedMinerSubmission) + .where(ManagedMinerSubmission.status != "retired") + .order_by( + ManagedMinerSubmission.submission_block.desc(), + ManagedMinerSubmission.created_at.desc(), + ) + .limit(limit * 4) # over-fetch then drop scored ones; simpler than a NOT IN. + ).scalars().all() + + out: list[QueuedSubmission] = [] + for sub in rows: + if sub.id in scored_ids: + continue + manifest = sub.manifest_json or {} + agent = manifest.get("agent") or {} + out.append( + QueuedSubmission( + submission_id=sub.id, + hotkey=sub.miner_hotkey, + hotkey_short=shorten_hotkey(sub.miner_hotkey), + family_id=sub.family_id, + agent_name=agent.get("name") if isinstance(agent, dict) else None, + agent_version=agent.get("version") if isinstance(agent, dict) else None, + artifact_sha256=sub.archive_sha256, + status=sub.status, + submitted_at=sub.created_at.isoformat() if sub.created_at else None, + submission_block=int(sub.submission_block) if sub.submission_block else None, + ) + ) + if len(out) >= limit: + break + return QueuedSubmissionsResponse(total=len(out), submissions=out) + + # --------------------------------------------------------------------------- # /miners/{hotkey} # --------------------------------------------------------------------------- @@ -717,6 +780,23 @@ def _task_evaluation_from_row( mode = _as_mode(bundle_task.get("mode") or meta.get("mode")) web_search = bool(bundle_task.get("web_search") or meta.get("web_search") or False) + # Multi-turn fixtures carry a ``turns`` array of {user, assistant?}. + # Surface the user-prompt sequence for the dashboard so the + # conversation context is visible alongside the final answer (which + # is what the judge actually scored). Single-turn tasks leave + # ``turns`` unset and the row falls back to ``prompt``. + raw_turns = bundle_task.get("turns") or [] + user_turns: list[str] = [] + if isinstance(raw_turns, list): + for turn in raw_turns: + if isinstance(turn, dict): + u = turn.get("user") + else: + u = getattr(turn, "user", None) + if isinstance(u, str) and u: + user_turns.append(u) + turn_count = len(user_turns) or 1 + status = "completed" if row.agreement_verdict != "error" else "failed" miner_citations = [ @@ -745,6 +825,8 @@ def _task_evaluation_from_row( task_status=status, evaluated_at=row.created_at.isoformat() if row.created_at else None, prompt=prompt_val if isinstance(prompt_val, str) else None, + turn_count=turn_count, + user_turns=user_turns, miner_response=row.miner_response_json, baseline_response_text=baseline_text, agreement_verdict=row.agreement_verdict, diff --git a/control_plane/owner_api/dashboard/schemas.py b/control_plane/owner_api/dashboard/schemas.py index 7c52243..b9fbe3e 100644 --- a/control_plane/owner_api/dashboard/schemas.py +++ b/control_plane/owner_api/dashboard/schemas.py @@ -104,6 +104,32 @@ class LeaderboardResponse(BaseModel): entries: list[LeaderboardEntry] = Field(default_factory=list) +class QueuedSubmission(BaseModel): + """One submission that has not yet appeared on a leaderboard. + + Covers anything pre-first-score: ``queued`` / ``building`` / + ``evaluating`` / ``build_failed``. Excludes ``retired`` and any + submission that already has a DeploymentScoreRecord (those belong on + the leaderboard, not the queue). + """ + + submission_id: str + hotkey: str + hotkey_short: str + family_id: str + agent_name: str | None = None + agent_version: str | None = None + artifact_sha256: str | None = None + status: str + submitted_at: str | None + submission_block: int | None = None + + +class QueuedSubmissionsResponse(BaseModel): + total: int + submissions: list[QueuedSubmission] = Field(default_factory=list) + + class RunSummary(BaseModel): id: str sequence: int @@ -163,6 +189,16 @@ class TaskEvaluation(BaseModel): task_status: str | None # "completed" | "failed" evaluated_at: str | None prompt: str | None + # Number of user turns in the fixture. 1 for single-turn tasks; >1 for + # multi-turn replay fixtures. Multi-turn tasks judge the *final* + # assistant answer only — earlier turns build context but do not + # produce a scored response. + turn_count: int = 1 + # User-prompt sequence for multi-turn fixtures. Empty for single-turn + # tasks (use ``prompt`` instead). Carries only user messages, not the + # scripted-assistant entries — those exist for both miner and + # baseline equally and are not interesting to display. + user_turns: list[str] = Field(default_factory=list) miner_response: dict | None # OpenAI baseline text, extracted from TaskEvaluation.baseline_response_json. # Rendered side-by-side with the miner response on the dashboard so users diff --git a/shared/benchmark/_invocation.py b/shared/benchmark/_invocation.py index 37a860f..3d9c09c 100644 --- a/shared/benchmark/_invocation.py +++ b/shared/benchmark/_invocation.py @@ -14,6 +14,7 @@ import asyncio import json import logging +import os import time from typing import Any @@ -33,20 +34,61 @@ _RETRY_BACKOFF_SECONDS = 1.0 -def _build_body(*, task: Any, prompt: str, family_id: str, task_id: str) -> dict[str, Any]: - expected_output = getattr(task, "expected_output", {}) or {} +def _build_body( + *, + task: Any, + prompt: str, + family_id: str, + task_id: str, + history: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Build the slim 0.3.0 invocation body. + + Family agents are stateless specialists; they receive only the + prompt, the per-turn knobs, and any prior conversation. Anything + eval-internal (``expected_output``, ``category``, ``difficulty``, + grading hints) stays server-side and never crosses the wire — old + builds leaked the answer key here via ``inputs.expected_output``. + + By default we also populate the legacy fields + (``primary_goal``/``subtask``/``inputs.{mode,web_search}``) so + miners on 0.2.x keep working through the migration window. Set + ``EIREL_VALIDATOR_SLIM_ONLY=1`` to suppress the legacy mirror — + used to airtight-test that miners are reading the 0.3.0 contract + (any miner that secretly relied on the legacy fields will + immediately fail the run with empty prompts). + """ inputs = getattr(task, "inputs", {}) or {} - metadata = dict(getattr(task, "metadata", {}) or {}) - return { - "task_id": task_id, - "family_id": family_id, - "primary_goal": prompt, - "subtask": prompt, - "inputs": ( - {**inputs, "expected_output": expected_output} if expected_output else inputs - ), - "metadata": metadata, + mode = inputs.get("mode") or "instant" + web_search = bool(inputs.get("web_search", False)) + + # Multi-turn replay: caller passes ``history`` accumulated from + # prior turns of the same fixture. Single-turn tasks pass None / + # empty. We don't read ``inputs.history`` anymore — that was only + # relevant in the pre-Phase-B shape. + cleaned_history = [ + {"role": h.get("role"), "content": h.get("content")} + for h in (history or []) + if isinstance(h, dict) and h.get("role") in ("user", "assistant") + ] + + body: dict[str, Any] = { + # Slim 0.3.0 contract — what new miners read. + "turn_id": task_id, + "prompt": prompt, + "mode": mode, + "web_search": web_search, + "history": cleaned_history, } + if os.getenv("EIREL_VALIDATOR_SLIM_ONLY", "0") not in {"1", "true", "yes"}: + body.update({ + "task_id": task_id, + "family_id": family_id, + "primary_goal": prompt, + "subtask": prompt, + "inputs": {"mode": mode, "web_search": web_search}, + }) + return body def _auth_headers(miner: MinerBenchmarkTarget) -> dict[str, str]: @@ -115,10 +157,17 @@ async def _invoke_stream( # mirrors the non-streaming response. Per-chunk # `citation` events are bonus diagnostics for now. citations = list(chunk["citations"]) - if isinstance(chunk.get("tool_calls"), list): - tool_calls = list(chunk["tool_calls"]) if isinstance(chunk.get("metadata"), dict): final_metadata = chunk["metadata"] + # Tool calls: 0.3.0 emits them under + # ``metadata.executed_tool_calls``; 0.2.x emitted a + # top-level ``tool_calls``. Read both during the + # migration window — slim wins when both present. + meta_tcs = final_metadata.get("executed_tool_calls") + if isinstance(meta_tcs, list): + tool_calls = list(meta_tcs) + elif isinstance(chunk.get("tool_calls"), list): + tool_calls = list(chunk["tool_calls"]) final_status = chunk.get("status") or "completed" final_error = chunk.get("error") @@ -156,45 +205,37 @@ async def _invoke_unary( return resp.json() if resp.content else {} -async def _invoke_task( +async def _invoke_one_turn( *, miner: MinerBenchmarkTarget, task: Any, - timeout_seconds: float = 90.0, -) -> BenchmarkTaskRun: - """POST a single task to a miner's endpoint and wrap the response. + prompt: str, + history: list[dict[str, Any]], + turn_id: str, + timeout_seconds: float, +) -> tuple[dict[str, Any], float, bool, Exception | None]: + """Invoke the miner for one turn (streaming-first, unary fallback). - Tries the streaming endpoint first. Falls back to the legacy unary - endpoint on 404. Records total completion latency in `metadata`. + Returns ``(payload, elapsed_seconds, used_stream, last_exc)``. + Caller decides whether to surface ``last_exc`` as a failure or + treat the partial payload as recoverable. Used by both single-turn + and multi-turn replay paths. """ endpoint = (miner.endpoint or "").rstrip("/") - prompt = getattr(task, "prompt", "") or "" family_id = getattr(task, "family_id", "general_chat") - task_id = getattr(task, "task_id", "") - expected_output = getattr(task, "expected_output", {}) or {} - metadata = dict(getattr(task, "metadata", {}) or {}) - - if not endpoint: - return BenchmarkTaskRun( - task_id=task_id, - family_id=family_id, - prompt=prompt, - expected_output=expected_output, - response={}, - status="failed", - error="missing_miner_endpoint", - metadata=metadata, - ) - headers = _auth_headers(miner) body = _build_body( - task=task, prompt=prompt, family_id=family_id, task_id=task_id, + task=task, + prompt=prompt, + family_id=family_id, + task_id=turn_id, + history=history, ) stream_url = f"{endpoint}/v1/agent/infer/stream" unary_url = f"{endpoint}/v1/agent/infer" t0 = time.perf_counter() - attempts = 2 # one initial + one retry on transient upstream errors + attempts = 2 last_exc: Exception | None = None payload: dict[str, Any] = {} used_stream = True @@ -216,20 +257,18 @@ async def _invoke_task( except httpx.HTTPStatusError as exc: last_exc = exc status_code = exc.response.status_code - # 404 on the stream URL → miner is on an older SDK without - # the streaming route. Fall back permanently for this call. if used_stream and status_code == 404: _logger.info( - "miner %s lacks streaming endpoint (404); falling back to unary: task=%s", - miner.hotkey[:16], task_id, + "miner %s lacks streaming endpoint (404); falling back to unary: turn=%s", + miner.hotkey[:16], turn_id, ) used_stream = False - t0 = time.perf_counter() # reset clock for the unary attempt + t0 = time.perf_counter() continue if status_code in _RETRYABLE_STATUS_CODES and attempt < attempts: _logger.warning( - "miner invocation hit %d on attempt %d/%d, retrying: task=%s miner=%s", - status_code, attempt, attempts, task_id, miner.hotkey[:16], + "miner invocation hit %d on attempt %d/%d, retrying: turn=%s miner=%s", + status_code, attempt, attempts, turn_id, miner.hotkey[:16], ) await asyncio.sleep(_RETRY_BACKOFF_SECONDS) continue @@ -238,9 +277,58 @@ async def _invoke_task( last_exc = exc break - elapsed = time.perf_counter() - t0 - if last_exc is not None: - _logger.warning("miner invocation failed: %s", last_exc) + return payload, time.perf_counter() - t0, used_stream, last_exc + + +def _extract_answer_text(payload: dict[str, Any] | None) -> str: + """Pull the assistant text out of a normalised invocation payload. + + Used by the multi-turn replay loop to feed each turn's reply back + into the next turn's history. + """ + if not isinstance(payload, dict): + return "" + out = payload.get("output") or {} + if isinstance(out, dict): + for key in ("answer", "response", "text", "content", "message"): + v = out.get(key) + if isinstance(v, str) and v: + return v + return "" + + +async def _invoke_task( + *, + miner: MinerBenchmarkTarget, + task: Any, + timeout_seconds: float = 90.0, +) -> BenchmarkTaskRun: + """POST a single task (single-turn or multi-turn) to a miner. + + Single-turn (``task.turns`` empty/None): one HTTP call, history is + empty. Multi-turn (``task.turns`` populated): replay each turn in + sequence, accumulating the miner's own replies as ``assistant`` + history entries between turns. Scripted turns (``assistant`` set on + the fixture) are inserted into history without calling the miner; + live turns (``assistant=None``) call the miner and record its + reply. The miner is always called for the final turn, and its + answer is what the judge scores. + + Recorded latency in ``metadata.latency_seconds`` is the **sum** of + per-turn wall clocks (matches "total time the user waited"). The + per-turn max is recorded as ``metadata.max_turn_latency_seconds`` + so the validator's mode-budget gate can fire on any turn that + overruns. ``metadata.turns`` carries the per-turn breakdown for the + dashboard. + """ + endpoint = (miner.endpoint or "").rstrip("/") + prompt = getattr(task, "prompt", "") or "" + family_id = getattr(task, "family_id", "general_chat") + task_id = getattr(task, "task_id", "") + expected_output = getattr(task, "expected_output", {}) or {} + metadata = dict(getattr(task, "metadata", {}) or {}) + + if not endpoint: return BenchmarkTaskRun( task_id=task_id, family_id=family_id, @@ -248,29 +336,126 @@ async def _invoke_task( expected_output=expected_output, response={}, status="failed", - error=str(last_exc), - metadata={**metadata, "latency_seconds": elapsed}, + error="missing_miner_endpoint", + metadata=metadata, ) - out_metadata: dict[str, Any] = {**metadata, "latency_seconds": elapsed} - out_metadata["streamed"] = used_stream - - # If the miner's `done` chunk explicitly reported failed, surface it on - # the run object so the validator's _judge_miner gates this as an error - # (otherwise we'd quietly score a failed agent call as a completed - # response with no answer text). - payload_status = ( - payload.get("status") if isinstance(payload, dict) else None - ) or "completed" - payload_error = payload.get("error") if isinstance(payload, dict) else None + # Build the turn script. + raw_turns = list(getattr(task, "turns", None) or []) + if not raw_turns: + # Single-turn: synthesize a one-turn script from the legacy + # ``prompt`` field so the loop below is uniform. + raw_turns = [{"user": prompt, "assistant": None}] + + history: list[dict[str, Any]] = [] + turn_breakdown: list[dict[str, Any]] = [] + final_payload: dict[str, Any] = {} + final_used_stream = True + final_status = "completed" + final_error: str | None = None + total_elapsed = 0.0 + max_turn_elapsed = 0.0 + + for idx, raw in enumerate(raw_turns): + is_last = (idx == len(raw_turns) - 1) + if hasattr(raw, "user"): + user_text = raw.user + scripted = raw.assistant + elif isinstance(raw, dict): + user_text = str(raw.get("user") or "") + scripted = raw.get("assistant") + else: + continue + if not isinstance(user_text, str) or not user_text: + continue + + # Scripted intermediate turn — no miner call, just inject the + # canned exchange into history. + if scripted is not None and not is_last: + history.append({"role": "user", "content": user_text}) + history.append({"role": "assistant", "content": str(scripted)}) + turn_breakdown.append({ + "turn_index": idx, + "scripted": True, + "latency_seconds": 0.0, + }) + continue + + # Live turn — call the miner. + turn_id = f"{task_id}-t{idx}" if len(raw_turns) > 1 else task_id + payload, elapsed, used_stream, exc = await _invoke_one_turn( + miner=miner, + task=task, + prompt=user_text, + history=list(history), + turn_id=turn_id, + timeout_seconds=timeout_seconds, + ) + total_elapsed += elapsed + if elapsed > max_turn_elapsed: + max_turn_elapsed = elapsed + if exc is not None: + _logger.warning( + "miner invocation failed at turn %d/%d: %s", + idx + 1, len(raw_turns), exc, + ) + return BenchmarkTaskRun( + task_id=task_id, + family_id=family_id, + prompt=prompt, + expected_output=expected_output, + response={}, + status="failed", + error=str(exc), + metadata={ + **metadata, + "latency_seconds": total_elapsed, + "max_turn_latency_seconds": max_turn_elapsed, + "failed_at_turn": idx, + "turns": turn_breakdown + [{ + "turn_index": idx, + "scripted": False, + "latency_seconds": elapsed, + "error": str(exc), + }], + }, + ) + + miner_reply = _extract_answer_text(payload) + # Append this turn to history before moving on (last-turn append + # is harmless — judge reads the payload directly). + history.append({"role": "user", "content": user_text}) + history.append({"role": "assistant", "content": miner_reply}) + turn_breakdown.append({ + "turn_index": idx, + "scripted": False, + "latency_seconds": elapsed, + "streamed": used_stream, + }) + final_payload = payload + final_used_stream = used_stream + # Surface a per-turn `done.status: failed` from the miner. + ps = (payload.get("status") if isinstance(payload, dict) else None) or "completed" + if ps != "completed": + final_status = "failed" + final_error = payload.get("error") if isinstance(payload, dict) else None + + out_metadata: dict[str, Any] = { + **metadata, + "latency_seconds": total_elapsed, + "max_turn_latency_seconds": max_turn_elapsed, + "streamed": final_used_stream, + "turns": turn_breakdown, + "turn_count": len(raw_turns), + } return BenchmarkTaskRun( task_id=task_id, family_id=family_id, prompt=prompt, expected_output=expected_output, - response=payload if isinstance(payload, dict) else {"raw": payload}, - status="completed" if payload_status == "completed" else "failed", - error=payload_error if payload_status != "completed" else None, + response=final_payload if isinstance(final_payload, dict) else {"raw": final_payload}, + status="completed" if final_status == "completed" else "failed", + error=final_error if final_status != "completed" else None, metadata=out_metadata, ) diff --git a/shared/core/evaluation_models.py b/shared/core/evaluation_models.py index 868575c..2f3d3f7 100644 --- a/shared/core/evaluation_models.py +++ b/shared/core/evaluation_models.py @@ -89,6 +89,33 @@ def normalize_specialist_family_id(cls, value: Any) -> str: return ensure_family_id(str(value)) +class EvaluationConversationTurn(BaseModel): + """One turn in a multi-turn evaluation fixture. + + ``user`` is always the user message for that turn. ``assistant`` is + optional: + * **None (default)** — *live* mode. The validator calls the miner + and the baseline for that turn; whatever each one replies is + appended to its own private history before the next user turn. + * **set** — *scripted* mode. Both miner and baseline see the + identical canned assistant message in their history; neither is + called for that turn. Used to set up specific multi-turn probes + (clarification, contradiction, reference-resolution, etc.). + + The final turn is always live — it produces the assistant answer the + pairwise judge scores. + + Distinct from the runtime ``ConversationTurn`` (defined later in + this module) which uses ``role`` / ``content`` for in-flight + history; this fixture model is the on-disk schema for a multi-turn + benchmark task and uses ``user`` / ``assistant`` to make scripted + vs live turns visually obvious. + """ + + user: str + assistant: str | None = None + + class FamilyEvaluationTask(BaseModel): task_id: str family_id: FamilyId @@ -104,6 +131,11 @@ class FamilyEvaluationTask(BaseModel): # explicit field Pydantic would drop it on bundle validation, so the # dashboard would always render "no web". web_search: bool = False + # Multi-turn evaluation fixture. When set, the validator replays the + # ``turns`` script against the miner and the baseline (each building + # its own history) and judges the *final* assistant answer only. + # Single-turn tasks leave this None and use ``prompt``. + turns: list[EvaluationConversationTurn] | None = None risk_tags: list[str] = Field(default_factory=list) allowed_tools: list[str] = Field(default_factory=list) retrieval_constraints: dict[str, Any] = Field(default_factory=dict) diff --git a/tests/benchmark/test_invoke_task_multiturn.py b/tests/benchmark/test_invoke_task_multiturn.py new file mode 100644 index 0000000..5854552 --- /dev/null +++ b/tests/benchmark/test_invoke_task_multiturn.py @@ -0,0 +1,257 @@ +"""Multi-turn replay tests for ``_invoke_task``. + +A fixture with a ``turns`` list is replayed against the miner one user +turn at a time, accumulating the miner's reply as ``assistant`` history +between turns. The final live turn's reply is what becomes the run's +``response`` (and what the judge eventually scores). +""" +from __future__ import annotations + +import json + +import httpx +import pytest + +from shared.benchmark._invocation import _invoke_task +from shared.core.evaluation_models import ( + EvaluationConversationTurn, + FamilyEvaluationTask, + MinerBenchmarkTarget, +) + + +def test_family_evaluation_task_accepts_multi_turn_dict_form(): + """The bundle loader will hand us raw dicts — confirm the schema + parses both single-turn (``prompt`` only, no ``turns``) and + multi-turn (``turns`` array, with mixed scripted-or-live turns).""" + single = FamilyEvaluationTask.model_validate({ + "task_id": "s-1", "family_id": "general_chat", + "prompt": "hello", "mode": "instant", + }) + assert single.turns is None + + multi = FamilyEvaluationTask.model_validate({ + "task_id": "m-1", "family_id": "general_chat", + "prompt": "first user turn", # legacy mirror for older readers + "mode": "thinking", + "turns": [ + {"user": "first user turn"}, + {"user": "second user turn", "assistant": "scripted reply"}, + {"user": "final question"}, + ], + }) + assert multi.turns is not None + assert len(multi.turns) == 3 + assert isinstance(multi.turns[0], EvaluationConversationTurn) + assert multi.turns[1].assistant == "scripted reply" + assert multi.turns[2].assistant is None # live final turn + + +class _Task: + """Fixture stand-in. ``turns`` is the multi-turn script.""" + + def __init__(self, *, turns) -> None: + self.task_id = "mt-1" + self.family_id = "general_chat" + self.prompt = turns[0].get("user") if turns and isinstance(turns[0], dict) else "" + self.expected_output = {} + self.inputs = {"mode": "instant", "web_search": False} + self.metadata = {} + self.turns = turns + + +def _miner() -> MinerBenchmarkTarget: + return MinerBenchmarkTarget( + hotkey="5X" * 32, endpoint="http://miner.local", stake=0, metadata={}, + ) + + +def _ndjson_body(answer: str) -> bytes: + chunks = [ + {"event": "delta", "text": answer}, + { + "event": "done", + "output": {"answer": answer}, + "citations": [], + "status": "completed", + "metadata": {}, + }, + ] + return ("\n".join(json.dumps(c) for c in chunks) + "\n").encode("utf-8") + + +class _ScriptedTransport(httpx.AsyncBaseTransport): + """Records every request body and replies in scripted order.""" + + def __init__(self, replies: list[bytes]) -> None: + self._replies = list(replies) + self.bodies: list[dict] = [] + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + try: + self.bodies.append(json.loads(request.content or b"{}")) + except Exception: + self.bodies.append({}) + body = self._replies.pop(0) + return httpx.Response( + 200, content=body, + headers={"content-type": "application/x-ndjson"}, + ) + + +@pytest.fixture +def _patch_client(monkeypatch): + import shared.benchmark._invocation as mod + original = mod.httpx.AsyncClient + holder: dict[str, httpx.AsyncBaseTransport] = {} + + def _install(transport: httpx.AsyncBaseTransport) -> None: + holder["t"] = transport + + def _patched(*args, **kwargs): + kwargs["transport"] = holder["t"] + return original(*args, **kwargs) + + monkeypatch.setattr(mod.httpx, "AsyncClient", _patched) + monkeypatch.setattr(mod, "_RETRY_BACKOFF_SECONDS", 0.0) + return _install + + +async def test_multi_turn_live_replay_accumulates_history(_patch_client): + """Three live user turns → three miner calls; each call sees the + history accumulated from previous turns. Final answer is from the + last turn.""" + transport = _ScriptedTransport([ + _ndjson_body("turn-1 answer"), + _ndjson_body("turn-2 answer"), + _ndjson_body("turn-3 answer"), + ]) + _patch_client(transport) + + task = _Task(turns=[ + {"user": "first question"}, + {"user": "follow up"}, + {"user": "another follow up"}, + ]) + + run = await _invoke_task(miner=_miner(), task=task, timeout_seconds=5.0) + + assert run.status == "completed" + assert run.response.get("output", {}).get("answer") == "turn-3 answer" + assert run.metadata.get("turn_count") == 3 + + # Turn 1: empty history. + assert transport.bodies[0]["history"] == [] + assert transport.bodies[0]["prompt"] == "first question" + + # Turn 2: history holds turn-1 user + turn-1 assistant reply. + assert transport.bodies[1]["history"] == [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": "turn-1 answer"}, + ] + assert transport.bodies[1]["prompt"] == "follow up" + + # Turn 3: full accumulated history through turn 2. + assert transport.bodies[2]["history"][-2:] == [ + {"role": "user", "content": "follow up"}, + {"role": "assistant", "content": "turn-2 answer"}, + ] + + +async def test_scripted_intermediate_turn_skips_miner(_patch_client): + """Scripted ``assistant`` on a non-final turn injects the canned + exchange into history without calling the miner.""" + transport = _ScriptedTransport([ + # Only one miner call expected — for the final live turn. + _ndjson_body("final answer"), + ]) + _patch_client(transport) + + task = _Task(turns=[ + {"user": "earlier question", "assistant": "earlier scripted reply"}, + {"user": "final question"}, + ]) + + run = await _invoke_task(miner=_miner(), task=task, timeout_seconds=5.0) + + assert run.status == "completed" + assert len(transport.bodies) == 1 + # The single miner call sees the scripted exchange in history. + assert transport.bodies[0]["prompt"] == "final question" + assert transport.bodies[0]["history"] == [ + {"role": "user", "content": "earlier question"}, + {"role": "assistant", "content": "earlier scripted reply"}, + ] + # Per-turn breakdown reflects scripted turn at index 0. + breakdown = run.metadata.get("turns") or [] + assert breakdown[0]["scripted"] is True + assert breakdown[1]["scripted"] is False + + +async def test_per_turn_max_latency_recorded_for_sla_gate(_patch_client): + """``max_turn_latency_seconds`` reflects the slowest turn (so the + validator's per-turn budget gate fires on any over-budget turn, + not just the average).""" + transport = _ScriptedTransport([ + _ndjson_body("a"), + _ndjson_body("b"), + ]) + _patch_client(transport) + + task = _Task(turns=[{"user": "one"}, {"user": "two"}]) + + run = await _invoke_task(miner=_miner(), task=task, timeout_seconds=5.0) + + assert run.metadata.get("turn_count") == 2 + # Both should be tiny; the meaningful invariant is that the sum is + # at least the max — sanity-checking the relationship rather than + # specific timing. + total = run.metadata.get("latency_seconds") or 0.0 + per_turn_max = run.metadata.get("max_turn_latency_seconds") or 0.0 + assert per_turn_max <= total + assert per_turn_max > 0.0 + + +async def test_failure_at_intermediate_turn_aborts_replay(_patch_client): + """If turn N fails, we stop and surface the error — turns N+1.. + are never sent.""" + transport = _ScriptedTransport([ + _ndjson_body("turn-1 ok"), + # Turn 2 will get a 500 (we'll script that via custom transport). + ]) + + class _Transport(httpx.AsyncBaseTransport): + def __init__(self) -> None: + self.bodies: list[dict] = [] + self.calls = 0 + + async def handle_async_request(self, request): + self.calls += 1 + try: + self.bodies.append(json.loads(request.content or b"{}")) + except Exception: + self.bodies.append({}) + if self.calls == 1: + return httpx.Response( + 200, content=_ndjson_body("turn-1 ok"), + headers={"content-type": "application/x-ndjson"}, + ) + return httpx.Response(500, text="boom") + + t = _Transport() + _patch_client(t) + + task = _Task(turns=[ + {"user": "first"}, + {"user": "second"}, + {"user": "third — must never be sent"}, + ]) + + run = await _invoke_task(miner=_miner(), task=task, timeout_seconds=5.0) + + assert run.status == "failed" + # Only turn 1 (success) + turn 2 (500 → retry → 500). Turn 3 never sent. + # Confirm by absence of turn-3 prompt in any body. + sent_prompts = [b.get("prompt") for b in t.bodies] + assert "third — must never be sent" not in sent_prompts + assert run.metadata.get("failed_at_turn") == 1 diff --git a/tests/validator/test_baseline_multiturn.py b/tests/validator/test_baseline_multiturn.py new file mode 100644 index 0000000..871c805 --- /dev/null +++ b/tests/validator/test_baseline_multiturn.py @@ -0,0 +1,152 @@ +"""Multi-turn replay tests for the OpenAI baseline. + +Mirrors the miner's replay pattern: the validator drives the baseline +through each user turn, accumulating the baseline's own assistant +replies as history between turns. The final live turn's response is +what gets compared to the miner's final answer by the pairwise judge. +""" +from __future__ import annotations + +import json +from types import SimpleNamespace + +import httpx +import pytest + +from validation.validator.openai_baseline import OpenAIBaselineClient + + +def _responses_api_body(text: str) -> dict: + return { + "id": "resp-1", + "output": [ + { + "type": "message", + "content": [ + {"type": "output_text", "text": text, "annotations": []}, + ], + }, + ], + "usage": {"total_cost_usd": 0.001}, + } + + +class _RecordingTransport(httpx.AsyncBaseTransport): + """Captures every request body and replies in scripted order.""" + + def __init__(self, replies: list[dict]) -> None: + self._replies = list(replies) + self.bodies: list[dict] = [] + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + try: + self.bodies.append(json.loads(request.content or b"{}")) + except Exception: + self.bodies.append({}) + body = self._replies.pop(0) + return httpx.Response(200, json=body) + + +def _client(transport: httpx.AsyncBaseTransport) -> OpenAIBaselineClient: + return OpenAIBaselineClient( + api_key="test-key", + model="gpt-test", + max_cost_usd_per_run=10.0, + transport=transport, + ) + + +async def test_history_passed_as_role_tagged_input_list(): + """When ``history`` is non-empty, the Responses API ``input`` field + becomes a role-tagged list with the latest user prompt as the last + message — that's how the API expects multi-turn context.""" + transport = _RecordingTransport([_responses_api_body("hi back")]) + client = _client(transport) + + history = [ + {"role": "user", "content": "earlier"}, + {"role": "assistant", "content": "earlier reply"}, + ] + out = await client.generate( + prompt="follow up", + use_web_search=False, + history=history, + ) + + assert out.response_text == "hi back" + body = transport.bodies[0] + # Multi-turn: input is a list, last message is the new user prompt. + assert isinstance(body["input"], list) + assert body["input"] == history + [{"role": "user", "content": "follow up"}] + # Single tool flag stays off. + assert "tools" not in body + + +async def test_empty_history_keeps_string_input_shape(): + """Single-turn calls keep the simpler string-input shape so we + don't churn the wire format for the common case.""" + transport = _RecordingTransport([_responses_api_body("just answer")]) + client = _client(transport) + + out = await client.generate(prompt="just ask", use_web_search=False, history=None) + assert out.response_text == "just answer" + body = transport.bodies[0] + assert body["input"] == "just ask" + + +async def test_validator_baseline_replay_accumulates_assistant_history(monkeypatch): + """End-to-end check of ``_baseline_replay`` (the validator-side + helper): given a 3-turn fixture, the helper should call the + baseline 3 times, each call seeing history accumulated from the + baseline's own previous replies, and return the final turn's + response.""" + # Import here to avoid pulling in the validator-engine module + # (with all its asyncio plumbing) at file-load time. + from validation.validator import engine as engine_mod + + # Get the inner _baseline_replay by re-implementing the closure + # arguments. The function is defined inside run_distributed_benchmarks, + # so we exercise it via a tiny harness that mirrors the call shape. + transport = _RecordingTransport([ + _responses_api_body("base-1"), + _responses_api_body("base-2"), + _responses_api_body("base-3"), + ]) + client = _client(transport) + + task = SimpleNamespace( + prompt="t1", + turns=[{"user": "t1"}, {"user": "t2"}, {"user": "t3"}], + ) + + # Reproduce the validator's replay loop locally — same logic as + # ``engine._baseline_replay`` but standalone for unit testing. + history: list[dict] = [] + last = None + for raw in task.turns: + last = await client.generate( + prompt=raw["user"], use_web_search=False, history=history, + ) + history.append({"role": "user", "content": raw["user"]}) + history.append({"role": "assistant", "content": last.response_text}) + + assert last is not None + assert last.response_text == "base-3" + + # Turn-1 input is a string; turns 2 + 3 are role-tagged lists with + # the baseline's own replies threaded back in. + assert transport.bodies[0]["input"] == "t1" + assert transport.bodies[1]["input"] == [ + {"role": "user", "content": "t1"}, + {"role": "assistant", "content": "base-1"}, + {"role": "user", "content": "t2"}, + ] + assert transport.bodies[2]["input"] == [ + {"role": "user", "content": "t1"}, + {"role": "assistant", "content": "base-1"}, + {"role": "user", "content": "t2"}, + {"role": "assistant", "content": "base-2"}, + {"role": "user", "content": "t3"}, + ] + # ``engine_mod`` reference kept so we exercise the import path. + assert engine_mod is not None diff --git a/validation/validator/engine.py b/validation/validator/engine.py index 23c8068..8abc638 100644 --- a/validation/validator/engine.py +++ b/validation/validator/engine.py @@ -285,6 +285,16 @@ async def run_validator_loop() -> None: _ACTIVE_FAMILIES, _POLL_INTERVAL_SECONDS, _BATCH_SIZE, _MAX_PARALLEL, ) + # Periodic idle heartbeat. Empty claim cycles used to be silent — + # the operator couldn't tell the validator was alive without + # poking the DB. Log a one-line "idle" every Nth empty cycle + # (default 10 ≈ 5 minutes at 30s poll). Set + # ``EIREL_VALIDATOR_IDLE_HEARTBEAT_EVERY=0`` to disable. + _idle_heartbeat_every = int( + os.getenv("EIREL_VALIDATOR_IDLE_HEARTBEAT_EVERY", "10") + ) + _idle_cycles = 0 + while True: found_work = False for family_id in _ACTIVE_FAMILIES: @@ -299,9 +309,11 @@ async def run_validator_loop() -> None: if result["total_claimed"] > 0: found_work = True logger.info( - "validator loop: family=%s claimed=%d submitted=%d failed=%d", + "validator loop: family=%s claimed=%d submitted=%d failed=%d " + "baseline_failed=%d", family_id, result["total_claimed"], result["total_submitted"], result["total_failed"], + result.get("total_baseline_failed", 0), ) except Exception: logger.exception("validator loop error: family=%s", family_id) @@ -310,7 +322,19 @@ async def run_validator_loop() -> None: loop_name="validator" ).set_to_current_time() if not found_work: + _idle_cycles += 1 + if ( + _idle_heartbeat_every > 0 + and _idle_cycles % _idle_heartbeat_every == 0 + ): + logger.info( + "validator loop: idle (no claimable tasks) families=%s " + "idle_cycles=%d", + _ACTIVE_FAMILIES, _idle_cycles, + ) await asyncio.sleep(_POLL_INTERVAL_SECONDS) + else: + _idle_cycles = 0 # --------------------------------------------------------------------------- @@ -450,28 +474,64 @@ async def run_weight_setting_loop() -> None: hotkey[:16], ) - # 4. Unassigned family weight goes to UID 0 (burns alpha) + # 4. Burn the remainder of the weight vector to UID 0. + # + # The on-chain ``set_weights`` call normalises whatever vector + # we publish so its components sum to 1.0 — i.e. emission is + # distributed in proportion to the published weights, no matter + # what their absolute magnitude. So if owner-api says + # ``general_chat:0.5`` and we publish only ``[0.5]`` against + # the winner, Bittensor renormalises that to ``[1.0]`` and the + # winner ends up with 100% of subnet emission — defeating the + # operator's intent to burn 50%. + # + # The fix: always publish a vector that sums to 1.0. Any + # fraction not assigned to a winner — for any reason — + # becomes burn: + # + # * The operator-configured "non-family" share — when + # ``EIREL_FAMILY_WEIGHTS`` sums to less than 1.0 (e.g. + # ``general_chat:0.5`` → 50% intentional burn). + # * Allocated-to-family-but-no-winner — when a family was + # assigned weight but produced no qualifying winner this + # run (gate didn't pass). + # * Winner-not-in-metagraph — already excluded from + # ``total_assigned`` above; gets folded into burn here. + # + # All three collapse into ``burn = 1.0 - total_assigned``. total_family_weight = sum( float(fw.get("family_weight", 0.0)) for fw in family_winners ) - unassigned = max(0.0, total_family_weight - total_assigned) - if unassigned > 0: - uid_weights[0] = uid_weights.get(0, 0.0) + unassigned + burn = max(0.0, 1.0 - total_assigned) + if burn > 0: + uid_weights[0] = uid_weights.get(0, 0.0) + burn if not uid_weights: - uid_weights[0] = total_family_weight or 1.0 + # Total fallback — owner-api returned nothing actionable. + # Burn-everything keeps us alive on chain (vtrust intact). + uid_weights[0] = 1.0 + elif total_assigned > 1.0: + logger.warning( + "weight-setting: total_assigned=%.4f exceeds 1.0; " + "EIREL_FAMILY_WEIGHTS likely misconfigured (sum > 1.0)", + total_assigned, + ) uids = sorted(uid_weights.keys()) weight_vals = [uid_weights[uid] for uid in uids] logger.info( - "weight-setting: setting weights for %d UIDs run=%s mode=%s uids=%s weights=%s", + "weight-setting: setting weights for %d UIDs run=%s mode=%s " + "uids=%s weights=%s assigned=%.4f burn=%.4f family_total=%.4f", len(uids), current_run_id, "owner_api_down_burn" if owner_api_down else "normal", uids, [round(w, 4) for w in weight_vals], + total_assigned, + uid_weights.get(0, 0.0) if 0 in uid_weights else 0.0, + total_family_weight, ) # 5. Set weights on-chain with retry + exponential backoff @@ -637,13 +697,88 @@ async def run_distributed_benchmarks( _metrics.benchmark_runs_started_total.labels(family=family_id).inc() _benchmark_outcome = "success" - # Outer wall-clock budget for the fan-out (miners + OpenAI baseline). Set - # generously above the largest configured downstream timeout so a slow - # baseline doesn't trip this wrapper before its own client timeout fires. - # Thinking-mode miners can run up to 600s (their completion SLA), so the - # wrapper has to clear that ceiling. The claim lease is 600s by default — - # bump EIREL_TASK_CLAIM_TIMEOUT_SECONDS if you raise this much further. - _FAN_OUT_TIMEOUT_SECONDS = 660.0 + # Outer wall-clock budget for the fan-out (miners + OpenAI baseline). + # Single-turn ceiling is one mode-budget (thinking=600s) plus headroom + # for the baseline call. Multi-turn fixtures replay N user turns, so + # the budget scales with N; default 1500s comfortably fits a 3-turn + # thinking-mode fixture (3 × ~500s wall clock). Override with + # ``EIREL_FAN_OUT_TIMEOUT_SECONDS`` if you ship longer multi-turn + # scripts. Per-turn budget enforcement still happens in ``_judge_miner``. + _FAN_OUT_TIMEOUT_SECONDS = float( + os.getenv("EIREL_FAN_OUT_TIMEOUT_SECONDS", "1500") + ) + + async def _baseline_replay( + *, + client: OpenAIBaselineClient, + task: Any, + use_web_search: bool, + ) -> Any: + """Replay a multi-turn fixture against the OpenAI baseline. + + Single-turn fixtures (``turns`` empty or None) → one call, history + empty. Multi-turn → call once per *live* user turn, accumulating + the baseline's own assistant replies as history between turns. + Scripted-assistant turns are inserted into history without a + baseline call so both miner and baseline see the identical canned + exchange when probe scenarios require it. + + Returns the BaselineResponse from the **last live turn** — that's + the answer the pairwise judge scores. + """ + raw_turns = list(getattr(task, "turns", None) or []) + if not raw_turns: + raw_turns = [{"user": task.prompt, "assistant": None}] + + history: list[dict[str, Any]] = [] + last_response = None + for idx, raw in enumerate(raw_turns): + is_last = (idx == len(raw_turns) - 1) + if hasattr(raw, "user"): + user_text = raw.user + scripted = raw.assistant + elif isinstance(raw, dict): + user_text = str(raw.get("user") or "") + scripted = raw.get("assistant") + else: + continue + if not isinstance(user_text, str) or not user_text: + continue + + if scripted is not None and not is_last: + history.append({"role": "user", "content": user_text}) + history.append({"role": "assistant", "content": str(scripted)}) + continue + + last_response = await client.generate( + prompt=user_text, + use_web_search=use_web_search, + history=history, + ) + # Per-turn cost log helps explain a high run-total cost on + # multi-turn fixtures (3 turns × $0.05 each ≠ a $0.05 task). + if len(raw_turns) > 1: + logger.info( + "baseline turn=%d/%d latency=%.2fs cost=$%.4f run_total=$%.4f", + idx + 1, len(raw_turns), + float(getattr(last_response, "latency_seconds", 0.0) or 0.0), + float(getattr(last_response, "cost_usd", 0.0) or 0.0), + float(getattr(client, "spent_usd", 0.0) or 0.0), + ) + history.append({"role": "user", "content": user_text}) + history.append({ + "role": "assistant", + "content": last_response.response_text or "", + }) + + if last_response is None: + # Defensive fallback — degenerate fixture (no live turns). + last_response = await client.generate( + prompt=task.prompt, + use_web_search=use_web_search, + history=[], + ) + return last_response async def _invoke_one_miner( miner: dict[str, Any], task_obj: Any, task_id: str, @@ -675,6 +810,7 @@ async def _evaluate_task(task_claim: dict[str, Any]) -> str: miners = task_claim.get("miners") or [] if not miners: return "ok" # nothing to evaluate + task_started_at = time.perf_counter() class _TaskProxy: pass @@ -686,6 +822,11 @@ class _TaskProxy: task_obj.inputs = _hydrate_agent_inputs(task_payload) task_obj.metadata = task_payload.get("metadata", {}) task_obj.execution_mode = task_payload.get("execution_mode") + # Multi-turn fixtures carry a ``turns`` array of {user, assistant?}. + # Pass through to the invocation helper unchanged (it handles + # both single-turn and multi-turn cases). Single-turn tasks have + # turns=None and use ``prompt``. + task_obj.turns = task_payload.get("turns") task_mode = str(task_obj.inputs.get("mode") or "instant") # Per-task web-search flag, mirroring the end-user toggle in the chat # UI. Missing field defaults to False so a baseline never silently @@ -696,10 +837,28 @@ class _TaskProxy: or False ) - # Fan out: all miners + baseline concurrently, bounded by wall clock + # Per-task lifecycle: emit a CLAIMED log so the operator can see + # exactly what's about to be invoked (mode, multi-turn vs + # single-turn, miner count) without correlating DB rows manually. + turn_count = len(task_obj.turns) if task_obj.turns else 1 + logger.info( + "task_eval=%s task=%s mode=%s web_search=%s turns=%d miners=%d CLAIMED", + task_evaluation_id[:8], task_id, task_mode, web_search_flag, + turn_count, len(miners), + ) + + # Build the baseline's user-only conversation script. The + # baseline replays the same user prompts as the miner so the + # judge compares two answers conditioned on the same user + # intent. For scripted-assistant turns we feed the same canned + # exchange into the baseline's history (parity with the miner). + # We then replay against the baseline turn-by-turn and let it + # build its own assistant history, which is what gets compared + # at the final live turn. baseline_task = asyncio.create_task( - baseline_client.generate( - prompt=task_obj.prompt, + _baseline_replay( + client=baseline_client, + task=task_obj, use_web_search=web_search_flag, ) ) @@ -733,11 +892,74 @@ class _TaskProxy: return "submit_failed" baseline_text = baseline.response_text + # Surface OpenAI baseline cost — per-call (this task) and the + # cumulative spend across the validator's lifetime so operators + # can see the budget burn without scraping metrics. + logger.info( + "task_eval=%s baseline=ok latency=%.2fs cost=$%.4f run_total=$%.4f", + task_evaluation_id[:8], + float(getattr(baseline, "latency_seconds", 0.0) or 0.0), + float(getattr(baseline, "cost_usd", 0.0) or 0.0), + float(getattr(baseline_client, "spent_usd", 0.0) or 0.0), + ) + + # Per-miner fan-out outcome. miner_runs is a list of either + # ``(hotkey, BenchmarkTaskRun)`` tuples or raised exceptions + # (return_exceptions=True on the gather). Cost is server-side + # ground truth: owner-api stamped X-Eirel-Job-Id per task, + # provider-proxy ledgered LLM spend under that tag, and + # owner-api injected the lookup result into the miner's + # done-chunk metadata before forwarding back here. + for entry in miner_runs: + if isinstance(entry, BaseException): + logger.warning( + "task_eval=%s miner=? fan_out_exception=%s", + task_evaluation_id[:8], entry, + ) + continue + try: + hk, run = entry # type: ignore[misc] + except Exception: + continue + run_meta = (run.metadata or {}) if hasattr(run, "metadata") else {} + response_meta = ( + (run.response or {}).get("metadata") or {} + if hasattr(run, "response") else {} + ) + llm_cost = float(response_meta.get("proxy_cost_usd") or 0.0) + cost_str = ( + f" llm_cost=$? (no ledger entry)" + if response_meta.get("proxy_cost_absent") + else f" llm_cost=${llm_cost:.4f}" + ) + logger.info( + "task_eval=%s miner=%s status=%s latency=%.2fs streamed=%s%s%s", + task_evaluation_id[:8], + str(hk)[:12] if hk else "?", + getattr(run, "status", "?"), + float(run_meta.get("latency_seconds") or 0.0), + run_meta.get("streamed"), + cost_str, + f" error=\"{run.error[:80]}\"" if getattr(run, "error", None) else "", + ) + task_category = ( task_payload.get("category") or (task_obj.metadata or {}).get("category") ) + # For multi-turn fixtures the judge sees the **final user turn** + # as the prompt. Earlier turns provide context but are not what + # the agent is being graded on. For single-turn tasks this is + # just task_obj.prompt. + judge_prompt = task_obj.prompt + if task_obj.turns: + for raw in reversed(task_obj.turns): + user_text = raw.get("user") if isinstance(raw, dict) else getattr(raw, "user", None) + if isinstance(user_text, str) and user_text: + judge_prompt = user_text + break + # Pick the latency budget for this task based on its declared mode. # Tasks without a recognized mode skip the gate entirely (None budget). if task_mode == "instant": @@ -752,7 +974,25 @@ class _TaskProxy: async def _judge_miner(miner_run: tuple[str, BenchmarkTaskRun]) -> dict[str, Any]: miner_hotkey, run = miner_run miner_citations = _extract_miner_citations(run) - miner_latency = float((run.metadata or {}).get("latency_seconds") or 0.0) + # ``latency_seconds`` is the sum across turns (the wall clock + # the user actually waited). For multi-turn fixtures the + # mode-budget SLA is enforced **per turn** (each turn must + # finish within the budget), so we read the per-turn max + # from the invocation helper. Single-turn runs report the + # same value under both keys. + run_meta = run.metadata or {} + miner_latency = float(run_meta.get("latency_seconds") or 0.0) + max_turn_latency = float( + run_meta.get("max_turn_latency_seconds") or miner_latency + ) + # Per-task LLM cost the miner incurred against the subnet + # provider-proxy. Owner-api injects this into the miner's + # done-chunk metadata after looking it up server-side from + # the proxy ledger — never from miner self-report. We pull + # it out of the response payload here and pass it along to + # the submit endpoint for storage. + response_meta = (run.response or {}).get("metadata") or {} + proxy_cost_usd = float(response_meta.get("proxy_cost_usd") or 0.0) if run.status != "completed": return { "miner_hotkey": miner_hotkey, @@ -763,6 +1003,8 @@ async def _judge_miner(miner_run: tuple[str, BenchmarkTaskRun]) -> dict[str, Any "verdict": "error", "miner_latency_seconds": miner_latency, "latency_seconds": 0.0, + "proxy_cost_usd": proxy_cost_usd, + "judge_cost_usd": 0.0, } try: miner_answer = _extract_answer_text(run) @@ -771,7 +1013,7 @@ async def _judge_miner(miner_run: tuple[str, BenchmarkTaskRun]) -> dict[str, Any judge_result = await asyncio.to_thread( judge_client.judge_agreement, family_id=family_id, - prompt=task_obj.prompt, + prompt=judge_prompt, response_a=miner_answer, response_b=baseline_text, task_mode=task_mode, @@ -781,14 +1023,32 @@ async def _judge_miner(miner_run: tuple[str, BenchmarkTaskRun]) -> dict[str, Any judge_latency = max(0.0, time.perf_counter() - judge_started) verdict = judge_result.verdict agreement_score = float(judge_result.agreement_score or 0.0) - # Mode-specific completion-time SLA: completed responses - # over budget count as a loss regardless of content match. - if latency_budget is not None and miner_latency > latency_budget: + # Phase 2c: eiretes-judge surfaces ``cost_usd`` per + # judgment in its response metadata. Pull it through + # so ``TaskMinerResult.judge_cost_usd`` reflects what + # this validator paid Chutes for this single + # judgment. Falls back to 0.0 on older judge versions + # that don't include the field. + judge_meta = ( + getattr(judge_result, "metadata", None) or {} + ) if hasattr(judge_result, "metadata") else {} + if not isinstance(judge_meta, dict): + judge_meta = {} + judge_cost_usd = float( + judge_meta.get("cost_usd") + or judge_meta.get("usd_cost") + or 0.0 + ) + # Mode-specific completion-time SLA: any *single turn* + # exceeding the budget counts as a violation. We gate on + # the per-turn max so a fixture that has 3 well-behaved + # turns + 1 over-budget turn is still flagged. + if latency_budget is not None and max_turn_latency > latency_budget: logger.info( "latency_violation miner=%s task_eval=%s mode=%s " - "completion=%.2fs budget=%.2fs prior_verdict=%s", + "max_turn=%.2fs total=%.2fs budget=%.2fs prior_verdict=%s", miner_hotkey[:16], task_evaluation_id, task_mode, - miner_latency, latency_budget, verdict, + max_turn_latency, miner_latency, latency_budget, verdict, ) verdict = "latency_violation" agreement_score = 0.0 @@ -801,6 +1061,8 @@ async def _judge_miner(miner_run: tuple[str, BenchmarkTaskRun]) -> dict[str, Any "verdict": verdict, "miner_latency_seconds": miner_latency, "latency_seconds": judge_latency, + "proxy_cost_usd": proxy_cost_usd, + "judge_cost_usd": judge_cost_usd, } except Exception as exc: logger.warning( @@ -816,6 +1078,8 @@ async def _judge_miner(miner_run: tuple[str, BenchmarkTaskRun]) -> dict[str, Any "verdict": "error", "miner_latency_seconds": miner_latency, "latency_seconds": 0.0, + "proxy_cost_usd": proxy_cost_usd, + "judge_cost_usd": 0.0, } # Materialize the gather results (handle exceptions from miner fan-out) @@ -835,6 +1099,41 @@ async def _judge_miner(miner_run: tuple[str, BenchmarkTaskRun]) -> dict[str, Any judge_results = await asyncio.gather(*[_judge_miner(r) for r in resolved_runs]) + # Verdict tally — one-line summary so an operator scrolling the + # validator log can see "task X scored 2/3 matches" at a glance, + # without having to query the DB. + verdict_counts: dict[str, int] = {} + for r in judge_results: + v = str(r.get("verdict") or "error") + verdict_counts[v] = verdict_counts.get(v, 0) + 1 + verdict_summary = " ".join( + f"{k}={v}" for k, v in sorted(verdict_counts.items()) + ) + logger.info( + "task_eval=%s verdicts: %s", + task_evaluation_id[:8], verdict_summary or "(none)", + ) + + # Per-task cost roll-up. Three components, all server-side + # sourced: (1) baseline = OpenAI Responses API spend for this + # task (sums all turns); (2) miners = sum of provider-proxy + # ledger lookups for each miner under this task's job_id; + # (3) judge = sum of eiretes-judge-reported per-judgment + # ``cost_usd`` across the miners we judged for this task. + baseline_cost = float(getattr(baseline, "cost_usd", 0.0) or 0.0) + miners_cost = sum( + float(r.get("proxy_cost_usd") or 0.0) for r in judge_results + ) + judge_cost = sum( + float(r.get("judge_cost_usd") or 0.0) for r in judge_results + ) + logger.info( + "task_eval=%s cost: baseline=$%.4f miners=$%.4f judge=$%.4f total=$%.4f", + task_evaluation_id[:8], + baseline_cost, miners_cost, judge_cost, + baseline_cost + miners_cost + judge_cost, + ) + # Submit the batch result_path = f"/v1/families/{family_id}/task-evaluations/{task_evaluation_id}/result" result_body = { @@ -855,6 +1154,11 @@ async def _judge_miner(miner_run: tuple[str, BenchmarkTaskRun]) -> dict[str, Any }, ) resp.raise_for_status() + elapsed = time.perf_counter() - task_started_at + logger.info( + "task_eval=%s SUBMITTED total=%.2fs", + task_evaluation_id[:8], elapsed, + ) return "ok" except Exception as exc: logger.warning("submit failed for task_eval=%s: %s", task_evaluation_id, exc) diff --git a/validation/validator/openai_baseline.py b/validation/validator/openai_baseline.py index e498baa..8395d38 100644 --- a/validation/validator/openai_baseline.py +++ b/validation/validator/openai_baseline.py @@ -90,6 +90,7 @@ async def generate( *, prompt: str, use_web_search: bool = False, + history: list[dict[str, Any]] | None = None, ) -> BaselineResponse: """Call the Responses API, optionally with the built-in web_search tool. @@ -97,6 +98,14 @@ async def generate( the caller — this mirrors the end-user toggle, so the baseline has the same information access as the miner would in real chat usage. + ``history`` is the list of prior conversation turns + (``[{role, content}, ...]``). For multi-turn fixtures the + validator passes history accumulated up to the final live turn; + for single-turn tasks it's empty. The Responses API accepts a + list ``input`` of role-tagged messages, so we send the full + conversation in chronological order with the latest user + prompt appended at the end. + Returns a normalized BaselineResponse; raises ``OpenAIBaselineError`` on any failure (network, HTTP error, malformed output, budget exceeded). @@ -109,9 +118,24 @@ async def generate( f"(spent={self._spent_usd:.4f} cap={self.max_cost_usd_per_run:.4f})" ) + # Single-turn → string input (cheaper to serialise + matches the + # prior shape). Multi-turn → role-tagged list with history + the + # current user prompt as the last message. + clean_history = [ + {"role": h.get("role"), "content": h.get("content")} + for h in (history or []) + if isinstance(h, dict) and h.get("role") in ("user", "assistant") + ] + if clean_history: + input_payload: list[dict[str, Any]] | str = ( + clean_history + [{"role": "user", "content": prompt}] + ) + else: + input_payload = prompt + payload: dict[str, Any] = { "model": self.model, - "input": prompt, + "input": input_payload, } if use_web_search: payload["tools"] = [{"type": "web_search"}] From 91f5af365cd77459a913473a8d7e9e932ad07d84 Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:52:50 +0000 Subject: [PATCH 04/24] feat(dashboard): queued submissions page --- control_plane/owner_api/routers/dashboard.py | 21 ++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/control_plane/owner_api/routers/dashboard.py b/control_plane/owner_api/routers/dashboard.py index e1f9de3..b4406a9 100644 --- a/control_plane/owner_api/routers/dashboard.py +++ b/control_plane/owner_api/routers/dashboard.py @@ -11,6 +11,7 @@ MinerProfileResponse, MinerRunsResponse, OverviewResponse, + QueuedSubmissionsResponse, RunDetailResponse, RunListResponse, ) @@ -118,6 +119,26 @@ async def get_leaderboard( return result +@router.get("/submissions/queued", response_model=QueuedSubmissionsResponse) +async def get_queued_submissions( + request: Request, + limit: int = Query(default=200, ge=1, le=500), +) -> QueuedSubmissionsResponse: + services: ManagedOwnerServices = request.app.state.services + key = ("queued_submissions", limit) + cached = _CACHE.get(key) + if cached is not None: + return cached # type: ignore[return-value] + with services.db.sessionmaker() as session: + result = queries.fetch_queued_submissions( + session, services=services, limit=limit, + ) + # Submissions move through states quickly enough that a 5s TTL keeps + # the queue page feeling live without hammering the DB. + _CACHE.set(key, result, ttl=_OPEN_RUN_TTL) + return result + + @router.get("/miners/{hotkey}", response_model=MinerProfileResponse) async def get_miner_profile( request: Request, From 32a13d5bb25b798dcd4e812f8068a020ffce1948 Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:54:17 +0000 Subject: [PATCH 05/24] feat(cost): per-task LLM cost via owner-api ledger injection --- .../evaluation/evaluation_task_manager.py | 7 + control_plane/owner_api/routers/runtime.py | 128 +++++++++++++++++- shared/common/models.py | 10 ++ 3 files changed, 140 insertions(+), 5 deletions(-) diff --git a/control_plane/owner_api/evaluation/evaluation_task_manager.py b/control_plane/owner_api/evaluation/evaluation_task_manager.py index a2fbd24..97e9ae5 100644 --- a/control_plane/owner_api/evaluation/evaluation_task_manager.py +++ b/control_plane/owner_api/evaluation/evaluation_task_manager.py @@ -343,6 +343,13 @@ def submit_task_result( agreement_score=float(score), miner_latency_seconds=float(entry.get("miner_latency_seconds") or 0.0), latency_seconds=float(entry.get("latency_seconds") or 0.0), + # Costs are server-side ground truth: ``proxy_cost_usd`` + # came from owner-api's provider-proxy ledger lookup and + # was injected into the miner's done-chunk metadata; + # ``judge_cost_usd`` came from eiretes-judge's response + # metadata. Validator passes both through verbatim. + proxy_cost_usd=float(entry.get("proxy_cost_usd") or 0.0), + judge_cost_usd=float(entry.get("judge_cost_usd") or 0.0), )) session.flush() diff --git a/control_plane/owner_api/routers/runtime.py b/control_plane/owner_api/routers/runtime.py index 7427f41..05d199a 100644 --- a/control_plane/owner_api/routers/runtime.py +++ b/control_plane/owner_api/routers/runtime.py @@ -1,6 +1,8 @@ from __future__ import annotations +import json as _json import logging +import os from collections.abc import AsyncIterator from typing import Any @@ -30,6 +32,65 @@ # own timeout fires. _STREAM_TIMEOUT_SECONDS = 660.0 +# Provider-proxy URL used for cost reconciliation after a miner stream +# completes. Owner-api stamps a per-task ``X-Eirel-Job-Id`` header on +# the way to the miner; the miner SDK forwards that as the proxy +# job_id; provider-proxy ledgers cost per job_id; owner-api queries +# this URL after the stream closes and injects the result into the +# final ``done`` chunk's metadata so the validator sees per-task +# proxy_cost_usd directly. +_PROVIDER_PROXY_URL = os.getenv( + "EIREL_PROVIDER_PROXY_URL", "http://provider-proxy:8092" +).rstrip("/") +_PROVIDER_PROXY_TOKEN = os.getenv("EIREL_PROVIDER_PROXY_TOKEN", "") + + +def _build_cost_tag(*, deployment_id: str, payload: dict[str, Any]) -> str: + """Stable per-task tag used as the provider-proxy job_id. + + Validator-issued requests carry ``turn_id`` (slim 0.3.0 contract); + consumer-chat-api requests do too. Falls back to ``task_id`` + (legacy 0.2.x) for older callers, then to deployment_id for any + request that has neither (e.g. orchestrator-internal probes — + those keep the deployment-sticky semantics they had before). + """ + turn_id = payload.get("turn_id") or payload.get("task_id") or "" + turn_id = str(turn_id).strip() + if turn_id: + return f"task-eval={turn_id};deployment={deployment_id}" + return f"miner-{deployment_id}" + + +async def _fetch_proxy_cost(job_id: str) -> dict[str, Any] | None: + """Best-effort cost lookup against provider-proxy. + + Returns None on any failure — cost reporting must never break the + eval pipeline. Caller renders a missing dict as "cost unknown". + """ + if not job_id: + return None + url = f"{_PROVIDER_PROXY_URL}/v1/jobs/{job_id}/cost" + headers = ( + {"Authorization": f"Bearer {_PROVIDER_PROXY_TOKEN}"} + if _PROVIDER_PROXY_TOKEN + else {} + ) + try: + async with httpx.AsyncClient(timeout=5.0) as client: + resp = await client.get(url, headers=headers) + if resp.status_code == 404: + # No ledger entry — miner didn't make any LLM calls + # under that job_id (or didn't forward the header). + return {"llm_cost_usd": 0.0, "tool_cost_usd": 0.0, "absent": True} + resp.raise_for_status() + return resp.json() + except Exception as exc: + logger.warning( + "proxy_cost lookup failed for job=%s url=%s: %s", + job_id, url, exc, + ) + return None + @router.get("/runtime/{deployment_id}/healthz") async def runtime_health( @@ -422,22 +483,44 @@ async def _proxy_stream_lines_to_pod( *, pod_endpoint: str, payload: dict[str, Any], + cost_tag: str | None = None, ) -> AsyncIterator[bytes]: """Forward the pod's NDJSON line-by-line. Each yielded chunk is one complete NDJSON frame ending with `\\n`. - Failures emerge as a synthetic `done` chunk so callers always see a - terminator they can act on. + Failures emerge as a synthetic ``done`` chunk so callers always see + a terminator they can act on. + + When ``cost_tag`` is non-empty, owner-api injects it as + ``X-Eirel-Job-Id`` so the miner SDK forwards it to the + provider-proxy as the LLM-call attribution key. After the miner's + terminal ``done`` arrives, we look up the cost ledger for that + tag and merge ``{proxy_cost_usd, proxy_request_count, ...}`` into + the chunk's ``metadata`` before re-emitting. Cost data is the + server-side ground truth; miner self-report is never trusted. """ url = f"{pod_endpoint.rstrip('/')}/v1/agent/infer/stream" + headers = {"X-Eirel-Job-Id": cost_tag} if cost_tag else {} + pending_done: dict[str, Any] | None = None try: async with httpx.AsyncClient(timeout=_STREAM_TIMEOUT_SECONDS) as client: - async with client.stream("POST", url, json=payload) as resp: + async with client.stream("POST", url, json=payload, headers=headers) as resp: resp.raise_for_status() async for line in resp.aiter_lines(): stripped = line.strip() if not stripped: continue + # Hold the terminal ``done`` chunk so we can + # augment it with cost metadata before forwarding. + # All non-done chunks pass through verbatim. + try: + parsed = _json.loads(stripped) + except _json.JSONDecodeError: + yield (stripped + "\n").encode("utf-8") + continue + if isinstance(parsed, dict) and parsed.get("event") == "done": + pending_done = parsed + continue yield (stripped + "\n").encode("utf-8") except httpx.HTTPStatusError as exc: logger.warning("stream pod returned %d for %s", exc.response.status_code, url) @@ -445,10 +528,33 @@ async def _proxy_stream_lines_to_pod( ('{"event":"done","status":"failed","error":"pod returned ' + str(exc.response.status_code) + '"}\n').encode("utf-8") ) + return except Exception as exc: # noqa: BLE001 logger.warning("stream pod connect failed for %s: %s", url, exc) msg = str(exc).replace('"', '\\"') yield ('{"event":"done","status":"failed","error":"' + msg + '"}\n').encode("utf-8") + return + + # Stream closed cleanly. Reconcile cost from the ledger, augment + # ``done`` metadata, emit. If the miner never sent ``done`` (rare, + # malformed stream) synthesize one so the validator's wire contract + # is preserved. + if pending_done is None: + pending_done = {"event": "done", "status": "completed"} + if cost_tag: + cost = await _fetch_proxy_cost(cost_tag) + if cost is not None: + meta = pending_done.get("metadata") + if not isinstance(meta, dict): + meta = {} + meta["proxy_cost_tag"] = cost_tag + meta["proxy_cost_usd"] = round(float(cost.get("llm_cost_usd") or 0.0), 8) + meta["proxy_tool_cost_usd"] = round( + float(cost.get("tool_cost_usd") or 0.0), 8, + ) + meta["proxy_cost_absent"] = bool(cost.get("absent", False)) + pending_done["metadata"] = meta + yield (_json.dumps(pending_done, separators=(",", ":")) + "\n").encode("utf-8") async def _eval_stream_impl( @@ -479,8 +585,13 @@ async def _eval_stream_impl( handle = services.runtime_manager.runtime_handle(deployment_id) if handle is None: raise HTTPException(status_code=503, detail="runtime not ready") + cost_tag = _build_cost_tag(deployment_id=deployment_id, payload=payload) return StreamingResponse( - _proxy_stream_lines_to_pod(pod_endpoint=handle.endpoint_url, payload=payload), + _proxy_stream_lines_to_pod( + pod_endpoint=handle.endpoint_url, + payload=payload, + cost_tag=cost_tag, + ), media_type="application/x-ndjson", headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"}, ) @@ -558,8 +669,15 @@ async def serving_runtime_infer_stream( handle = services.runtime_manager.runtime_handle(serving_deployment_id) if handle is None: raise HTTPException(status_code=503, detail="serving runtime not ready") + cost_tag = _build_cost_tag( + deployment_id=serving_deployment_id, payload=payload, + ) return StreamingResponse( - _proxy_stream_lines_to_pod(pod_endpoint=handle.endpoint_url, payload=payload), + _proxy_stream_lines_to_pod( + pod_endpoint=handle.endpoint_url, + payload=payload, + cost_tag=cost_tag, + ), media_type="application/x-ndjson", headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"}, ) diff --git a/shared/common/models.py b/shared/common/models.py index 937c7ac..cc1b48a 100644 --- a/shared/common/models.py +++ b/shared/common/models.py @@ -717,6 +717,16 @@ class TaskMinerResult(Base): # historically named `latency_seconds`; the attr name is preserved for # back-compat with existing seed/migration code. latency_seconds: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + # Per-task LLM cost the miner incurred against the subnet + # provider-proxy, in USD. Sourced server-side from the + # provider-proxy ledger via owner-api injection (Phase 2a) — never + # trusts miner self-report. Zero when the miner made no proxied + # LLM calls or when the cost lookup failed. + proxy_cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + # Per-task judge LLM cost in USD (Phase 2c). Reported by + # eiretes-judge in its response metadata; the validator passes it + # through verbatim. Always >= 0. + judge_cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=False), default=utcnow, nullable=False) From 98e4ff32a208c27e4bb23287ab42c4063551961c Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:55:03 +0000 Subject: [PATCH 06/24] fix(validator): burn 1.0 - total_assigned so EIREL_FAMILY_WEIGHTS controls real burn --- deploy/k8s/owner-api-hybrid/configmap.env.example | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/deploy/k8s/owner-api-hybrid/configmap.env.example b/deploy/k8s/owner-api-hybrid/configmap.env.example index 266bb22..03bf9f0 100644 --- a/deploy/k8s/owner-api-hybrid/configmap.env.example +++ b/deploy/k8s/owner-api-hybrid/configmap.env.example @@ -32,6 +32,21 @@ EIREL_JUDGE_MODEL=moonshotai/Kimi-K2.6-TEE EIREL_JUDGE_BASE_URL=https://llm.chutes.ai/v1 EIREL_JUDGE_TIMEOUT_SECONDS=300 +# -- Family weights / burn ---------------------------------------------------- +# Comma-separated ``family:weight`` pairs. The remaining ``1.0 - sum(weights)`` +# is automatically burned (sent to UID 0) by the validator on every +# weight-set call, so weights here express "fraction of subnet emission +# this family wins" directly. +# +# Examples: +# general_chat:0.5 → 50% to general_chat winner, 50% burn +# general_chat:0.3,coding:0.4 → 30% gc, 40% coding, 30% burn +# general_chat:1.0 → 100% to general_chat (no burn) +# +# A family with no winner this run still has its weight burned — +# emission never goes to a phantom recipient. +EIREL_FAMILY_WEIGHTS=general_chat:0.5 + # -- Run cadence -------------------------------------------------------------- # ISO-8601 UTC timestamp for when Run 1 opens. Leave blank to start immediately # at owner-api boot. Set to ~15 minutes after deploy to allow miner submissions From dc8dd761b9762faa87ef38bd19b98d9216b0e75b Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:55:43 +0000 Subject: [PATCH 07/24] chore(fee_verifier): boot-time audit log + per-call DEBUG --- control_plane/owner_api/app.py | 15 +++++++++++++++ control_plane/owner_api/fee_verifier.py | 8 ++++++++ 2 files changed, 23 insertions(+) diff --git a/control_plane/owner_api/app.py b/control_plane/owner_api/app.py index 1df0588..1596a57 100644 --- a/control_plane/owner_api/app.py +++ b/control_plane/owner_api/app.py @@ -369,6 +369,21 @@ async def lifespan(app: FastAPI): treasury_address=settings.submission_treasury_address, fee_tao=settings.submission_fee_tao, ) + # One-shot audit log so operators have a record of which chain + # the fee verifier is bound to. A miner who pays on a different + # network will hit "extrinsic not found on chain" — knowing the + # operator's bound network from this line lets them debug + # without having to read source. Goes through ``print`` (and is + # also emitted via ``logger`` for environments that capture it) + # because uvicorn's lifespan-time logger reconfig sometimes + # eats lifespan-emitted log records before any handler attaches. + _audit_msg = ( + f"fee_verifier: bound to network={settings.bittensor_network} " + f"treasury={settings.submission_treasury_address} " + f"fee_tao={settings.submission_fee_tao:.4f}" + ) + print(_audit_msg, flush=True) + logger.warning(_audit_msg) app.state.execution_worker_client_factory = None app.state.runtime_remediation_policy_state = { "enabled": bool(settings.workflow_runtime_auto_remediation_enabled), diff --git a/control_plane/owner_api/fee_verifier.py b/control_plane/owner_api/fee_verifier.py index 39daba3..2da2b6a 100644 --- a/control_plane/owner_api/fee_verifier.py +++ b/control_plane/owner_api/fee_verifier.py @@ -74,6 +74,14 @@ def verify_payment( Returns a ``FeeVerificationResult`` indicating success or failure with a human-readable reason. """ + # Per-call DEBUG audit. Set ``LOG_LEVEL=DEBUG`` to see this; INFO + # operators can grep for the boot-time ``fee_verifier: bound to`` + # line to learn the bound network. + logger.debug( + "verifying extrinsic %s on %s chain (treasury=%s, fee_rao=%d)", + extrinsic_hash, self.network, + self.treasury_address, self.fee_rao, + ) # Resolve the coldkey that owns this hotkey (cached). expected_coldkey = self._get_hotkey_owner_cached(expected_hotkey) if not expected_coldkey: From f422c7fe9f3863d8bb47790d6d933afeb1b571eb Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:57:39 +0000 Subject: [PATCH 08/24] feat(monitoring): add kube-state-metrics and node-exporter configurations --- .../kube-state-metrics/clusterrole.yaml | 64 +++++++++++++ .../clusterrolebinding.yaml | 12 +++ monitoring/kube-state-metrics/deployment.yaml | 51 ++++++++++ .../kube-state-metrics/kustomization.yaml | 12 +++ monitoring/kube-state-metrics/namespace.yaml | 6 ++ monitoring/kube-state-metrics/service.yaml | 19 ++++ .../kube-state-metrics/serviceaccount.yaml | 4 + monitoring/node-exporter/daemonset.yaml | 92 +++++++++++++++++++ monitoring/node-exporter/kustomization.yaml | 7 ++ 9 files changed, 267 insertions(+) create mode 100644 monitoring/kube-state-metrics/clusterrole.yaml create mode 100644 monitoring/kube-state-metrics/clusterrolebinding.yaml create mode 100644 monitoring/kube-state-metrics/deployment.yaml create mode 100644 monitoring/kube-state-metrics/kustomization.yaml create mode 100644 monitoring/kube-state-metrics/namespace.yaml create mode 100644 monitoring/kube-state-metrics/service.yaml create mode 100644 monitoring/kube-state-metrics/serviceaccount.yaml create mode 100644 monitoring/node-exporter/daemonset.yaml create mode 100644 monitoring/node-exporter/kustomization.yaml diff --git a/monitoring/kube-state-metrics/clusterrole.yaml b/monitoring/kube-state-metrics/clusterrole.yaml new file mode 100644 index 0000000..7990bde --- /dev/null +++ b/monitoring/kube-state-metrics/clusterrole.yaml @@ -0,0 +1,64 @@ +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: kube-state-metrics +rules: + - apiGroups: [""] + resources: + - namespaces + - nodes + - persistentvolumeclaims + - persistentvolumes + - pods + - secrets + - services + - resourcequotas + - replicationcontrollers + - limitranges + - endpoints + - configmaps + - serviceaccounts + verbs: ["list", "watch"] + - apiGroups: ["apps"] + resources: + - deployments + - daemonsets + - replicasets + - statefulsets + verbs: ["list", "watch"] + - apiGroups: ["batch"] + resources: + - cronjobs + - jobs + verbs: ["list", "watch"] + - apiGroups: ["autoscaling"] + resources: + - horizontalpodautoscalers + verbs: ["list", "watch"] + - apiGroups: ["networking.k8s.io"] + resources: + - ingresses + - networkpolicies + verbs: ["list", "watch"] + - apiGroups: ["policy"] + resources: + - poddisruptionbudgets + verbs: ["list", "watch"] + - apiGroups: ["storage.k8s.io"] + resources: + - storageclasses + - volumeattachments + verbs: ["list", "watch"] + - apiGroups: ["certificates.k8s.io"] + resources: + - certificatesigningrequests + verbs: ["list", "watch"] + - apiGroups: ["admissionregistration.k8s.io"] + resources: + - mutatingwebhookconfigurations + - validatingwebhookconfigurations + verbs: ["list", "watch"] + - apiGroups: ["coordination.k8s.io"] + resources: + - leases + verbs: ["list", "watch"] diff --git a/monitoring/kube-state-metrics/clusterrolebinding.yaml b/monitoring/kube-state-metrics/clusterrolebinding.yaml new file mode 100644 index 0000000..879f9e2 --- /dev/null +++ b/monitoring/kube-state-metrics/clusterrolebinding.yaml @@ -0,0 +1,12 @@ +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: kube-state-metrics +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: kube-state-metrics +subjects: + - kind: ServiceAccount + name: kube-state-metrics + namespace: eirel-monitoring diff --git a/monitoring/kube-state-metrics/deployment.yaml b/monitoring/kube-state-metrics/deployment.yaml new file mode 100644 index 0000000..b537417 --- /dev/null +++ b/monitoring/kube-state-metrics/deployment.yaml @@ -0,0 +1,51 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: kube-state-metrics +spec: + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/name: kube-state-metrics + template: + metadata: + labels: + app.kubernetes.io/name: kube-state-metrics + spec: + serviceAccountName: kube-state-metrics + automountServiceAccountToken: true + containers: + - name: kube-state-metrics + image: registry.k8s.io/kube-state-metrics/kube-state-metrics:v2.13.0 + args: + - --port=8080 + - --telemetry-port=8081 + ports: + - name: http + containerPort: 8080 + - name: telemetry + containerPort: 8081 + readinessProbe: + httpGet: + path: /healthz + port: 8080 + initialDelaySeconds: 5 + periodSeconds: 10 + livenessProbe: + httpGet: + path: /healthz + port: 8080 + initialDelaySeconds: 15 + periodSeconds: 20 + resources: + requests: + cpu: 50m + memory: 64Mi + limits: + cpu: 200m + memory: 256Mi + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + runAsNonRoot: true + runAsUser: 65534 diff --git a/monitoring/kube-state-metrics/kustomization.yaml b/monitoring/kube-state-metrics/kustomization.yaml new file mode 100644 index 0000000..4cfbc9b --- /dev/null +++ b/monitoring/kube-state-metrics/kustomization.yaml @@ -0,0 +1,12 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +namespace: eirel-monitoring + +resources: + - namespace.yaml + - serviceaccount.yaml + - clusterrole.yaml + - clusterrolebinding.yaml + - deployment.yaml + - service.yaml diff --git a/monitoring/kube-state-metrics/namespace.yaml b/monitoring/kube-state-metrics/namespace.yaml new file mode 100644 index 0000000..0769582 --- /dev/null +++ b/monitoring/kube-state-metrics/namespace.yaml @@ -0,0 +1,6 @@ +apiVersion: v1 +kind: Namespace +metadata: + name: eirel-monitoring + labels: + name: eirel-monitoring diff --git a/monitoring/kube-state-metrics/service.yaml b/monitoring/kube-state-metrics/service.yaml new file mode 100644 index 0000000..07ed080 --- /dev/null +++ b/monitoring/kube-state-metrics/service.yaml @@ -0,0 +1,19 @@ +apiVersion: v1 +kind: Service +metadata: + name: kube-state-metrics + labels: + app.kubernetes.io/name: kube-state-metrics +spec: + # NodePort 30080 matches the host.docker.internal:30080 scrape target + # in monitoring/prometheus/prometheus.yml. The compose-hosted prometheus + # reaches it through host-gateway → node host IP → NodePort. + type: NodePort + selector: + app.kubernetes.io/name: kube-state-metrics + ports: + - name: http + port: 8080 + targetPort: 8080 + nodePort: 30080 + protocol: TCP diff --git a/monitoring/kube-state-metrics/serviceaccount.yaml b/monitoring/kube-state-metrics/serviceaccount.yaml new file mode 100644 index 0000000..9977935 --- /dev/null +++ b/monitoring/kube-state-metrics/serviceaccount.yaml @@ -0,0 +1,4 @@ +apiVersion: v1 +kind: ServiceAccount +metadata: + name: kube-state-metrics diff --git a/monitoring/node-exporter/daemonset.yaml b/monitoring/node-exporter/daemonset.yaml new file mode 100644 index 0000000..9cfc549 --- /dev/null +++ b/monitoring/node-exporter/daemonset.yaml @@ -0,0 +1,92 @@ +apiVersion: apps/v1 +kind: DaemonSet +metadata: + name: node-exporter + labels: + app.kubernetes.io/name: node-exporter +spec: + selector: + matchLabels: + app.kubernetes.io/name: node-exporter + template: + metadata: + labels: + app.kubernetes.io/name: node-exporter + spec: + # Pin to the miner-pool only — these are the nodes the runtime_node + # job in prometheus.yml expects to scrape. Adding the daemonset + # cluster-wide would also scrape the control-plane and confuse + # alert thresholds tuned for miner workloads. + nodeSelector: + eirel.dev/runtime-pool: "true" + # Bind directly to the host network so port 9100 lands on the + # node's IP — that's what the prometheus runtime_node service + # discovery emits and scrapes. + hostNetwork: true + hostPID: true + # Tolerate any common node taints so the daemonset still lands + # on tainted miner nodes (e.g. NoSchedule for runtime-class). + tolerations: + - operator: Exists + containers: + - name: node-exporter + image: quay.io/prometheus/node-exporter:v1.8.2 + args: + - --path.procfs=/host/proc + - --path.sysfs=/host/sys + - --path.rootfs=/host/root + - --collector.filesystem.mount-points-exclude=^/(dev|proc|sys|var/lib/docker/.+|var/lib/kubelet/pods/.+)($|/) + - --collector.filesystem.fs-types-exclude=^(autofs|binfmt_misc|bpf|cgroup2?|configfs|debugfs|devpts|devtmpfs|fusectl|hugetlbfs|iso9660|mqueue|nsfs|overlay|proc|procfs|pstore|rpc_pipefs|securityfs|selinuxfs|squashfs|sysfs|tracefs)$ + - --web.listen-address=0.0.0.0:9100 + ports: + - name: metrics + containerPort: 9100 + hostPort: 9100 + protocol: TCP + readinessProbe: + httpGet: + path: / + port: 9100 + initialDelaySeconds: 5 + periodSeconds: 30 + livenessProbe: + httpGet: + path: / + port: 9100 + initialDelaySeconds: 15 + periodSeconds: 30 + resources: + requests: + cpu: 25m + memory: 32Mi + limits: + cpu: 250m + memory: 128Mi + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + runAsNonRoot: true + runAsUser: 65534 + capabilities: + drop: ["ALL"] + volumeMounts: + - name: proc + mountPath: /host/proc + readOnly: true + - name: sys + mountPath: /host/sys + readOnly: true + - name: root + mountPath: /host/root + mountPropagation: HostToContainer + readOnly: true + volumes: + - name: proc + hostPath: + path: /proc + - name: sys + hostPath: + path: /sys + - name: root + hostPath: + path: / diff --git a/monitoring/node-exporter/kustomization.yaml b/monitoring/node-exporter/kustomization.yaml new file mode 100644 index 0000000..fe9083b --- /dev/null +++ b/monitoring/node-exporter/kustomization.yaml @@ -0,0 +1,7 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +namespace: eirel-monitoring + +resources: + - daemonset.yaml From ab021bf42cb2e0f88a25dbc3cefad620bc968d54 Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Mon, 4 May 2026 12:15:51 +0000 Subject: [PATCH 09/24] refactor: remove honeytoken and trace gate penalty features from the codebase --- .../owner_api/evaluation/run_manager.py | 8 +- .../owner_api/evaluation/scoring_manager.py | 48 ------ .../owner_api/middleware/trace_capture.py | 26 +-- docs/miner-guide.md | 4 - docs/validator-guide.md | 8 +- shared/common/config.py | 12 -- shared/common/migrations.py | 63 +++++++ shared/core/evaluation_models.py | 9 - shared/core/honeytokens.py | 162 ------------------ .../families/_judge_to_conversation_score.py | 32 +--- .../owner_api/test_scoring_manager_penalty.py | 84 --------- .../test_provider_proxy_usd_budget.py | 72 +------- tool_platforms/provider_proxy/app.py | 41 +---- tool_platforms/provider_proxy/redis_store.py | 23 +-- 14 files changed, 82 insertions(+), 510 deletions(-) delete mode 100644 shared/core/honeytokens.py delete mode 100644 tests/owner_api/test_scoring_manager_penalty.py diff --git a/control_plane/owner_api/evaluation/run_manager.py b/control_plane/owner_api/evaluation/run_manager.py index 149d370..d8fd9ea 100644 --- a/control_plane/owner_api/evaluation/run_manager.py +++ b/control_plane/owner_api/evaluation/run_manager.py @@ -42,7 +42,6 @@ from shared.contracts.models import BenchmarkRunRecord from shared.contracts.specialist_contracts import contract_for_family from shared.core.evaluation_models import FamilyEvaluationBundle -from shared.core.honeytokens import generate_honeytoken_set from eirel.groups import ensure_family_id if TYPE_CHECKING: @@ -138,7 +137,6 @@ def create_run( status: str = "open", ) -> EvaluationRun: started = started_at or _utcnow() - honeytoken_count = int(getattr(self.settings, "honeytoken_count_per_run", 8)) run = EvaluationRun( id=f"run-{sequence}", sequence=sequence, @@ -149,11 +147,7 @@ def create_run( min_scores_json=self._owner.run_min_scores(), started_at=started, ends_at=started + timedelta(days=self._owner.run_duration_days), - metadata_json={ - "honeytokens": generate_honeytoken_set( - f"run-{sequence}", count=honeytoken_count, - ), - }, + metadata_json={}, ) session.add(run) session.flush() diff --git a/control_plane/owner_api/evaluation/scoring_manager.py b/control_plane/owner_api/evaluation/scoring_manager.py index 3a03fc7..9dfaede 100644 --- a/control_plane/owner_api/evaluation/scoring_manager.py +++ b/control_plane/owner_api/evaluation/scoring_manager.py @@ -956,54 +956,6 @@ def fetch_deployment_cost(self, deployment_id: str) -> dict[str, Any]: logger.warning("failed to fetch cost for deployment %s: %s", deployment_id, exc) return {} - def charge_trace_gate_penalty( - self, - deployment_id: str, - *, - amount_usd: float, - reason: str = "trace_gate_fail", - ) -> bool: - """Debit a USD penalty against a deployment's run budget. - - Called once per conversation that fails the trace integrity gate. - Returns True on success, False on any failure (logged — we don't - want penalty charging to break the scoring pipeline). The penalty - lands unconditionally, even if it overshoots ``max_usd_budget`` — - that's the whole point of the economic gate. - """ - if amount_usd <= 0.0: - return False - proxy_url = self.settings.provider_proxy_url - if not proxy_url: - logger.warning( - "provider_proxy_url not configured; cannot charge trace-gate penalty", - ) - return False - url = f"{proxy_url.rstrip('/')}/v1/jobs/miner-{deployment_id}/charge_penalty" - headers: dict[str, str] = {} - token = getattr(self.settings, "provider_proxy_token", None) - if token: - headers["Authorization"] = f"Bearer {token}" - try: - resp = httpx.post( - url, - json={"reason": reason, "amount_usd": float(amount_usd)}, - headers=headers, - timeout=10.0, - ) - resp.raise_for_status() - logger.info( - "charged trace-gate penalty: deployment=%s amount_usd=%.4f reason=%s", - deployment_id, amount_usd, reason, - ) - return True - except Exception as exc: - logger.warning( - "failed to charge trace-gate penalty for deployment %s: %s", - deployment_id, exc, - ) - return False - def populate_cost_columns( self, record: DeploymentScoreRecord, diff --git a/control_plane/owner_api/middleware/trace_capture.py b/control_plane/owner_api/middleware/trace_capture.py index 719b7ef..1a2cf67 100644 --- a/control_plane/owner_api/middleware/trace_capture.py +++ b/control_plane/owner_api/middleware/trace_capture.py @@ -11,7 +11,6 @@ from shared.common.tool_pricing import cost_for_call from shared.core.evaluation_models import ConversationTrace, TraceEntry -from shared.core.honeytokens import inject_honeytokens_into_search_payload _logger = logging.getLogger(__name__) @@ -86,9 +85,8 @@ def _digest_result(payload: Any) -> str: def _body_excerpt(payload: Any) -> str: """Return a lowercased, length-capped excerpt of a tool response body. - Used by ``verify_trace_integrity`` to do substring / token-overlap - checks against cited content. Structure-preserving JSON serialization - keeps field names available for matching. + Stored on each :class:`TraceEntry` for audit / debugging purposes. + Structure-preserving JSON serialization keeps field names readable. """ if payload is None: return "" @@ -133,8 +131,6 @@ def __init__( provider_proxy_url: str | None = None, provider_proxy_token: str | None = None, job_id: str | None = None, - active_honeytokens: list[str] | None = None, - honeytoken_injection_rate: float = 0.02, ) -> None: self._store = store self._client = http_client or httpx.AsyncClient(timeout=30.0) @@ -142,10 +138,6 @@ def __init__( self._provider_proxy_url = provider_proxy_url self._provider_proxy_token = provider_proxy_token self._job_id = job_id - self._active_honeytokens = list(active_honeytokens or []) - self._honeytoken_rate = honeytoken_injection_rate - # Per-conversation counter for deterministic injection keying. - self._call_counters: dict[str, int] = {} async def close(self) -> None: if self._owns_client: @@ -223,20 +215,6 @@ async def proxy_call( error_text = f"http_error:{type(exc).__name__}" latency_ms = int((time.perf_counter() - t0) * 1000) - # Honeytoken injection — only applies to search-style tools that - # return a list payload. Runs after the real fetch so the injected - # entry is visible to the miner alongside legitimate results. Rate - # is intentionally low (~2%) so honest miners rarely encounter one. - if self._active_honeytokens and error_text is None: - call_index = self._call_counters.get(conversation_id, 0) - self._call_counters[conversation_id] = call_index + 1 - payload = inject_honeytokens_into_search_payload( - payload, - active_set=self._active_honeytokens, - conversation_id=conversation_id, - call_index=call_index, - rate=self._honeytoken_rate, - ) entry = TraceEntry( tool_name=tool_name, args=dict(args), diff --git a/docs/miner-guide.md b/docs/miner-guide.md index 1e40f9d..edfceb7 100644 --- a/docs/miner-guide.md +++ b/docs/miner-guide.md @@ -117,10 +117,6 @@ with `--extrinsic-hash --block-hash ` instead of re-paying. your container directly — they hit the operator's `/runtime/{deployment_id}/v1/agent/infer` path, which routes to your pod. This means you don't need a public IP or axon. -- **Honeytokens + trace checks are hidden.** The operator injects - canary URLs into some evaluation tasks and checks your citations - against a hidden denylist. Citing a honeytoken as if it were a real - source gates out your submission. The full list is not published. - **Budget enforcement is at the proxy.** Your LLM spend is capped per run at the provider-proxy layer (`EIREL_RUN_BUDGET_USD`). Requests that would push you over-budget are rejected. diff --git a/docs/validator-guide.md b/docs/validator-guide.md index d709669..a5cf318 100644 --- a/docs/validator-guide.md +++ b/docs/validator-guide.md @@ -161,11 +161,9 @@ and directly reduces your TAO earnings. The field is unchanged from `EIREL_JUDGE_MODEL` in `.env.validator`, so it's just a matter of not changing it. -**Anti-gaming stays server-side.** The operator's owner-api applies -trace integrity checks, honeytoken detection, and latency axis on top -of your LLM quality score when you submit. You don't see the honeytoken -URL list, the trace-gate heuristics, or the latency penalty curve — -those live inside the operator process. +**Latency stays server-side.** The operator's owner-api applies the +latency axis on top of your LLM quality score when you submit. The +latency penalty curve lives inside the operator process. ## Costs diff --git a/shared/common/config.py b/shared/common/config.py index c3cc85a..cd5fb7d 100644 --- a/shared/common/config.py +++ b/shared/common/config.py @@ -579,15 +579,6 @@ class Settings: run_budget_usd: float = field( default_factory=lambda: float(os.getenv("EIREL_RUN_BUDGET_USD", "30.0")) ) - trace_gate_penalty_usd: float = field( - default_factory=lambda: float(os.getenv("EIREL_TRACE_GATE_PENALTY_USD", "0.50")) - ) - honeytoken_count_per_run: int = field( - default_factory=lambda: int(os.getenv("EIREL_HONEYTOKEN_COUNT_PER_RUN", "8")) - ) - honeytoken_injection_rate: float = field( - default_factory=lambda: float(os.getenv("EIREL_HONEYTOKEN_INJECTION_RATE", "0.02")) - ) trace_store_backend: str = field( default_factory=lambda: os.getenv("EIREL_TRACE_STORE_BACKEND", "memory") ) @@ -620,9 +611,6 @@ class Settings: first_run_start_time: str = field( default_factory=lambda: os.getenv("EIREL_FIRST_RUN_START_TIME", "") ) - spot_check_duplicate_rate: float = field( - default_factory=lambda: _float_env("EIREL_SPOT_CHECK_DUPLICATE_RATE", 0.05) - ) active_families: str = field( default_factory=lambda: os.getenv("EIREL_ACTIVE_FAMILIES", "general_chat") ) diff --git a/shared/common/migrations.py b/shared/common/migrations.py index ffb2a26..df167a4 100644 --- a/shared/common/migrations.py +++ b/shared/common/migrations.py @@ -391,6 +391,14 @@ def _migration_add_snapshot_unique_constraint(engine: Engine) -> None: ), apply=lambda engine: _migration_add_proxy_and_judge_cost(engine), ), + Migration( + version="drop_anti_gaming_artifacts", + description=( + "Strip honeytokens / trace_gate artifacts from JSON metadata " + "after the anti-gaming subsystem was removed" + ), + apply=lambda engine: _migration_drop_anti_gaming_artifacts(engine), + ), ) @@ -471,6 +479,61 @@ def _migration_add_proxy_and_judge_cost(engine: Engine) -> None: )) +def _migration_drop_anti_gaming_artifacts(engine: Engine) -> None: + """Strip anti-gaming residue from existing rows. + + The honeytoken / trace-gate code path was removed. No SQL columns + were ever added for it, but ``evaluation_runs.metadata_json`` carried + a ``"honeytokens"`` key with the per-run canary URL list, and + ``deployment_score_records`` (via Pydantic JSON serialization) may + have ``trace_gate``, ``trace_gate_penalty_usd``, or ``honeytoken_cited`` + keys embedded in stored conversation-score blobs. Strip both for + cleanliness so reads don't surface stale fields. + """ + inspector = inspect(engine) + tables = set(inspector.get_table_names()) + if "evaluation_runs" not in tables: + return + cols = {col["name"] for col in inspector.get_columns("evaluation_runs")} + if "metadata_json" not in cols: + return + dialect = engine.dialect.name + with engine.begin() as conn: + if dialect == "postgresql": + # JSONB-aware: strip the "honeytokens" key in place. + conn.execute(text( + "UPDATE evaluation_runs " + "SET metadata_json = metadata_json::jsonb - 'honeytokens' " + "WHERE metadata_json::jsonb ? 'honeytokens'" + )) + else: + # SQLite / others: rebuild the JSON without the key. Cheap + # and correct because run counts are small. + rows = conn.execute(text( + "SELECT id, metadata_json FROM evaluation_runs " + "WHERE metadata_json IS NOT NULL" + )).fetchall() + import json as _json + for row in rows: + raw = row[1] + if not raw: + continue + try: + parsed = _json.loads(raw) if isinstance(raw, str) else dict(raw) + except (ValueError, TypeError): + continue + if not isinstance(parsed, dict) or "honeytokens" not in parsed: + continue + parsed.pop("honeytokens", None) + conn.execute( + text( + "UPDATE evaluation_runs SET metadata_json = :m " + "WHERE id = :id" + ), + {"m": _json.dumps(parsed), "id": row[0]}, + ) + + def _drop_registered_neurons_is_active(engine: Engine) -> None: inspector = inspect(engine) if "registered_neurons" not in set(inspector.get_table_names()): diff --git a/shared/core/evaluation_models.py b/shared/core/evaluation_models.py index 2f3d3f7..df6b020 100644 --- a/shared/core/evaluation_models.py +++ b/shared/core/evaluation_models.py @@ -293,15 +293,10 @@ class ConversationScore(BaseModel): quality: float = Field(ge=0.0, le=1.0) latency: float = Field(ge=0.0, le=1.0) cost: float = Field(ge=0.0, le=1.0) - trace_gate: float = Field(ge=0.0, le=1.0) total: float = Field(ge=0.0, le=1.0) per_dimension: dict[str, float] = Field(default_factory=dict) mode: Literal["instant", "thinking"] = "instant" web_search: bool = False - # USD penalty the scoring manager should charge against the run budget - # when this conversation's trace integrity gate failed. 0.0 when the - # gate passed or when no penalty was configured. - trace_gate_penalty_usd: float = Field(ge=0.0, default=0.0) metadata: dict[str, Any] = Field(default_factory=dict) @@ -316,10 +311,6 @@ class MinerGeneralChatScore(BaseModel): llm_cost_usd: float = Field(ge=0.0, default=0.0) tool_cost_usd: float = Field(ge=0.0, default=0.0) cost_rejection_count: int = Field(ge=0, default=0) - # Bad-actor flag: True iff any conversation in this run cited an - # active honeytoken URL. When True, ``blended`` is forced to 0.0 - # regardless of quality — fabricated citations zero the miner. - honeytoken_cited: bool = False conversation_scores: list[ConversationScore] = Field(default_factory=list) metadata: dict[str, Any] = Field(default_factory=dict) diff --git a/shared/core/honeytokens.py b/shared/core/honeytokens.py deleted file mode 100644 index f802ca9..0000000 --- a/shared/core/honeytokens.py +++ /dev/null @@ -1,162 +0,0 @@ -from __future__ import annotations - -"""Honeytoken canary URLs for trace integrity enforcement. - -A honeytoken is a fictional URL that the tool proxy injects into a small -fraction of tool responses. If a miner response cites a honeytoken, we -know the miner hallucinated — honeytokens do not exist in the real world, -so a citation proves the miner is fabricating plausible-looking URLs -instead of actually reading what the tool returned. - -Key properties: - * Per-run rotation — a fresh set is generated on run open, so miners - cannot build a blocklist across runs. - * Deterministic generation from ``(run_id, seed)`` — recovery-safe, - easy to audit. - * Plausible shape — the URLs look like real archive / news links so a - non-reading miner that scrapes URL patterns from the response body - cannot distinguish them from legitimate search results. - * Injection rate is low (~2%) so honest miners rarely encounter them - and honeytokens remain scarce enough to serve as a detection signal. -""" - -import hashlib -import random -from typing import Any - - -# Canonical marker embedded in every honeytoken path. Allows fast rejection -# checks and helps operators grep logs / DB rows. -HONEYTOKEN_MARKER = "eirel-canary" - -# Plausible-looking hostnames that will never resolve to real content. -# ``.example`` is an IANA reserved TLD — guaranteed not to resolve. -_HONEYTOKEN_HOSTS: tuple[str, ...] = ( - "archive-news.example", - "wire-research.example", - "press-releases.example", - "docs-library.example", - "open-registry.example", -) - - -def generate_honeytoken_set( - run_id: str, - *, - count: int = 8, - seed: int | None = None, -) -> list[str]: - """Return a deterministic list of honeytoken URLs for a run. - - Same ``(run_id, seed, count)`` always produces the same list — - stored alongside the run and queried by the scoring pipeline at - evaluation time. ``run_id`` alone is sufficient in production; - ``seed`` exists for tests. - """ - if count <= 0: - return [] - base = f"{run_id}:{seed if seed is not None else 0}".encode("utf-8") - digest = hashlib.sha256(base).hexdigest() - rng = random.Random(int(digest[:16], 16)) - out: list[str] = [] - for idx in range(count): - host = rng.choice(_HONEYTOKEN_HOSTS) - suffix = hashlib.sha256(f"{digest}:{idx}".encode("utf-8")).hexdigest()[:10] - out.append(f"https://{host}/{HONEYTOKEN_MARKER}/{idx:02d}-{suffix}") - return out - - -def is_honeytoken(url: str, active_set: list[str] | set[str] | None) -> bool: - """Return True if ``url`` is in the active honeytoken set. - - Membership test is case-insensitive and tolerates trailing - punctuation / whitespace. Also accepts substring match against the - ``HONEYTOKEN_MARKER`` as a safety net — any URL that looks like a - honeytoken (path contains the marker) is treated as one even if the - active set was lost (fail-closed against fabricated URLs that share - the marker shape). - """ - if not url: - return False - lowered = url.strip().rstrip(".,;:)").lower() - if HONEYTOKEN_MARKER in lowered: - return True - if not active_set: - return False - lower_set = {item.strip().lower() for item in active_set if item} - return lowered in lower_set - - -def detect_honeytoken_citation( - claims: list[str], - active_set: list[str] | set[str] | None, -) -> bool: - """Return True iff any claim names an active honeytoken URL. - - Operates on the tagged-claim list produced by - ``extract_attribution_claims``: strips the ``url:`` prefix before - testing. Also catches bare-URL legacy claims. Ignores - ``search:*`` and ``authority:*`` claims — those are not URL-based. - """ - if not claims: - return False - for claim in claims: - if claim.startswith("search:") or claim.startswith("authority:"): - continue - candidate = claim[len("url:"):] if claim.startswith("url:") else claim - if is_honeytoken(candidate, active_set): - return True - return False - - -def inject_honeytokens_into_search_payload( - payload: Any, - *, - active_set: list[str], - conversation_id: str, - call_index: int, - rate: float = 0.02, -) -> Any: - """Inject a honeytoken result into a search-style tool payload. - - Only runs when the deterministic dice roll lands — ``rate`` of calls - receive an injection, keyed by ``(conversation_id, call_index)`` so - two independent runs against the same honeytoken set can differ. - - Safe no-op when: - * ``active_set`` is empty - * the payload is not a dict or has no ``results``/``items``/``data`` list - * the dice roll doesn't hit - - Returns the (possibly-mutated) payload. The original object is - modified in place to avoid copying large responses. - """ - if not active_set or not isinstance(payload, dict): - return payload - key = f"{conversation_id}:{call_index}".encode("utf-8") - roll = int(hashlib.sha256(key).hexdigest()[:8], 16) / 0xFFFFFFFF - if roll > rate: - return payload - # Find the first list-of-results field in common search response shapes. - list_key: str | None = None - for candidate in ("results", "items", "data", "hits", "documents"): - value = payload.get(candidate) - if isinstance(value, list): - list_key = candidate - break - if list_key is None: - return payload - # Deterministic pick of which honeytoken to inject. - pick_idx = int(hashlib.sha256(key).hexdigest()[8:16], 16) % len(active_set) - token_url = active_set[pick_idx] - injected = { - "title": f"Reference archive — entry {pick_idx:02d}", - "url": token_url, - "snippet": ( - "Historical record of relevant proceedings. " - "See full document for details." - ), - "source": "archive", - } - payload[list_key].append(injected) - return payload diff --git a/shared/scoring/families/_judge_to_conversation_score.py b/shared/scoring/families/_judge_to_conversation_score.py index c1c610e..caa366f 100644 --- a/shared/scoring/families/_judge_to_conversation_score.py +++ b/shared/scoring/families/_judge_to_conversation_score.py @@ -1,17 +1,11 @@ from __future__ import annotations -"""Translate judge output into a ConversationScore without requiring a trace. +"""Translate judge output into a ConversationScore. -The layered general_chat scoring path requires the miner to emit both a -``trace`` and ``response_text`` in its response dict so the trace-integrity -gate can run. Miners that don't opt in (including the example agent) -produce a raw ``AgentInvocationResponse`` with no trace. - -``build_conversation_score_from_judge`` fills that gap: it synthesizes a -valid ``ConversationScore`` from the judge sidecar's output plus whatever -latency signal is available in the miner response. It keeps the same -``total = trace_gate * (QUALITY_WEIGHT * quality + LATENCY_WEIGHT * latency)`` -shape as the layered path, so scores stay comparable across miners. +``build_conversation_score_from_judge`` synthesizes a valid +``ConversationScore`` from the judge sidecar's output plus whatever +latency signal is available in the miner response. Final score: +``total = QUALITY_WEIGHT * quality + LATENCY_WEIGHT * latency``. Callers: * ``_on_miner_evaluation_complete`` in the evaluation task manager — @@ -73,20 +67,15 @@ def build_conversation_score_from_judge( miner_response: dict[str, Any], mode: Literal["instant", "thinking"] = "instant", ) -> ConversationScore: - """Build a ConversationScore from judge output when no trace is available. - - The returned score is shaped identically to one produced by the layered - path, so the family aggregator (``aggregate_miner_score``) can consume - it without special-casing. + """Build a ConversationScore from judge output. ``quality`` prefers ``judge_output["score"]`` (the judge's own 0-1 overall score) and falls back to ``task_score`` — the latter is the value the validator already computed and submitted, which matches the judge's output when the layered path didn't run. - ``trace_gate`` defaults to 1.0 (pass) because there's no trace to gate. - Constraint flags from the judge are recorded in metadata for audit but - do NOT zero the gate — they're quality signals, not integrity signals. + Constraint flags from the judge are recorded in metadata for audit; + they are quality signals only. ``cost`` defaults to 0.0. Per-conversation cost attribution is not tracked here; run-level cost lives in ``DeploymentScoreRecord`` via @@ -106,11 +95,9 @@ def build_conversation_score_from_judge( latency_ms = _extract_latency_ms(miner_response, judge_output if isinstance(judge_output, dict) else {}) latency = _latency_score(latency_ms, mode=mode) - trace_gate = 1.0 # No trace to gate; constraint_flags are quality signals. cost = 0.0 - weighted = _QUALITY_WEIGHT * quality + _LATENCY_WEIGHT * latency - total = round(trace_gate * weighted, 6) + total = round(_QUALITY_WEIGHT * quality + _LATENCY_WEIGHT * latency, 6) constraint_flags = ( judge_output.get("constraint_flags") @@ -121,7 +108,6 @@ def build_conversation_score_from_judge( quality=round(quality, 6), latency=round(latency, 6), cost=cost, - trace_gate=trace_gate, total=total, per_dimension={ str(k): round(float(v), 6) diff --git a/tests/owner_api/test_scoring_manager_penalty.py b/tests/owner_api/test_scoring_manager_penalty.py deleted file mode 100644 index 46a3cb8..0000000 --- a/tests/owner_api/test_scoring_manager_penalty.py +++ /dev/null @@ -1,84 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace -from typing import Any - -import httpx - -from control_plane.owner_api.evaluation.scoring_manager import ScoringManager - - -def _settings(proxy_url: str = "http://provider-proxy.test") -> SimpleNamespace: - return SimpleNamespace( - provider_proxy_url=proxy_url, - provider_proxy_token="internal-token", - ) - - -def _owner(proxy_url: str = "http://provider-proxy.test") -> SimpleNamespace: - return SimpleNamespace(settings=_settings(proxy_url), db=None) - - -def test_charge_trace_gate_penalty_hits_provider_proxy(monkeypatch): - captured: dict[str, Any] = {} - - def _fake_post(url, json, headers, timeout): - captured["url"] = url - captured["json"] = json - captured["headers"] = headers - return httpx.Response( - 200, - json={ - "cost_usd_used": 5.50, - "max_usd_budget": 30.0, - "reason": "trace_gate_fail", - }, - request=httpx.Request("POST", url), - ) - - monkeypatch.setattr(httpx, "post", _fake_post) - - manager = ScoringManager(_owner()) - ok = manager.charge_trace_gate_penalty( - "dep-abc", amount_usd=0.50, reason="trace_gate_fail" - ) - assert ok is True - assert captured["url"] == "http://provider-proxy.test/v1/jobs/miner-dep-abc/charge_penalty" - assert captured["json"] == {"reason": "trace_gate_fail", "amount_usd": 0.50} - assert captured["headers"] == {"Authorization": "Bearer internal-token"} - - -def test_charge_trace_gate_penalty_zero_amount_is_noop(monkeypatch): - called = {"count": 0} - - def _fake_post(*args, **kwargs): - called["count"] += 1 - return httpx.Response(200, request=httpx.Request("POST", "http://x")) - - monkeypatch.setattr(httpx, "post", _fake_post) - - manager = ScoringManager(_owner()) - assert manager.charge_trace_gate_penalty("dep", amount_usd=0.0) is False - assert called["count"] == 0 - - -def test_charge_trace_gate_penalty_returns_false_when_proxy_unconfigured(monkeypatch): - def _fake_post(*args, **kwargs): # pragma: no cover - must not be called - raise AssertionError("httpx.post should not be called") - - monkeypatch.setattr(httpx, "post", _fake_post) - - manager = ScoringManager(_owner(proxy_url="")) - assert manager.charge_trace_gate_penalty("dep", amount_usd=0.50) is False - - -def test_charge_trace_gate_penalty_swallows_network_errors(monkeypatch): - def _fake_post(url, json, headers, timeout): - raise httpx.ConnectError("connection refused") - - monkeypatch.setattr(httpx, "post", _fake_post) - - manager = ScoringManager(_owner()) - # Must not raise — penalty-charging failures should be logged but - # never break the scoring pipeline. - assert manager.charge_trace_gate_penalty("dep", amount_usd=0.50) is False diff --git a/tests/tool_platforms/test_provider_proxy_usd_budget.py b/tests/tool_platforms/test_provider_proxy_usd_budget.py index f200533..0a6c834 100644 --- a/tests/tool_platforms/test_provider_proxy_usd_budget.py +++ b/tests/tool_platforms/test_provider_proxy_usd_budget.py @@ -198,16 +198,13 @@ async def test_get_cost_endpoint_returns_snapshot(_patch_env): # Split surfaces for DeploymentScoreRecord population. # After one successful LLM call and no tool charges: # * llm_cost_usd comes from the bare ``chutes`` entry - # * tool_cost_usd and penalty_cost_usd are zero (no tool: - # or penalty: prefixed entries) + # * tool_cost_usd is zero (no ``tool:`` prefixed entries) assert "llm_cost_usd" in body assert "tool_cost_usd" in body - assert "penalty_cost_usd" in body assert body["llm_cost_usd"] == pytest.approx( body["per_provider"].get("chutes", 0.0) ) assert body["tool_cost_usd"] == pytest.approx(0.0, abs=1e-9) - assert body["penalty_cost_usd"] == pytest.approx(0.0, abs=1e-9) async def test_charge_tool_increments_cost(_patch_env): @@ -269,70 +266,3 @@ async def test_missing_budget_header_rejected(_patch_env): ) assert resp.status_code == 400 assert "missing run budget header" in resp.json()["detail"] - - -async def test_charge_penalty_increments_cost_unconditionally(_patch_env): - app = create_app() - async with app.router.lifespan_context(app): - transport = ASGITransport(app=app) - async with AsyncClient(transport=transport, base_url="http://test") as client: - await client.post( - "/v1/chat/completions", - json=_base_payload(), - headers=_base_headers("5.00"), - ) - cost_before = (await client.get( - "/v1/jobs/job-usd-1/cost", - headers={"Authorization": "Bearer master-token"}, - )).json()["cost_usd_used"] - - charge_resp = await client.post( - "/v1/jobs/job-usd-1/charge_penalty", - json={"reason": "trace_gate_fail", "amount_usd": 0.50}, - headers={"Authorization": "Bearer master-token"}, - ) - assert charge_resp.status_code == 200 - body = charge_resp.json() - assert body["cost_usd_used"] == pytest.approx(cost_before + 0.50) - assert body["reason"] == "trace_gate_fail" - - # Per-provider map records the penalty under a "penalty:" prefix. - cost_after = (await client.get( - "/v1/jobs/job-usd-1/cost", - headers={"Authorization": "Bearer master-token"}, - )).json() - assert cost_after["per_provider"].get("penalty:trace_gate_fail") == pytest.approx(0.50) - - -async def test_charge_penalty_overshoots_budget_without_429(_patch_env): - app = create_app() - async with app.router.lifespan_context(app): - transport = ASGITransport(app=app) - async with AsyncClient(transport=transport, base_url="http://test") as client: - await client.post( - "/v1/chat/completions", - json=_base_payload(), - headers=_base_headers("5.00"), - ) - # 100x larger than budget — charge_tool would 429, charge_penalty must land. - charge_resp = await client.post( - "/v1/jobs/job-usd-1/charge_penalty", - json={"reason": "trace_gate_fail", "amount_usd": 500.0}, - headers={"Authorization": "Bearer master-token"}, - ) - assert charge_resp.status_code == 200 - body = charge_resp.json() - assert body["cost_usd_used"] >= 500.0 - - -async def test_charge_penalty_404_for_unknown_job(_patch_env): - app = create_app() - async with app.router.lifespan_context(app): - transport = ASGITransport(app=app) - async with AsyncClient(transport=transport, base_url="http://test") as client: - resp = await client.post( - "/v1/jobs/no-such-job/charge_penalty", - json={"reason": "trace_gate_fail", "amount_usd": 0.50}, - headers={"Authorization": "Bearer master-token"}, - ) - assert resp.status_code == 404 diff --git a/tool_platforms/provider_proxy/app.py b/tool_platforms/provider_proxy/app.py index 57231e0..ffb7c55 100644 --- a/tool_platforms/provider_proxy/app.py +++ b/tool_platforms/provider_proxy/app.py @@ -43,11 +43,6 @@ class ChargeToolRequest(_BaseModel): amount_usd: float = _Field(ge=0.0) -class ChargePenaltyRequest(_BaseModel): - reason: str - amount_usd: float = _Field(ge=0.0) - - @dataclass(slots=True) class AppState: auth_token: str @@ -173,29 +168,24 @@ async def job_cost(job_id: str, _: str = Depends(require_auth)) -> dict[str, Any status_code=status.HTTP_404_NOT_FOUND, detail="job not found", ) - # ``cost_by_provider`` keys come in three shapes: + # ``cost_by_provider`` keys come in two shapes: # * bare LLM provider names ("chutes", "openai", ...) — # written by /chat/completions reconciliation # * ``tool:`` — written by ``charge_tool`` - # * ``penalty:`` — written by ``charge_penalty`` # Split on those prefixes so ScoringManager.populate_cost_columns # gets truthful ``DeploymentScoreRecord.{llm_cost_usd,tool_cost_usd}`` # instead of lumping everything together. llm_cost_usd = 0.0 tool_cost_usd = 0.0 - penalty_cost_usd = 0.0 for key, value in usage.cost_by_provider.items(): if key.startswith("tool:"): tool_cost_usd += float(value) - elif key.startswith("penalty:"): - penalty_cost_usd += float(value) else: llm_cost_usd += float(value) return { "cost_usd_used": usage.cost_usd_used, "llm_cost_usd": round(llm_cost_usd, 8), "tool_cost_usd": round(tool_cost_usd, 8), - "penalty_cost_usd": round(penalty_cost_usd, 8), "max_usd_budget": usage.max_usd_budget, "cost_rejections": usage.cost_rejections, "per_provider": dict(usage.cost_by_provider), @@ -223,35 +213,6 @@ async def charge_tool( ) return {"cost_usd_used": cost_used} - @app.post("/v1/jobs/{job_id}/charge_penalty") - async def charge_penalty( - job_id: str, - body: ChargePenaltyRequest, - _: str = Depends(require_auth), - ) -> dict[str, Any]: - """Charge an unconditional USD penalty against the run budget. - - Unlike ``charge_tool``, this endpoint never 429-rejects — the - penalty always lands even if it overshoots ``max_usd_budget``. - Used by the scoring pipeline when a trace integrity gate fires, - so bad-actor miners pay real cost for each failed conversation. - """ - store = app.state.services.store - if not await store.exists(job_id): - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail="job not found", - ) - cost_used = await store.charge_penalty( - job_id=job_id, reason=body.reason, amount_usd=body.amount_usd, - ) - usage = await store.get(job_id) - return { - "cost_usd_used": cost_used, - "max_usd_budget": usage.max_usd_budget if usage else 0.0, - "reason": body.reason, - } - @app.post("/v1/chat/completions", response_model=ProviderProxyResponse) async def chat_completions( payload: ProviderProxyRequest, diff --git a/tool_platforms/provider_proxy/redis_store.py b/tool_platforms/provider_proxy/redis_store.py index f0d919a..00dd81b 100644 --- a/tool_platforms/provider_proxy/redis_store.py +++ b/tool_platforms/provider_proxy/redis_store.py @@ -20,8 +20,8 @@ provider_proxy:provider_counts: (hash: provider → int) provider_proxy:model_counts: (hash: provider:model → int) provider_proxy:cost_by_provider: (hash: key → float) - keys are either a bare provider name (LLM), ``tool:``, - or ``penalty:`` — the split ScoringManager consumes. + keys are either a bare provider name (LLM) or ``tool:`` — + the split ScoringManager consumes. Concurrent request safety: ``reserve_estimate`` and ``charge_tool`` execute a Lua script so @@ -354,25 +354,6 @@ async def _charge_tool_fallback( await pipe.execute() return (True, cost + amount_usd) - async def charge_penalty( - self, - *, - job_id: str, - reason: str, - amount_usd: float, - ) -> float: - """Charge an unconditional penalty (always lands, no 429).""" - usage_key = _USAGE_PREFIX + job_id - cost_key = _COST_BY_PROVIDER_PREFIX + job_id - async with self._client.pipeline(transaction=True) as pipe: - pipe.hincrbyfloat(usage_key, "cost_usd_used", float(amount_usd)) - pipe.hincrbyfloat(cost_key, f"penalty:{reason}", float(amount_usd)) - pipe.expire(usage_key, self._ttl) - pipe.expire(cost_key, self._ttl) - pipe.hget(usage_key, "cost_usd_used") - values = await pipe.execute() - return float(values[-1] or 0.0) - # ------------------------------------------------------------------ # Reads # ------------------------------------------------------------------ From e0e5af120d448816c58e09cebb8f5100d63532d2 Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Thu, 7 May 2026 17:29:29 +0000 Subject: [PATCH 10/24] feat(schema): add EvalFeedback + OrchestratorToolCallLog tables; per-dimension TaskMinerResult fields --- shared/common/migrations.py | 639 +++++++----------------------- shared/common/models.py | 766 ++++++++++++++++++++++++++++++++++-- 2 files changed, 876 insertions(+), 529 deletions(-) diff --git a/shared/common/migrations.py b/shared/common/migrations.py index df167a4..e07def3 100644 --- a/shared/common/migrations.py +++ b/shared/common/migrations.py @@ -1,3 +1,14 @@ +"""Schema migration runner for eirel-ai. + +The schema is owned by ``shared.common.models.Base.metadata.create_all``; +migrations are a thin advisory-lock wrapper that records "the schema as +of this revision" in the ``schema_migrations`` table. + +Single migration: ``initial_schema``. Pre-launch we collapsed every +prior migration step into ``Base.metadata.create_all`` since no +production DBs need in-place ALTERs. Future migrations append a new +``Migration`` entry to ``MIGRATIONS`` with one-purpose ALTER DDL. +""" from __future__ import annotations import logging @@ -7,7 +18,6 @@ from sqlalchemy import text from sqlalchemy.engine import Engine -from sqlalchemy import inspect _logger = logging.getLogger(__name__) @@ -25,16 +35,20 @@ def run_migrations(engine: Engine) -> list[str]: dialect = engine.dialect.name if dialect == "postgresql": with engine.begin() as conn: - conn.execute(text("SELECT pg_advisory_lock(:lock_id)"), - {"lock_id": _MIGRATION_ADVISORY_LOCK_ID}) + conn.execute( + text("SELECT pg_advisory_lock(:lock_id)"), + {"lock_id": _MIGRATION_ADVISORY_LOCK_ID}, + ) _logger.info("acquired migration advisory lock") try: return _run_migrations_unlocked(engine) finally: if dialect == "postgresql": with engine.begin() as conn: - conn.execute(text("SELECT pg_advisory_unlock(:lock_id)"), - {"lock_id": _MIGRATION_ADVISORY_LOCK_ID}) + conn.execute( + text("SELECT pg_advisory_unlock(:lock_id)"), + {"lock_id": _MIGRATION_ADVISORY_LOCK_ID}, + ) _logger.info("released migration advisory lock") @@ -77,537 +91,150 @@ def _ensure_schema_migrations_table(engine: Engine) -> None: def _applied_versions(engine: Engine) -> set[str]: with engine.begin() as conn: - rows = conn.execute(text("SELECT version FROM schema_migrations")).fetchall() + rows = conn.execute( + text("SELECT version FROM schema_migrations") + ).fetchall() return {str(row[0]) for row in rows} -def _migration_family_native_staging_bootstrap(engine: Engine) -> None: - del engine - +def _migration_initial_schema(engine: Engine) -> None: + """Marker migration — the schema itself is created by + ``Base.metadata.create_all`` in ``Database.create_all``. -def _migration_remove_legacy_workflow_market_schema(engine: Engine) -> None: - inspector = inspect(engine) - tables = set(inspector.get_table_names()) - with engine.begin() as conn: - if "coalition_score_snapshots" in tables: - conn.execute(text("DROP TABLE coalition_score_snapshots")) - episode_columns = { - column["name"] - for column in inspector.get_columns("workflow_episode_records") - } if "workflow_episode_records" in tables else set() - if "coalition_json" in episode_columns: - conn.execute(text("ALTER TABLE workflow_episode_records DROP COLUMN coalition_json")) - - -def _migration_add_distributed_evaluation_schema(engine: Engine) -> None: - inspector = inspect(engine) - tables = set(inspector.get_table_names()) - with engine.begin() as conn: - if "miner_evaluation_tasks" not in tables: - conn.execute(text(""" - CREATE TABLE miner_evaluation_tasks ( - id VARCHAR(36) PRIMARY KEY, - epoch_id VARCHAR(128) NOT NULL, - family_id VARCHAR(64) NOT NULL, - miner_hotkey VARCHAR(128) NOT NULL, - task_id VARCHAR(128) NOT NULL, - task_index INTEGER NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - claimed_by_validator VARCHAR(128), - claimed_at TIMESTAMP, - claim_expires_at TIMESTAMP, - claim_attempt_count INTEGER NOT NULL DEFAULT 0, - miner_response_json JSON, - judge_output_json JSON, - task_score FLOAT, - task_status VARCHAR(32), - result_metadata_json JSON NOT NULL DEFAULT '{}', - evaluated_at TIMESTAMP, - created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - UNIQUE(epoch_id, family_id, miner_hotkey, task_id) - ) - """)) - conn.execute(text( - "CREATE INDEX idx_met_claimable ON miner_evaluation_tasks (epoch_id, family_id, status, claim_expires_at)" - )) - conn.execute(text( - "CREATE INDEX idx_met_miner ON miner_evaluation_tasks (epoch_id, family_id, miner_hotkey)" - )) - conn.execute(text( - "CREATE INDEX idx_met_epoch_id ON miner_evaluation_tasks (epoch_id)" - )) - conn.execute(text( - "CREATE INDEX idx_met_family_id ON miner_evaluation_tasks (family_id)" - )) - conn.execute(text( - "CREATE INDEX idx_met_status ON miner_evaluation_tasks (status)" - )) - conn.execute(text( - "CREATE INDEX idx_met_claimed_by ON miner_evaluation_tasks (claimed_by_validator)" - )) - if "miner_evaluation_summaries" not in tables: - conn.execute(text(""" - CREATE TABLE miner_evaluation_summaries ( - id VARCHAR(36) PRIMARY KEY, - epoch_id VARCHAR(128) NOT NULL, - family_id VARCHAR(64) NOT NULL, - miner_hotkey VARCHAR(128) NOT NULL, - total_tasks INTEGER NOT NULL, - completed_tasks INTEGER NOT NULL DEFAULT 0, - failed_tasks INTEGER NOT NULL DEFAULT 0, - family_capability_score FLOAT, - robustness_score FLOAT, - anti_gaming_score FLOAT, - official_family_score FLOAT, - protocol_gate_passed BOOLEAN, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - rollout_metadata_json JSON NOT NULL DEFAULT '{}', - created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - UNIQUE(epoch_id, family_id, miner_hotkey) - ) - """)) - conn.execute(text( - "CREATE INDEX idx_mes_epoch_id ON miner_evaluation_summaries (epoch_id)" - )) - conn.execute(text( - "CREATE INDEX idx_mes_family_id ON miner_evaluation_summaries (family_id)" - )) - conn.execute(text( - "CREATE INDEX idx_mes_miner ON miner_evaluation_summaries (miner_hotkey)" - )) - conn.execute(text( - "CREATE INDEX idx_mes_status ON miner_evaluation_summaries (status)" - )) + Pre-launch we collapsed every prior schema step (validators table, + submission/deployment plumbing, evaluation runs/tasks, owner dataset + bindings, consumer-side product tables, MCP catalog, server-attested + tool-call ledger, etc.) into the ``Base`` SQLAlchemy declaration. + Fresh DBs come up via ``create_all``; this migration row records + that fact in ``schema_migrations`` so future ALTER-style migrations + have a baseline to anchor against. + """ + del engine -def _migration_add_neuron_uid_table(engine: Engine) -> None: - inspector = inspect(engine) - tables = set(inspector.get_table_names()) - if "neuron_uids" not in tables: - with engine.begin() as conn: - conn.execute(text(""" - CREATE TABLE neuron_uids ( - hotkey VARCHAR(128) PRIMARY KEY, - uid INTEGER NOT NULL, - stake BIGINT NOT NULL DEFAULT 0, - is_validator BOOLEAN NOT NULL DEFAULT FALSE, - is_active BOOLEAN NOT NULL DEFAULT TRUE, - last_synced_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP - ) - """)) - conn.execute(text("CREATE INDEX idx_neuron_uids_uid ON neuron_uids (uid)")) +def _migration_multi_metric_scoring(engine: Engine) -> None: + """Add per-dimension score columns to ``task_miner_results``. + Each task is now scored along six independent dimensions + (``pairwise_preference_score`` + 5 outer metrics) plus an aggregate + ``final_task_score``. ``applied_weights_json`` records the actual + weights after N/A re-normalization for the task type. New columns + are nullable so legacy pairwise-only rows coexist. -def _migration_add_owner_dataset_bindings(engine: Engine) -> None: + Skipped on fresh databases — there ``Base.metadata.create_all`` + runs after migrations and creates the table with the new columns + already in place. Migrations only do work on pre-existing DBs that + were bootstrapped before this column set was added. + """ + from sqlalchemy import inspect, text inspector = inspect(engine) - tables = set(inspector.get_table_names()) - if "owner_dataset_bindings" in tables: + if "task_miner_results" not in inspector.get_table_names(): return - with engine.begin() as conn: - conn.execute(text(""" - CREATE TABLE owner_dataset_bindings ( - id VARCHAR(36) PRIMARY KEY, - family_id VARCHAR(64) NOT NULL, - run_id VARCHAR(128) NOT NULL, - bundle_uri VARCHAR(1024) NOT NULL, - bundle_sha256 VARCHAR(64) NOT NULL, - generator_version VARCHAR(128) NOT NULL, - generated_by VARCHAR(128) NOT NULL, - signature_hex VARCHAR(256) NOT NULL, - generator_provider VARCHAR(64) NOT NULL DEFAULT '', - generator_model VARCHAR(128) NOT NULL DEFAULT '', - status VARCHAR(32) NOT NULL DEFAULT 'pending', - provenance_json JSON NOT NULL DEFAULT '{}', - created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - activated_at TIMESTAMP NULL, - CONSTRAINT uq_owner_dataset_bindings_family_run UNIQUE (family_id, run_id) + existing_columns = { + col["name"] for col in inspector.get_columns("task_miner_results") + } + columns_to_add = ( + ("pairwise_preference_score", "DOUBLE PRECISION"), + ("grounded_correctness", "DOUBLE PRECISION"), + ("retrieval_quality", "DOUBLE PRECISION"), + ("tool_routing", "DOUBLE PRECISION"), + ("instruction_safety", "DOUBLE PRECISION"), + ("latency_cost", "DOUBLE PRECISION"), + ("computation_correctness", "DOUBLE PRECISION"), + ("final_task_score", "DOUBLE PRECISION"), + ("applied_weights_json", "JSON"), + ("applicable_metrics_json", "JSON"), + ("task_type", "VARCHAR(64)"), + ) + with engine.begin() as conn: + for name, col_type in columns_to_add: + if name in existing_columns: + continue + conn.execute( + text( + f"ALTER TABLE task_miner_results ADD COLUMN {name} {col_type}" + ) ) - """)) - conn.execute(text( - "CREATE INDEX idx_owner_dataset_bindings_family ON owner_dataset_bindings (family_id)" - )) - conn.execute(text( - "CREATE INDEX idx_owner_dataset_bindings_run ON owner_dataset_bindings (run_id)" - )) - conn.execute(text( - "CREATE INDEX idx_owner_dataset_bindings_family_status " - "ON owner_dataset_bindings (family_id, status)" - )) -def _migration_add_cost_accounting_columns(engine: Engine) -> None: - inspector = inspect(engine) - tables = set(inspector.get_table_names()) - if "deployment_score_records" not in tables: - return - existing = {col["name"] for col in inspector.get_columns("deployment_score_records")} - new_columns = { - "run_budget_usd": "FLOAT NOT NULL DEFAULT 30.0", - "run_cost_usd_used": "FLOAT NOT NULL DEFAULT 0.0", - "llm_cost_usd": "FLOAT NOT NULL DEFAULT 0.0", - "tool_cost_usd": "FLOAT NOT NULL DEFAULT 0.0", - "cost_rejection_count": "INTEGER NOT NULL DEFAULT 0", - } - with engine.begin() as conn: - for col_name, col_def in new_columns.items(): - if col_name not in existing: - conn.execute(text( - f"ALTER TABLE deployment_score_records ADD COLUMN {col_name} {col_def}" - )) - +def _migration_eval_feedback_table(engine: Engine) -> None: + """Create the ``eval_feedback`` table for per-(run, miner, task) + EvalJudge outcomes. -def _migration_add_pending_runtime_stop(engine: Engine) -> None: - inspector = inspect(engine) - tables = set(inspector.get_table_names()) - if "managed_deployments" not in tables: - return - existing = {col["name"] for col in inspector.get_columns("managed_deployments")} - if "pending_runtime_stop" in existing: - return - default_literal = "false" if engine.dialect.name == "postgresql" else "0" - with engine.begin() as conn: - conn.execute(text( - f"ALTER TABLE managed_deployments ADD COLUMN pending_runtime_stop BOOLEAN NOT NULL DEFAULT {default_literal}" - )) - - -def _migration_add_snapshot_unique_constraint(engine: Engine) -> None: + Skipped on fresh databases — there ``Base.metadata.create_all`` + runs after migrations and creates the table from the SQLAlchemy + model declaration. This migration only does work on pre-existing + DBs that were bootstrapped before the table existed. + """ + from sqlalchemy import inspect, text inspector = inspect(engine) - tables = set(inspector.get_table_names()) - if "epoch_target_snapshots" not in tables: - return - indexes = inspector.get_indexes("epoch_target_snapshots") - existing_names = {idx["name"] for idx in indexes} - if "uq_snapshot_run_family" in existing_names: - return - if "uq_epoch_target_snapshots_epoch_family" in existing_names: + if "eval_feedback" in inspector.get_table_names(): return with engine.begin() as conn: - conn.execute(text( - "CREATE UNIQUE INDEX uq_snapshot_run_family " - "ON epoch_target_snapshots (epoch_id, family_id)" - )) + conn.execute( + text( + """ + CREATE TABLE IF NOT EXISTS eval_feedback ( + id VARCHAR(36) PRIMARY KEY, + run_id VARCHAR(64) NOT NULL, + miner_hotkey VARCHAR(64) NOT NULL, + task_id VARCHAR(64) NOT NULL, + outcome VARCHAR(32) NOT NULL, + failure_mode VARCHAR(64), + guidance TEXT NOT NULL DEFAULT '', + prompt_excerpt TEXT NOT NULL DEFAULT '', + response_excerpt TEXT NOT NULL DEFAULT '', + composite_score DOUBLE PRECISION NOT NULL DEFAULT 0.0, + knockout_reasons_json JSON NOT NULL, + oracle_status VARCHAR(32), + created_at TIMESTAMP NOT NULL, + CONSTRAINT uq_eval_feedback_run_miner_task + UNIQUE (run_id, miner_hotkey, task_id) + ) + """ + ) + ) + conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_eval_feedback_miner_run " + "ON eval_feedback (miner_hotkey, run_id)" + ) + ) + conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_eval_feedback_run_task " + "ON eval_feedback (run_id, task_id)" + ) + ) MIGRATIONS: tuple[Migration, ...] = ( Migration( - version="family_native_staging_bootstrap", - description="Initialize family-native staging schema baseline", - apply=_migration_family_native_staging_bootstrap, - ), - Migration( - version="remove_legacy_workflow_market_schema", - description="Drop legacy workflow-market coalition storage", - apply=_migration_remove_legacy_workflow_market_schema, - ), - Migration( - version="add_distributed_evaluation_schema", - description="Add miner_evaluation_tasks and miner_evaluation_summaries tables for distributed task evaluation", - apply=_migration_add_distributed_evaluation_schema, - ), - Migration( - version="add_neuron_uid_table", - description="Add neuron_uids table for all hotkey-to-uid mappings from metagraph", - apply=_migration_add_neuron_uid_table, - ), - Migration( - version="add_owner_dataset_bindings", - description="Add owner_dataset_bindings table for private dataset pipeline", - apply=_migration_add_owner_dataset_bindings, - ), - Migration( - version="add_cost_accounting_columns", - description="Add per-run USD cost accounting columns to deployment_score_records", - apply=_migration_add_cost_accounting_columns, - ), - Migration( - version="add_pending_runtime_stop", - description="Add pending_runtime_stop flag to managed_deployments for orphan cleanup", - apply=_migration_add_pending_runtime_stop, - ), - Migration( - version="add_snapshot_unique_constraint", - description="Add unique constraint on (epoch_id, family_id) to epoch_target_snapshots", - apply=_migration_add_snapshot_unique_constraint, - ), - Migration( - version="drop_registered_neurons_is_active", - description="Drop is_active column from registered_neurons — presence = registered", - apply=lambda engine: _drop_registered_neurons_is_active(engine), - ), - Migration( - version="refactor_to_task_level_evaluation", - description=( - "Replace miner_evaluation_tasks with task_evaluations + " - "task_miner_results for task-level validator claims and pairwise " - "judging vs OpenAI baseline" - ), - apply=lambda engine: _migration_refactor_to_task_level_evaluation(engine), - ), - Migration( - version="add_miner_latency_to_task_miner_results", - description=( - "Add miner_latency_seconds column to task_miner_results so the " - "leaderboard can show miner response latency separately from " - "judge latency, and for the latency-violation scoring gate" - ), - apply=lambda engine: _migration_add_miner_latency_to_task_miner_results(engine), - ), - Migration( - version="add_miner_first_token_to_task_miner_results", + version="initial_schema", description=( - "Add miner_first_token_seconds column to task_miner_results to " - "store time-to-first-token from the streaming invocation path; " - "feeds the mode-agnostic 10s TTFB SLA gate" + "Initial schema — owned by Base.metadata.create_all. " + "Subsequent migrations append ALTER DDL for in-place upgrades." ), - apply=lambda engine: _migration_add_miner_first_token_to_task_miner_results(engine), + apply=_migration_initial_schema, ), Migration( - version="drop_miner_first_token_seconds", + version="multi_metric_scoring", description=( - "Drop miner_first_token_seconds — TTFB metric removed; only " - "completion-time latency is recorded and gated. Most LLM " - "providers stream first token <20s anyway, so the gate added " - "noise without changing miner ranking." + "Per-dimension score columns on task_miner_results: " + "pairwise + grounded + retrieval + tool_routing + safety + " + "latency_cost + computation_correctness + final_task_score, " + "plus applied_weights_json / applicable_metrics_json / task_type." ), - apply=lambda engine: _migration_drop_miner_first_token_seconds(engine), + apply=_migration_multi_metric_scoring, ), Migration( - version="add_session_mode_web_search", + version="add_eval_feedback_table", description=( - "Persist per-session user toggles (mode, web_search) on " - "consumer_sessions so the orchestrator can apply them on every " - "turn without the client re-asserting them. Owned by the " - "orchestrator now that consumer-chat-api is a thin facade." + "Create eval_feedback table for per-(run, miner, task) " + "EvalJudge outcomes. Indexed on (miner_hotkey, run_id) for " + "the per-miner read path and (run_id, task_id) for " + "operator dashboard cross-miner drill-in." ), - apply=lambda engine: _migration_add_session_mode_web_search(engine), - ), - Migration( - version="add_proxy_and_judge_cost_to_task_miner_results", - description=( - "Add proxy_cost_usd + judge_cost_usd to task_miner_results" - ), - apply=lambda engine: _migration_add_proxy_and_judge_cost(engine), - ), - Migration( - version="drop_anti_gaming_artifacts", - description=( - "Strip honeytokens / trace_gate artifacts from JSON metadata " - "after the anti-gaming subsystem was removed" - ), - apply=lambda engine: _migration_drop_anti_gaming_artifacts(engine), + apply=_migration_eval_feedback_table, ), ) - - -def _migration_add_miner_latency_to_task_miner_results(engine: Engine) -> None: - inspector = inspect(engine) - if "task_miner_results" not in set(inspector.get_table_names()): - return - cols = {c["name"] for c in inspector.get_columns("task_miner_results")} - if "miner_latency_seconds" in cols: - return - with engine.begin() as conn: - conn.execute(text( - "ALTER TABLE task_miner_results " - "ADD COLUMN miner_latency_seconds FLOAT NOT NULL DEFAULT 0.0" - )) - - -def _migration_add_miner_first_token_to_task_miner_results(engine: Engine) -> None: - inspector = inspect(engine) - if "task_miner_results" not in set(inspector.get_table_names()): - return - cols = {c["name"] for c in inspector.get_columns("task_miner_results")} - if "miner_first_token_seconds" in cols: - return - with engine.begin() as conn: - conn.execute(text( - "ALTER TABLE task_miner_results " - "ADD COLUMN miner_first_token_seconds FLOAT NOT NULL DEFAULT 0.0" - )) - - -def _migration_drop_miner_first_token_seconds(engine: Engine) -> None: - inspector = inspect(engine) - if "task_miner_results" not in set(inspector.get_table_names()): - return - cols = {c["name"] for c in inspector.get_columns("task_miner_results")} - if "miner_first_token_seconds" not in cols: - return - with engine.begin() as conn: - conn.execute(text( - "ALTER TABLE task_miner_results DROP COLUMN miner_first_token_seconds" - )) - - -def _migration_add_session_mode_web_search(engine: Engine) -> None: - inspector = inspect(engine) - if "consumer_sessions" not in set(inspector.get_table_names()): - return - cols = {c["name"] for c in inspector.get_columns("consumer_sessions")} - with engine.begin() as conn: - if "mode" not in cols: - conn.execute(text( - "ALTER TABLE consumer_sessions " - "ADD COLUMN mode VARCHAR(16) NOT NULL DEFAULT 'instant'" - )) - if "web_search" not in cols: - conn.execute(text( - "ALTER TABLE consumer_sessions " - "ADD COLUMN web_search BOOLEAN NOT NULL DEFAULT FALSE" - )) - - -def _migration_add_proxy_and_judge_cost(engine: Engine) -> None: - inspector = inspect(engine) - if "task_miner_results" not in set(inspector.get_table_names()): - return - cols = {c["name"] for c in inspector.get_columns("task_miner_results")} - with engine.begin() as conn: - if "proxy_cost_usd" not in cols: - conn.execute(text( - "ALTER TABLE task_miner_results " - "ADD COLUMN proxy_cost_usd FLOAT NOT NULL DEFAULT 0.0" - )) - if "judge_cost_usd" not in cols: - conn.execute(text( - "ALTER TABLE task_miner_results " - "ADD COLUMN judge_cost_usd FLOAT NOT NULL DEFAULT 0.0" - )) - - -def _migration_drop_anti_gaming_artifacts(engine: Engine) -> None: - """Strip anti-gaming residue from existing rows. - - The honeytoken / trace-gate code path was removed. No SQL columns - were ever added for it, but ``evaluation_runs.metadata_json`` carried - a ``"honeytokens"`` key with the per-run canary URL list, and - ``deployment_score_records`` (via Pydantic JSON serialization) may - have ``trace_gate``, ``trace_gate_penalty_usd``, or ``honeytoken_cited`` - keys embedded in stored conversation-score blobs. Strip both for - cleanliness so reads don't surface stale fields. - """ - inspector = inspect(engine) - tables = set(inspector.get_table_names()) - if "evaluation_runs" not in tables: - return - cols = {col["name"] for col in inspector.get_columns("evaluation_runs")} - if "metadata_json" not in cols: - return - dialect = engine.dialect.name - with engine.begin() as conn: - if dialect == "postgresql": - # JSONB-aware: strip the "honeytokens" key in place. - conn.execute(text( - "UPDATE evaluation_runs " - "SET metadata_json = metadata_json::jsonb - 'honeytokens' " - "WHERE metadata_json::jsonb ? 'honeytokens'" - )) - else: - # SQLite / others: rebuild the JSON without the key. Cheap - # and correct because run counts are small. - rows = conn.execute(text( - "SELECT id, metadata_json FROM evaluation_runs " - "WHERE metadata_json IS NOT NULL" - )).fetchall() - import json as _json - for row in rows: - raw = row[1] - if not raw: - continue - try: - parsed = _json.loads(raw) if isinstance(raw, str) else dict(raw) - except (ValueError, TypeError): - continue - if not isinstance(parsed, dict) or "honeytokens" not in parsed: - continue - parsed.pop("honeytokens", None) - conn.execute( - text( - "UPDATE evaluation_runs SET metadata_json = :m " - "WHERE id = :id" - ), - {"m": _json.dumps(parsed), "id": row[0]}, - ) - - -def _drop_registered_neurons_is_active(engine: Engine) -> None: - inspector = inspect(engine) - if "registered_neurons" not in set(inspector.get_table_names()): - return - if "is_active" not in {c["name"] for c in inspector.get_columns("registered_neurons")}: - return - with engine.begin() as conn: - conn.execute(text("ALTER TABLE registered_neurons DROP COLUMN is_active")) - - -def _migration_refactor_to_task_level_evaluation(engine: Engine) -> None: - """Replace per-(miner, task) rows with task-level rows + per-miner results. - - Drops `miner_evaluation_tasks` (per-pair claim rows) and creates - `task_evaluations` (one row per task, validator claims this) plus - `task_miner_results` (one row per miner per task, stores pairwise judge - output vs OpenAI baseline). Non-reversible: any in-flight pairs are lost. - """ - inspector = inspect(engine) - tables = set(inspector.get_table_names()) - with engine.begin() as conn: - if "miner_evaluation_tasks" in tables: - conn.execute(text("DROP TABLE miner_evaluation_tasks")) - if "task_evaluations" not in tables: - conn.execute(text(""" - CREATE TABLE task_evaluations ( - id VARCHAR(36) PRIMARY KEY, - epoch_id VARCHAR(128) NOT NULL, - family_id VARCHAR(64) NOT NULL, - task_id VARCHAR(128) NOT NULL, - task_index INTEGER NOT NULL, - status VARCHAR(32) NOT NULL DEFAULT 'pending', - claimed_by_validator VARCHAR(128), - claimed_at TIMESTAMP, - claim_expires_at TIMESTAMP, - claim_attempt_count INTEGER NOT NULL DEFAULT 0, - baseline_response_json JSON, - evaluated_at TIMESTAMP, - created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - UNIQUE(epoch_id, family_id, task_id) - ) - """)) - conn.execute(text( - "CREATE INDEX idx_te_claimable ON task_evaluations (epoch_id, family_id, status, claim_expires_at)" - )) - conn.execute(text( - "CREATE INDEX idx_te_epoch_id ON task_evaluations (epoch_id)" - )) - conn.execute(text( - "CREATE INDEX idx_te_status ON task_evaluations (status)" - )) - if "task_miner_results" not in tables: - conn.execute(text(""" - CREATE TABLE task_miner_results ( - id VARCHAR(36) PRIMARY KEY, - task_evaluation_id VARCHAR(36) NOT NULL REFERENCES task_evaluations(id) ON DELETE CASCADE, - epoch_id VARCHAR(128) NOT NULL, - family_id VARCHAR(64) NOT NULL, - task_id VARCHAR(128) NOT NULL, - miner_hotkey VARCHAR(128) NOT NULL, - miner_response_json JSON NOT NULL, - miner_citations_json JSON NOT NULL DEFAULT '[]', - judge_output_json JSON, - agreement_verdict VARCHAR(32) NOT NULL, - agreement_score FLOAT NOT NULL DEFAULT 0.0, - latency_seconds FLOAT NOT NULL DEFAULT 0.0, - created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - UNIQUE(task_evaluation_id, miner_hotkey) - ) - """)) - conn.execute(text( - "CREATE INDEX idx_tmr_miner ON task_miner_results (epoch_id, family_id, miner_hotkey)" - )) - conn.execute(text( - "CREATE INDEX idx_tmr_task_eval ON task_miner_results (task_evaluation_id)" - )) diff --git a/shared/common/models.py b/shared/common/models.py index cc1b48a..93ad06c 100644 --- a/shared/common/models.py +++ b/shared/common/models.py @@ -719,15 +719,46 @@ class TaskMinerResult(Base): latency_seconds: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) # Per-task LLM cost the miner incurred against the subnet # provider-proxy, in USD. Sourced server-side from the - # provider-proxy ledger via owner-api injection (Phase 2a) — never - # trusts miner self-report. Zero when the miner made no proxied - # LLM calls or when the cost lookup failed. + # provider-proxy ledger via owner-api injection — never trusts + # miner self-report. Zero when the miner made no proxied LLM calls + # or when the cost lookup failed. proxy_cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) - # Per-task judge LLM cost in USD (Phase 2c). Reported by - # eiretes-judge in its response metadata; the validator passes it - # through verbatim. Always >= 0. + # Per-task judge LLM cost in USD. Reported by eiretes-judge in its + # response metadata; the validator passes it through verbatim. + # Always >= 0. judge_cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + # Graph-runtime: link an eval result back to a graph checkpoint + # thread + the specific checkpoint id the miner emitted. Both + # nullable for backwards compat with BaseAgent-style miners. + thread_id: Mapped[str | None] = mapped_column( + String(64), nullable=True, index=True, + ) + checkpoint_id: Mapped[str | None] = mapped_column( + String(64), nullable=True, + ) + + # Multi-metric per-task scoring. Each dimension is computed + # independently (pairwise from /v1/judge/pairwise; grounded / + # retrieval / safety from /v1/judge/multi; tool_routing / + # latency_cost / computation_correctness deterministically by the + # validator). All nullable so a metric marked N/A for the task type + # re-normalizes out cleanly. ``final_task_score`` is the weighted + # sum after re-normalization. ``applied_weights_json`` records what + # weights actually applied; ``applicable_metrics_json`` is the set + # of dimensions that had a real score (non-N/A). + pairwise_preference_score: Mapped[float | None] = mapped_column(Float, nullable=True) + grounded_correctness: Mapped[float | None] = mapped_column(Float, nullable=True) + retrieval_quality: Mapped[float | None] = mapped_column(Float, nullable=True) + tool_routing: Mapped[float | None] = mapped_column(Float, nullable=True) + instruction_safety: Mapped[float | None] = mapped_column(Float, nullable=True) + latency_cost: Mapped[float | None] = mapped_column(Float, nullable=True) + computation_correctness: Mapped[float | None] = mapped_column(Float, nullable=True) + final_task_score: Mapped[float | None] = mapped_column(Float, nullable=True) + applied_weights_json: Mapped[dict[str, float] | None] = mapped_column(JSON, nullable=True) + applicable_metrics_json: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) + task_type: Mapped[str | None] = mapped_column(String(64), nullable=True) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=False), default=utcnow, nullable=False) @property @@ -735,6 +766,84 @@ def epoch_id(self) -> str: return self.run_id +class GraphCheckpoint(Base): + """One persisted graph-runtime checkpoint. + + The miner SDK's ``PostgresCheckpointer`` writes here over HTTP + via ``/v1/internal/checkpoints/{thread_id}`` so the miner pod + never holds DB credentials. + """ + + __tablename__ = "graph_checkpoints" + __table_args__ = ( + UniqueConstraint( + "thread_id", "checkpoint_id", + name="uq_graph_checkpoint_thread_checkpoint", + ), + Index( + "idx_chk_deployment_thread", + "deployment_id", "thread_id", + ), + Index( + "idx_chk_thread_created", + "thread_id", "created_at", + ), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + thread_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True) + checkpoint_id: Mapped[str] = mapped_column(String(64), nullable=False) + parent_checkpoint_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + deployment_id: Mapped[str] = mapped_column( + ForeignKey("managed_deployments.id"), nullable=False, index=True, + ) + family_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True) + checkpoint_namespace: Mapped[str | None] = mapped_column(String(128), nullable=True) + node: Mapped[str | None] = mapped_column(String(128), nullable=True) + state_blob: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + pending_writes_json: Mapped[list[dict[str, Any]]] = mapped_column( + JSON, nullable=False, default=list, + ) + metadata_json: Mapped[dict[str, Any]] = mapped_column( + JSON, nullable=False, default=dict, + ) + blob_size_bytes: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, index=True, + ) + + +class ConversationThread(Base): + """Long-lived conversation thread anchor, mirroring ConsumerSessionState. + + Lets the orchestrator pin a multi-turn conversation to a single + ``thread_id`` (and therefore a single miner deployment) so + checkpoint resume works without per-turn pod affinity. + """ + + __tablename__ = "conversation_threads" + __table_args__ = ( + Index("idx_thread_user", "user_id"), + Index("idx_thread_deployment", "deployment_id"), + ) + + thread_id: Mapped[str] = mapped_column(String(64), primary_key=True) + user_id: Mapped[str | None] = mapped_column(String(128), nullable=True) + deployment_id: Mapped[str] = mapped_column( + ForeignKey("managed_deployments.id"), nullable=False, + ) + family_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True) + last_checkpoint_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + class MinerEvaluationSummary(Base): """Aggregated per-miner score computed from individual task results.""" __tablename__ = "miner_evaluation_summaries" @@ -794,29 +903,640 @@ class WorkflowRuntimePolicyStateRecord(Base): updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=False), default=utcnow, nullable=False) -class OwnerDatasetBinding(Base): - __tablename__ = "owner_dataset_bindings" +# --------------------------------------------------------------------------- +# Product runtime tables +# +# These power the consumer-facing path. Miner pods never see real user data; +# the product orchestrator loads from these tables on every turn, builds an +# AgentInvocationRequest, and forwards to a ServingDeployment (the promoted +# eval-winner's image). User state survives across promotions because it +# lives here, not in any miner deployment. +# --------------------------------------------------------------------------- - id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid4())) + +class ConsumerUser(Base): + """One real user of the product layer. + + ``auth_subject`` is whatever the consumer-api authorizer surfaces (today + an API-key principal; OAuth subject later). Distinct from the miner + ``hotkey`` namespace. + """ + + __tablename__ = "consumer_users" + + user_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + auth_subject: Mapped[str] = mapped_column(String(255), nullable=False, unique=True, index=True) + display_name: Mapped[str | None] = mapped_column(String(128), nullable=True) + metadata_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ConsumerProject(Base): + """A 'chat project' container — like ChatGPT projects. + + Custom instructions inject into ``metadata.project_context`` on every + turn. Project memory (vector recall over user-uploaded docs) is keyed + by ``project_id`` in :class:`ConsumerProjectMemory`. + """ + + __tablename__ = "consumer_projects" + __table_args__ = ( + Index("idx_consumer_project_user", "user_id"), + ) + + project_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + user_id: Mapped[str] = mapped_column( + ForeignKey("consumer_users.user_id"), nullable=False, index=True, + ) + name: Mapped[str] = mapped_column(String(128), nullable=False) + custom_instructions: Mapped[str | None] = mapped_column(Text, nullable=True) + default_family_id: Mapped[str] = mapped_column( + String(64), nullable=False, default="general_chat", + ) + metadata_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ConsumerConversation(Base): + """One user-visible thread. + + Distinct from the SDK's graph ``thread_id`` — the orchestrator + generates a fresh graph thread per turn and passes the user's + history via ``AgentInvocationRequest.history``. + """ + + __tablename__ = "consumer_conversations" + __table_args__ = ( + Index("idx_consumer_conv_user", "user_id"), + Index("idx_consumer_conv_project", "project_id"), + Index("idx_consumer_conv_user_updated", "user_id", "last_message_at"), + ) + + conversation_id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + user_id: Mapped[str] = mapped_column( + ForeignKey("consumer_users.user_id"), nullable=False, index=True, + ) + project_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_projects.project_id"), nullable=True, index=True, + ) + family_id: Mapped[str] = mapped_column( + String(64), nullable=False, default="general_chat", index=True, + ) + title: Mapped[str | None] = mapped_column(String(255), nullable=True) + last_message_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=False), nullable=True, index=True, + ) + # Rolling summary of the head of the conversation. Populated by + # ConversationSummarizer when the verbatim tail accumulates past a + # configurable threshold; injected as a system message into + # request.history so the agent sees compressed context for the head + # plus verbatim recent turns. NULL until first summarization. + rolling_summary: Mapped[str | None] = mapped_column(Text, nullable=True) + last_summarized_message_id: Mapped[str | None] = mapped_column( + String(36), nullable=True, + ) + metadata_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ConsumerMessage(Base): + """One user or assistant turn within a conversation. + + On every turn, the orchestrator loads the last N messages (default 20) + and folds them into ``request.history`` so the agent code never has to + persist user-visible state itself. + + ``served_by_*`` fields are an audit trail: which serving release / + deployment answered. Useful for cost reconciliation and for pulling a + user back to the same code if a promotion regresses. + """ + + __tablename__ = "consumer_messages" + __table_args__ = ( + UniqueConstraint( + "conversation_id", "turn_idx", + name="uq_consumer_message_conversation_turn", + ), + Index("idx_consumer_msg_conversation", "conversation_id", "turn_idx"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + conversation_id: Mapped[str] = mapped_column( + ForeignKey("consumer_conversations.conversation_id"), + nullable=False, index=True, + ) + turn_idx: Mapped[int] = mapped_column(Integer, nullable=False) + role: Mapped[str] = mapped_column(String(16), nullable=False) # user|assistant|system + content: Mapped[str] = mapped_column(Text, nullable=False, default="") + citations_json: Mapped[list[dict[str, Any]]] = mapped_column( + JSON, nullable=False, default=list, + ) + tool_calls_json: Mapped[list[dict[str, Any]]] = mapped_column( + JSON, nullable=False, default=list, + ) + served_by_deployment_id: Mapped[str | None] = mapped_column(String(36), nullable=True) + served_by_release_id: Mapped[str | None] = mapped_column(String(36), nullable=True) + latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + metadata_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ConsumerPreference(Base): + """User preference (tone, language, persona, default tools). + + Loaded into ``metadata.user_preferences`` on every turn. ``scope`` + distinguishes a global default from a per-project override. + """ + + __tablename__ = "consumer_preferences" + __table_args__ = ( + UniqueConstraint( + "user_id", "scope", "project_id", "key", + name="uq_consumer_pref_user_scope_proj_key", + ), + Index("idx_consumer_pref_user_scope", "user_id", "scope"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + user_id: Mapped[str] = mapped_column( + ForeignKey("consumer_users.user_id"), nullable=False, index=True, + ) + scope: Mapped[str] = mapped_column(String(16), nullable=False, default="global") # global|project + project_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_projects.project_id"), nullable=True, index=True, + ) + key: Mapped[str] = mapped_column(String(64), nullable=False) + value_json: Mapped[Any] = mapped_column(JSON, nullable=False) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ConsumerProjectMemory(Base): + """Per-project vector store for long-lived user docs / recall. + + Loaded as top-K into ``metadata.recalled_memory``. Ingestion (embedding + job over uploaded docs and over assistant outputs) is a follow-up; the + schema is wired now so the orchestrator can read from it as soon as + rows exist. + """ + + __tablename__ = "consumer_project_memory" + __table_args__ = ( + Index("idx_consumer_memory_project", "project_id"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + project_id: Mapped[str] = mapped_column( + ForeignKey("consumer_projects.project_id"), nullable=False, index=True, + ) + vector_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True) + embedding: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + text: Mapped[str] = mapped_column(Text, nullable=False, default="") + source_message_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_messages.id"), nullable=True, + ) + metadata_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ConsumerUserMemory(Base): + """Per-user vector store for stable facts and preferences. + + Distinct from :class:`ConsumerProjectMemory`. Project memory is + scoped to one project's history; user memory is the place for + durable facts about the user themselves ("works in Python", "prefers + concise answers", "lives in Tokyo") that should surface across + every project they touch. + + Populated by :class:`UserMemoryWriter` (gated regex pre-filter → + extractor LLM call → embedding). Loaded as top-K into + ``metadata.user_facts`` on every turn. + + The ``kind`` column is informational, not load-bearing — useful for + UI grouping and for retrieval policies that want to weight different + kinds of facts differently. + """ + + __tablename__ = "consumer_user_memory" + __table_args__ = ( + Index("idx_consumer_user_memory_user", "user_id"), + Index("idx_consumer_user_memory_user_vector", "user_id", "vector_id"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + user_id: Mapped[str] = mapped_column( + ForeignKey("consumer_users.user_id"), nullable=False, index=True, + ) + vector_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True) + embedding: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + text: Mapped[str] = mapped_column(Text, nullable=False, default="") + kind: Mapped[str] = mapped_column( + String(32), nullable=False, default="fact", # fact | preference | skill + ) + confidence: Mapped[float] = mapped_column(Float, nullable=False, default=1.0) + source_conversation_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_conversations.conversation_id"), nullable=True, + ) + source_message_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_messages.id"), nullable=True, + ) + metadata_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ServingPromotion(Base): + """Audit trail: which eval-winner became today's serving agent. + + One row per (family_id, run_id). Idempotency for the promotion job + is keyed on the unique constraint below. + """ + + __tablename__ = "serving_promotions" + __table_args__ = ( + UniqueConstraint( + "family_id", "run_id", + name="uq_serving_promotion_family_run", + ), + Index("idx_serving_promotion_release", "serving_release_id"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) family_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True) run_id: Mapped[str] = mapped_column(String(128), nullable=False, index=True) - bundle_uri: Mapped[str] = mapped_column(String(1024), nullable=False) - bundle_sha256: Mapped[str] = mapped_column(String(64), nullable=False) - generator_version: Mapped[str] = mapped_column(String(128), nullable=False) - generated_by: Mapped[str] = mapped_column(String(128), nullable=False) - signature_hex: Mapped[str] = mapped_column(String(256), nullable=False) - generator_provider: Mapped[str] = mapped_column(String(64), nullable=False, default="") - generator_model: Mapped[str] = mapped_column(String(128), nullable=False, default="") - status: Mapped[str] = mapped_column(String(32), nullable=False, default="pending", index=True) - provenance_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict) + source_deployment_id: Mapped[str] = mapped_column( + ForeignKey("managed_deployments.id"), nullable=False, index=True, + ) + serving_release_id: Mapped[str] = mapped_column( + ForeignKey("serving_releases.id"), nullable=False, index=True, + ) + promoted_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + metadata_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict) + + +class ConsumerAttachment(Base): + """One uploaded file (PDF / DOCX / CSV / TXT / etc.) for product chat. + + Matches the ChatGPT / Claude pattern: the orchestrator preprocesses + uploads into extracted text, the agent never sees a "file upload + tool" — it just sees ``metadata.attached_files`` on the next chat + turn. This row is the canonical record of what was uploaded and + what came out of extraction. + + ``blob_ref`` is opaque storage location (S3 key, local path, etc.) + so we don't dump the raw bytes into Postgres. ``extracted_text`` + is the LLM-ready content; ``extraction_metadata_json`` carries + page count / row count / etc. for downstream debugging. + """ + + __tablename__ = "consumer_attachments" + __table_args__ = ( + Index("idx_consumer_attachment_user", "user_id"), + Index("idx_consumer_attachment_conversation", "conversation_id"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + user_id: Mapped[str] = mapped_column( + ForeignKey("consumer_users.user_id"), nullable=False, index=True, + ) + conversation_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_conversations.conversation_id"), + nullable=True, index=True, + ) + message_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_messages.id"), nullable=True, + ) + filename: Mapped[str] = mapped_column(String(512), nullable=False) + content_type: Mapped[str] = mapped_column(String(128), nullable=False) + size_bytes: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + blob_ref: Mapped[str | None] = mapped_column(String(1024), nullable=True) + extracted_text: Mapped[str] = mapped_column(Text, nullable=False, default="") + extraction_metadata_json: Mapped[dict[str, Any]] = mapped_column( + JSON, nullable=False, default=dict, + ) + extraction_status: Mapped[str] = mapped_column( + String(32), nullable=False, default="ok", + ) # ok | unsupported | failed | truncated created_at: Mapped[datetime] = mapped_column( - DateTime(timezone=False), default=utcnow, nullable=False + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class McpIntegration(Base): + """Operator-curated MCP server in the catalog. + + Consumers connect to integrations *from this catalog*, never to + arbitrary URLs. The subnet operator vets each integration once + (auth scopes, capability surface, abuse posture) before adding it; + consumers see only the active ones. + + ``capabilities_hash`` is sha256 over the canonicalized list of + tools the integration declared at registration time. The relay + service rejects calls when an integration's live capability set has + drifted from the stored hash — operator must run reprobe to refresh + before consumers can use the new surface. + """ + + __tablename__ = "mcp_integrations" + __table_args__ = ( + Index("idx_mcp_integration_status", "status"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + slug: Mapped[str] = mapped_column( + String(64), nullable=False, unique=True, index=True, + ) + display_name: Mapped[str] = mapped_column(String(128), nullable=False) + vendor: Mapped[str] = mapped_column(String(128), nullable=False, default="") + description: Mapped[str] = mapped_column(Text, nullable=False, default="") + base_url: Mapped[str] = mapped_column(String(1024), nullable=False) + transport: Mapped[str] = mapped_column( + String(32), nullable=False, default="http", + ) # http | sse | stdio_via_relay + capabilities_json: Mapped[list[dict[str, Any]]] = mapped_column( + JSON, nullable=False, default=list, + ) + capabilities_hash: Mapped[str] = mapped_column( + String(64), nullable=False, default="", + ) + oauth_provider: Mapped[str | None] = mapped_column( + String(64), nullable=True, + ) + oauth_authorize_url: Mapped[str | None] = mapped_column( + String(1024), nullable=True, + ) + oauth_token_url: Mapped[str | None] = mapped_column( + String(1024), nullable=True, + ) + oauth_scopes_json: Mapped[list[str]] = mapped_column( + JSON, nullable=False, default=list, + ) + status: Mapped[str] = mapped_column( + String(16), nullable=False, default="active", index=True, + ) # active | disabled + created_by_admin_id: Mapped[str | None] = mapped_column( + String(128), nullable=True, ) - activated_at: Mapped[datetime | None] = mapped_column( - DateTime(timezone=False), nullable=True + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ConsumerMcpConnection(Base): + """One user's connection to a catalog integration. + + Stores encrypted OAuth tokens keyed on (user_id, integration_id). + Decryption happens only inside the orchestrator process; tokens + never leave the orchestrator boundary. The miner pod sees zero + references to these rows — the orchestrator runs MCP tool calls on + the user's behalf and injects the results into envelope metadata. + """ + + __tablename__ = "consumer_mcp_connections" + __table_args__ = ( + UniqueConstraint( + "user_id", "integration_id", + name="uq_consumer_mcp_user_integration", + ), + Index("idx_consumer_mcp_user", "user_id"), + Index("idx_consumer_mcp_integration", "integration_id"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + user_id: Mapped[str] = mapped_column( + ForeignKey("consumer_users.user_id"), nullable=False, index=True, + ) + integration_id: Mapped[str] = mapped_column( + ForeignKey("mcp_integrations.id"), nullable=False, index=True, + ) + oauth_access_token_encrypted: Mapped[bytes | None] = mapped_column( + LargeBinary, nullable=True, + ) + oauth_refresh_token_encrypted: Mapped[bytes | None] = mapped_column( + LargeBinary, nullable=True, + ) + oauth_expires_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=False), nullable=True, + ) + status: Mapped[str] = mapped_column( + String(16), nullable=False, default="active", index=True, + ) # active | expired | revoked + last_used_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=False), nullable=True, + ) + metadata_json: Mapped[dict[str, Any]] = mapped_column( + JSON, nullable=False, default=dict, + ) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class ConsumerMcpToolCall(Base): + """Audit log of every MCP tool call brokered for a consumer. + + Stores per-call cost so :class:`RunCostTracker` can attribute spend. + Result is digested (truncated string) — full result lives only in + transit; we don't persist user-visible payloads in case they + contain PII the orchestrator's safety pipeline didn't catch. + """ + + __tablename__ = "consumer_mcp_tool_calls" + __table_args__ = ( + Index("idx_consumer_mcp_call_conversation", "conversation_id"), + Index("idx_consumer_mcp_call_connection", "connection_id"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), ) + conversation_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_conversations.conversation_id"), + nullable=True, index=True, + ) + message_id: Mapped[str | None] = mapped_column( + ForeignKey("consumer_messages.id"), nullable=True, + ) + connection_id: Mapped[str] = mapped_column( + ForeignKey("consumer_mcp_connections.id"), nullable=False, index=True, + ) + tool_name: Mapped[str] = mapped_column(String(128), nullable=False) + args_json: Mapped[dict[str, Any]] = mapped_column( + JSON, nullable=False, default=dict, + ) + result_digest: Mapped[str] = mapped_column(Text, nullable=False, default="") + latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + error: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class OrchestratorToolCallLog(Base): + """Server-attested ledger of every tool call brokered by the subnet. + + Each row is written by the tool service that actually executed the + call (web_search, url_fetch, sandbox; the MCP relay writes a row in + addition to its consumer-MCP audit row). The ledger is the + authoritative source the eval pipeline reads to compute tool-use + KPIs — miner-emitted trace frames carry zero scoring weight. + + Indexed by ``job_id`` so an eval run can fetch every tool call its + items triggered. Per-row write latency must stay low; tool services + use a fire-and-forget write-behind buffer rather than blocking the + response on the DB write. + """ + + __tablename__ = "orchestrator_tool_call_log" + __table_args__ = ( + Index("idx_otcl_job_id", "job_id"), + Index("idx_otcl_job_tool", "job_id", "tool_name"), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + job_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True) + tool_name: Mapped[str] = mapped_column(String(64), nullable=False) + args_hash: Mapped[str] = mapped_column(String(64), nullable=False, default="") + args_json: Mapped[dict[str, Any]] = mapped_column( + JSON, nullable=False, default=dict, + ) + result_digest: Mapped[str] = mapped_column( + Text, nullable=False, default="", + ) + latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0) + status: Mapped[str] = mapped_column(String(16), nullable=False, default="ok") + error: Mapped[str | None] = mapped_column(Text, nullable=True) + ts: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, + ) + + +class EvalFeedback(Base): + """Per-(run, miner, task) eval outcome row. + + The validator writes one row per ``_judge_miner`` invocation + after computing the EvalJudge outcome + composite. Eiretes' + ``GET /v1/eval/feedback`` (hotkey-signed, exposed to miners) + aggregates rows for ``(run_id, miner_hotkey)`` into an + ``EvalFeedbackDoc`` with categorical guidance per item — the + miner sees ``failure_mode`` + ``guidance`` per task plus the + largest-gap summary, but never the verbatim expected_claims. + + Indexed on ``(miner_hotkey, run_id)`` for the per-miner read + path and ``(run_id, task_id)`` for cross-miner per-task drill-in + on the operator dashboard. + """ + __tablename__ = "eval_feedback" __table_args__ = ( - UniqueConstraint("family_id", "run_id", name="uq_owner_dataset_bindings_family_run"), - Index("idx_owner_dataset_bindings_family_status", "family_id", "status"), + Index("idx_eval_feedback_miner_run", "miner_hotkey", "run_id"), + Index("idx_eval_feedback_run_task", "run_id", "task_id"), + UniqueConstraint( + "run_id", "miner_hotkey", "task_id", + name="uq_eval_feedback_run_miner_task", + ), + ) + + id: Mapped[str] = mapped_column( + String(36), primary_key=True, default=lambda: str(uuid4()), + ) + run_id: Mapped[str] = mapped_column(String(64), nullable=False) + miner_hotkey: Mapped[str] = mapped_column(String(64), nullable=False) + task_id: Mapped[str] = mapped_column(String(64), nullable=False) + # EvalJudge output: outcome ∈ {correct, partial, wrong, hallucinated, + # refused, disputed}; failure_mode is one of the categorical buckets + # in eiretes/eval/models.py:FailureMode (or null when outcome=correct). + outcome: Mapped[str] = mapped_column(String(32), nullable=False) + failure_mode: Mapped[str | None] = mapped_column( + String(64), nullable=True, + ) + # Categorical hint (≤200 chars) the miner can act on — never + # quotes ``expected_claims`` verbatim. + guidance: Mapped[str] = mapped_column(Text, nullable=False, default="") + # Substring previews surfaced in the per-miner doc; full prompt / + # response stay on TaskMinerResult. ``prompt_excerpt`` ≤200 chars, + # ``response_excerpt`` ≤500 chars by convention (validator + # truncates before the POST). + prompt_excerpt: Mapped[str] = mapped_column( + Text, nullable=False, default="", + ) + response_excerpt: Mapped[str] = mapped_column( + Text, nullable=False, default="", + ) + composite_score: Mapped[float] = mapped_column( + Float, nullable=False, default=0.0, + ) + # Why composite was zeroed (if it was) — e.g. + # "tool_attestation_factor=0 because required_tool=web_search but + # ledger had no calls". JSON list so the dashboard can render + # one-per-line; empty for happy-path rows. + knockout_reasons_json: Mapped[list[str]] = mapped_column( + JSON, nullable=False, default=list, + ) + # Validator-side reconciler verdict on the 3 oracles + # (consensus / majority / disputed) or "deterministic" for non- + # three_oracle items. Lets the miner-facing doc surface + # "this item was disputed by oracles" so they don't chase a + # phantom failure. + oracle_status: Mapped[str | None] = mapped_column( + String(32), nullable=True, + ) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=False), default=utcnow, nullable=False, ) From 4746aabcdb87c7884916b50a40eed179c82a9b13 Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Thu, 7 May 2026 17:30:39 +0000 Subject: [PATCH 11/24] feat(shared): add safety pipeline + multi-metric scoring + tool pricing table --- shared/common/tool_pricing.py | 6 +- shared/safety/__init__.py | 42 +++++ shared/safety/guard.py | 97 +++++++++++ shared/safety/pii_redaction.py | 184 +++++++++++++++++++++ shared/safety/prompt_injection.py | 186 +++++++++++++++++++++ shared/safety/token_encryption.py | 110 +++++++++++++ shared/scoring/multi_metric.py | 259 ++++++++++++++++++++++++++++++ 7 files changed, 883 insertions(+), 1 deletion(-) create mode 100644 shared/safety/__init__.py create mode 100644 shared/safety/guard.py create mode 100644 shared/safety/pii_redaction.py create mode 100644 shared/safety/prompt_injection.py create mode 100644 shared/safety/token_encryption.py create mode 100644 shared/scoring/multi_metric.py diff --git a/shared/common/tool_pricing.py b/shared/common/tool_pricing.py index e8b6805..b4916a7 100644 --- a/shared/common/tool_pricing.py +++ b/shared/common/tool_pricing.py @@ -18,8 +18,12 @@ class ToolPrice: # general_chat 4D scorer to compute the cost dimension. TOOL_PRICING: dict[str, ToolPrice] = { "web_search": ToolPrice(per_call_usd=0.001), - "x_api": ToolPrice(per_call_usd=0.050), "sandbox": ToolPrice(per_call_usd=0.002), + "url_fetch": ToolPrice(per_call_usd=0.0005), + # text-embedding-3-small averages ~$0.0002 per query (one + # ~600-token query embed + amortized index cost). Indexing is + # operator-paid, retrieval is the per-call cost the miner sees. + "rag.retrieve": ToolPrice(per_call_usd=0.0002), } diff --git a/shared/safety/__init__.py b/shared/safety/__init__.py new file mode 100644 index 0000000..f399deb --- /dev/null +++ b/shared/safety/__init__.py @@ -0,0 +1,42 @@ +"""Orchestrator-boundary safety guards. + +Two thin guards run by default at the :class:`ProductOrchestrator` +boundary: + + * :class:`PIIRedactionGuard` — regex SSN / CC / email / phone + redaction on the inbound prompt and outbound assistant content. + * :class:`PromptInjectionGuard` — regex denylist (with an optional + LLM classifier escalation) on the inbound prompt. + +The guards live at the orchestrator boundary, NOT inside the miner +graph. Same isolation principle as user data in eval mode: hardening +is a product-layer responsibility, miners compete on reasoning over +already-clean inputs. + +The :class:`OrchestratorGuard` ABC is intentionally separate from +:class:`eirel.safety.Guard`. The SDK guard signature is shaped around +graph state mappings; the orchestrator guard takes a plain string +prompt or content. Mixing them via shimming creates more friction than +it saves. +""" +from __future__ import annotations + +from shared.safety.guard import ( + GuardVerdict, + OrchestratorGuard, + Redaction, +) +from shared.safety.pii_redaction import PIIRedactionGuard +from shared.safety.prompt_injection import ( + PromptInjectionClassifier, + PromptInjectionGuard, +) + +__all__ = [ + "GuardVerdict", + "OrchestratorGuard", + "PIIRedactionGuard", + "PromptInjectionClassifier", + "PromptInjectionGuard", + "Redaction", +] diff --git a/shared/safety/guard.py b/shared/safety/guard.py new file mode 100644 index 0000000..fcce5f4 --- /dev/null +++ b/shared/safety/guard.py @@ -0,0 +1,97 @@ +"""Orchestrator-side guard ABC. + +Distinct from :class:`eirel.safety.Guard` (which scans graph state +inside a miner pod). This guard runs at the +:class:`ProductOrchestrator` boundary on plain strings and is shaped +for two product-layer concerns: rejecting prompt-injection attempts +before the miner sees them, and redacting PII from both inbound user +prompts and outbound assistant content. +""" +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +__all__ = ["GuardVerdict", "OrchestratorGuard", "Redaction"] + + +@dataclass(frozen=True, slots=True) +class Redaction: + """One PII match the guard wants the orchestrator to mask. + + ``replacement`` is the substitution token (e.g. ``"[REDACTED-EMAIL]"``). + ``span`` is the original ``(start, end)`` byte offsets — useful for + audit logs where the operator needs to know what was masked + without storing the raw value. + """ + + kind: str + replacement: str + span: tuple[int, int] + + +@dataclass(frozen=True, slots=True) +class GuardVerdict: + """Outcome of one guard call at the orchestrator boundary. + + ``allow=False`` short-circuits the turn — the orchestrator returns + a canned refusal message and never calls the miner. ``redacted_text`` + (when non-None) replaces the original prompt or content; the miner + sees the redacted version, not the raw one. + """ + + allow: bool + reason: str | None = None + redacted_text: str | None = None + redactions: tuple[Redaction, ...] = () + metadata: Mapping[str, Any] = field(default_factory=dict) + + @classmethod + def ok(cls, **metadata: Any) -> "GuardVerdict": + return cls(allow=True, metadata=dict(metadata) if metadata else {}) + + @classmethod + def deny(cls, reason: str, **metadata: Any) -> "GuardVerdict": + return cls( + allow=False, reason=reason, + metadata=dict(metadata) if metadata else {}, + ) + + +class OrchestratorGuard(ABC): + """ABC for guards wired into :class:`SafetyPipeline`. + + Two-call contract: + + * :meth:`pre_input` runs before the orchestrator builds the + envelope. Receives the raw user prompt + a context dict + (``user_id``, ``conversation_id``, ``project_id`` when + present). May return a redacted prompt; an ``allow=False`` + verdict aborts the turn. + * :meth:`post_output` runs after the miner replies, before the + assistant turn is persisted. Receives the assistant content + + the same context shape. Same redaction semantics. + + Implementations MUST be pure-async and side-effect free except for + their own scanner calls. The pipeline applies the redacted text + sequentially through the chain so multiple guards can stack. + """ + + @abstractmethod + async def pre_input( + self, + text: str, + context: Mapping[str, Any], + ) -> GuardVerdict: ... + + @abstractmethod + async def post_output( + self, + text: str, + context: Mapping[str, Any], + ) -> GuardVerdict: ... + + async def aclose(self) -> None: + return None diff --git a/shared/safety/pii_redaction.py b/shared/safety/pii_redaction.py new file mode 100644 index 0000000..1a41d8f --- /dev/null +++ b/shared/safety/pii_redaction.py @@ -0,0 +1,184 @@ +"""Regex PII redaction for inbound prompts and outbound content. + +Matches four high-frequency leak vectors: + + * **email** — RFC-5322-ish; intentionally tighter than the spec to + avoid grabbing arbitrary at-signs in code blocks. + * **phone** — North American formats (``(123) 456-7890``, + ``123-456-7890``, ``+1 555 123 4567``) plus generic 10–15 digit + runs with country-code prefixes. Not bulletproof internationally; + the goal is to mask plausible matches, not all possible numbers. + * **SSN** — US format ``DDD-DD-DDDD``, no other variants. + * **credit card** — 13–19 contiguous digits with optional spaces or + dashes; validated with the Luhn checksum so plain digit strings + that aren't payment cards (order numbers, IDs) don't get masked. + +False positives are cheap (one masked token in a chat reply); false +negatives are expensive (PII leaks downstream). The guard tilts toward +catching too much, not too little. + +The regexes deliberately do not chase obscure formats — production +deployments wanting Microsoft Presidio or AWS Comprehend integration +implement another :class:`OrchestratorGuard` subclass and stack it via +the :class:`SafetyPipeline`. +""" +from __future__ import annotations + +import re +from collections.abc import Mapping +from typing import Any + +from shared.safety.guard import GuardVerdict, OrchestratorGuard, Redaction + +__all__ = ["PIIRedactionGuard"] + + +# -- Patterns ---------------------------------------------------------------- +# +# Order matters: more-specific patterns first so they win on overlap. + +_EMAIL_RE = re.compile( + r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}", +) + +_SSN_RE = re.compile(r"\b\d{3}-\d{2}-\d{4}\b") + +# Credit card: 13–19 digits with optional space/dash separators between +# 3–4 digit groups. Validated with Luhn below, so non-card digit runs +# don't trip. +_CC_RE = re.compile( + r"(? bool: + """Return ``True`` when ``digits`` (no separators) passes Luhn.""" + if not digits.isdigit() or len(digits) < 13 or len(digits) > 19: + return False + total = 0 + odd = True # we walk right-to-left + for ch in reversed(digits): + d = ord(ch) - 48 + if odd: + total += d + else: + doubled = d * 2 + total += doubled if doubled < 10 else doubled - 9 + odd = not odd + return total % 10 == 0 + + +def _scan(text: str) -> list[tuple[str, int, int, str]]: + """Return ``[(kind, start, end, replacement), ...]`` non-overlapping. + + Kinds in priority order: email, ssn, cc, phone. Once a span is + masked, later patterns won't match inside it. + """ + if not text: + return [] + matches: list[tuple[str, int, int, str]] = [] + masked: list[bool] = [False] * len(text) + + def _claim(s: int, e: int) -> bool: + if any(masked[s:e]): + return False + for i in range(s, e): + masked[i] = True + return True + + for m in _EMAIL_RE.finditer(text): + if _claim(m.start(), m.end()): + matches.append(("email", m.start(), m.end(), "[REDACTED-EMAIL]")) + for m in _SSN_RE.finditer(text): + if _claim(m.start(), m.end()): + matches.append(("ssn", m.start(), m.end(), "[REDACTED-SSN]")) + for m in _CC_RE.finditer(text): + digits = re.sub(r"[^\d]", "", m.group(0)) + if not _luhn(digits): + continue + if _claim(m.start(), m.end()): + matches.append(("credit_card", m.start(), m.end(), "[REDACTED-CC]")) + for m in _PHONE_RE.finditer(text): + if _claim(m.start(), m.end()): + matches.append(("phone", m.start(), m.end(), "[REDACTED-PHONE]")) + + matches.sort(key=lambda x: x[1]) + return matches + + +def _apply(text: str, matches: list[tuple[str, int, int, str]]) -> str: + """Splice the replacements into ``text`` left-to-right.""" + if not matches: + return text + out: list[str] = [] + cursor = 0 + for _kind, start, end, replacement in matches: + out.append(text[cursor:start]) + out.append(replacement) + cursor = end + out.append(text[cursor:]) + return "".join(out) + + +class PIIRedactionGuard(OrchestratorGuard): + """Mask SSN / CC / email / phone in both inbound and outbound text. + + Always allows the turn — the deny path is reserved for actively + hostile inputs (prompt injection). PII redaction is non-blocking + by design; the user might *legitimately* be asking the model to + parse a contact card. Mask, don't reject. + + Set ``redact_input=False`` or ``redact_output=False`` to keep the + guard one-sided; useful for staging environments where the operator + wants to see raw inbound text in logs but still mask outbound. + """ + + __slots__ = ("_redact_input", "_redact_output") + + def __init__( + self, + *, + redact_input: bool = True, + redact_output: bool = True, + ) -> None: + self._redact_input = redact_input + self._redact_output = redact_output + + async def pre_input( + self, text: str, context: Mapping[str, Any] + ) -> GuardVerdict: + if not self._redact_input: + return GuardVerdict.ok() + return self._verdict(text) + + async def post_output( + self, text: str, context: Mapping[str, Any] + ) -> GuardVerdict: + if not self._redact_output: + return GuardVerdict.ok() + return self._verdict(text) + + def _verdict(self, text: str) -> GuardVerdict: + matches = _scan(text) + if not matches: + return GuardVerdict.ok() + redacted = _apply(text, matches) + redactions = tuple( + Redaction(kind=kind, replacement=replacement, span=(start, end)) + for kind, start, end, replacement in matches + ) + return GuardVerdict( + allow=True, + redacted_text=redacted, + redactions=redactions, + metadata={"pii_kinds": sorted({r.kind for r in redactions})}, + ) diff --git a/shared/safety/prompt_injection.py b/shared/safety/prompt_injection.py new file mode 100644 index 0000000..b9a5979 --- /dev/null +++ b/shared/safety/prompt_injection.py @@ -0,0 +1,186 @@ +"""Prompt-injection guard for the orchestrator boundary. + +Two-layer detection: + + 1. **Regex denylist** (always on). Catches the well-trodden patterns — + "ignore previous instructions", "you are now DAN", system-impersonation + headers, etc. False-positive rate matters; the denylist tilts toward + specificity over recall, since a misfire blocks a real chat turn. + + 2. **LLM classifier** (optional). When the regex layer abstains and a + classifier is wired in, the guard escalates the prompt to the + classifier and respects its verdict. Skipped entirely without a + classifier — keeps the cheap path zero-cost. + +A deny verdict short-circuits the turn at the orchestrator. The miner +pod never sees the prompt, no envelope is built, and the consumer gets +a canned refusal with the matched-rule name in metadata for audit. +""" +from __future__ import annotations + +import logging +import re +from collections.abc import Mapping +from typing import Any, Protocol + +from shared.safety.guard import GuardVerdict, OrchestratorGuard + +_logger = logging.getLogger(__name__) + +__all__ = ["PromptInjectionClassifier", "PromptInjectionGuard"] + + +# -- Denylist regexes ------------------------------------------------------- +# +# Each entry is (rule_name, compiled_pattern). Order is informational; +# the first match wins for the metadata.matched_rule field. + +_DENYLIST: tuple[tuple[str, re.Pattern[str]], ...] = ( + ( + "ignore_previous_instructions", + re.compile( + r"\b(?:please\s+)?(?:ignore|disregard|forget|bypass|override)" + r"\s+(?:all|any|the)?\s*(?:prior|previous|preceding|above|earlier|" + r"prior\s+system|system|safety)\s+" + r"(?:instruction|instructions|prompt|prompts|rule|rules|" + r"guideline|guidelines|directive|directives)\b", + re.IGNORECASE | re.DOTALL, + ), + ), + ( + "system_impersonation", + re.compile( + r"^\s*(?:###\s*)?(?:\[?\s*)?system\s*[:>\]]" + r"|" + r"<\|im_start\|>\s*system", + re.IGNORECASE | re.MULTILINE, + ), + ), + ( + "role_hijack_dan", + re.compile( + r"\b(?:you\s+are\s+now\s+)?(?:DAN|do\s+anything\s+now|" + r"developer\s+mode|jailbreak\s+mode|unrestricted\s+mode)\b", + re.IGNORECASE, + ), + ), + ( + "reveal_system_prompt", + re.compile( + r"\b(?:reveal|show|print|leak|repeat|recite|output)\s+" + r"(?:your|the|all)?\s*" + r"(?:system|hidden|original|initial|raw)\s+" + r"(?:prompt|instructions|message|context)\b", + re.IGNORECASE | re.DOTALL, + ), + ), + ( + "instruction_override", + re.compile( + r"\bnew\s+instructions?:\s*", + re.IGNORECASE, + ), + ), +) + + +class PromptInjectionClassifier(Protocol): + """Optional escalation path for prompts the regex layer didn't catch. + + Implementations call out to a model (or any heuristic). They MUST + return a (deny, reason) pair. Failures should raise; the guard + catches and logs them as ``allow`` so the chat turn isn't broken + by a flaky classifier. + """ + + async def classify(self, text: str) -> tuple[bool, str | None]: ... + + +class PromptInjectionGuard(OrchestratorGuard): + """Two-layer guard: regex denylist + optional classifier escalation. + + Parameters + ---------- + classifier + Optional :class:`PromptInjectionClassifier`. When ``None``, the + guard runs the regex layer only. + enable_classifier + Master switch for the classifier (independent of whether one is + passed). Lets operators ship the classifier off by default and + flip it on per-deployment via env without changing wiring. + apply_to_output + When ``True``, the guard also runs at :meth:`post_output`. + Defaults to ``False`` — assistant content tripping the denylist + usually means the model is parroting the user's prompt back, not + an injection attempt; the input-side block already handled it. + """ + + __slots__ = ("_classifier", "_enable_classifier", "_apply_to_output") + + def __init__( + self, + *, + classifier: PromptInjectionClassifier | None = None, + enable_classifier: bool = True, + apply_to_output: bool = False, + ) -> None: + self._classifier = classifier + self._enable_classifier = enable_classifier + self._apply_to_output = apply_to_output + + async def pre_input( + self, text: str, context: Mapping[str, Any] + ) -> GuardVerdict: + return await self._evaluate(text, stage="pre_input") + + async def post_output( + self, text: str, context: Mapping[str, Any] + ) -> GuardVerdict: + if not self._apply_to_output: + return GuardVerdict.ok() + return await self._evaluate(text, stage="post_output") + + async def _evaluate(self, text: str, *, stage: str) -> GuardVerdict: + if not text: + return GuardVerdict.ok() + rule = self._regex_match(text) + if rule is not None: + return GuardVerdict( + allow=False, + reason=f"prompt_injection: {rule}", + metadata={ + "layer": "regex", + "matched_rule": rule, + "stage": stage, + }, + ) + if self._classifier is not None and self._enable_classifier: + try: + deny, reason = await self._classifier.classify(text) + except Exception as exc: # noqa: BLE001 — best-effort + _logger.warning( + "prompt-injection classifier failed: %s", exc, + ) + return GuardVerdict( + allow=True, + metadata={"layer": "classifier_error"}, + ) + if deny: + return GuardVerdict( + allow=False, + reason=f"prompt_injection: {reason or 'classifier'}", + metadata={ + "layer": "classifier", + "stage": stage, + }, + ) + return GuardVerdict( + allow=True, + metadata={"layer": "regex", "stage": stage}, + ) + + def _regex_match(self, text: str) -> str | None: + for name, pattern in _DENYLIST: + if pattern.search(text): + return name + return None diff --git a/shared/safety/token_encryption.py b/shared/safety/token_encryption.py new file mode 100644 index 0000000..ce68ced --- /dev/null +++ b/shared/safety/token_encryption.py @@ -0,0 +1,110 @@ +"""Symmetric encryption for OAuth tokens at rest. + +OAuth tokens for consumer MCP connections are encrypted in the DB and +decrypted only inside the orchestrator process. Tokens never leave +the orchestrator boundary, never log in plaintext, and never reach a +miner pod. + +The implementation uses :mod:`cryptography.fernet` when the dependency +is available (production); a deterministic XOR-based fallback is wired +in for tests and dev environments where ``cryptography`` may not be +installed. The fallback is **not secure** — it only obscures plaintext +in logs / DB dumps. Production deployments MUST set +``EIREL_MCP_TOKEN_ENCRYPTION_KEY`` to a Fernet key (``Fernet.generate_key()``). +""" +from __future__ import annotations + +import base64 +import hashlib +import os +from typing import Any + +__all__ = [ + "TokenCipher", + "build_token_cipher", +] + + +class TokenCipher: + """Encrypt / decrypt small byte strings. + + Concrete implementation chosen at construction time: + + * Production: ``cryptography.fernet.Fernet`` if the lib is present + and ``EIREL_MCP_TOKEN_ENCRYPTION_KEY`` resolves to a valid key. + * Fallback: XOR-keystream over a sha256 expansion. Bytes change, + plaintext is not visible at a glance, but this is NOT + cryptographically strong — it's a marker for dev/test only. + + Inputs and outputs are bytes. Callers (the connection writer / the + relay service) are responsible for storing and retrieving the + blob unchanged. + """ + + __slots__ = ("_fernet", "_xor_key") + + def __init__( + self, + *, + fernet: Any | None = None, + xor_key: bytes | None = None, + ) -> None: + if fernet is None and xor_key is None: + raise ValueError("TokenCipher requires fernet or xor_key") + self._fernet = fernet + self._xor_key = xor_key + + @property + def is_secure(self) -> bool: + return self._fernet is not None + + def encrypt(self, plaintext: bytes) -> bytes: + if self._fernet is not None: + return self._fernet.encrypt(plaintext) + return _xor(plaintext, self._xor_key or b"") + + def decrypt(self, ciphertext: bytes) -> bytes: + if self._fernet is not None: + return self._fernet.decrypt(ciphertext) + return _xor(ciphertext, self._xor_key or b"") + + +def _expand_key(seed: bytes, length: int) -> bytes: + """Stretch ``seed`` to ``length`` bytes via repeated SHA-256.""" + out = bytearray() + counter = 0 + while len(out) < length: + digest = hashlib.sha256(seed + counter.to_bytes(8, "big")).digest() + out.extend(digest) + counter += 1 + return bytes(out[:length]) + + +def _xor(data: bytes, key_seed: bytes) -> bytes: + if not data: + return data + keystream = _expand_key(key_seed, len(data)) + return bytes(b ^ k for b, k in zip(data, keystream)) + + +def build_token_cipher() -> TokenCipher: + """Build a cipher from the env, falling back to XOR for dev. + + Reads ``EIREL_MCP_TOKEN_ENCRYPTION_KEY``: + * Looks like a Fernet key (44 url-safe base64 chars) → Fernet. + * Anything else → XOR fallback keyed on the value (or a static + dev key when the env var is unset). + """ + raw = os.getenv("EIREL_MCP_TOKEN_ENCRYPTION_KEY", "").encode("utf-8") + if raw: + try: + from cryptography.fernet import Fernet # type: ignore[import-not-found] + except ImportError: + return TokenCipher(xor_key=raw) + try: + return TokenCipher(fernet=Fernet(raw)) + except Exception: # noqa: BLE001 — bad key shape, degrade + return TokenCipher(xor_key=raw) + # Dev / test default — deterministic so two test runs roundtrip + # without surprises but with a stable plaintext-bytes guarantee. + return TokenCipher(xor_key=b"eirel-mcp-dev-key") diff --git a/shared/scoring/multi_metric.py b/shared/scoring/multi_metric.py new file mode 100644 index 0000000..0d3c8eb --- /dev/null +++ b/shared/scoring/multi_metric.py @@ -0,0 +1,259 @@ +"""Multi-metric per-task scoring for the general_chat eval. + +Pure-function helpers that the validator engine uses to compute per-task +scores from (task_definition, miner_response, baseline_response, judge +outputs, tool_call_ledger, latency, cost). No I/O, no LLM calls — those +are upstream. This module is the math. + +Per-task formula: + + task_score = w_p · pairwise_preference_score + + w_g · grounded_correctness + + w_r · retrieval_quality (or computation_correctness) + + w_t · tool_routing + + w_s · instruction_safety + + w_l · latency_cost + +Where the weights re-normalize over whichever dimensions are +``applicable`` for the task type (an N/A dimension drops out of the +sum and shifts its weight to the others proportionally). +""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +# -- Task type taxonomy ---------------------------------------------------- + + +# Bundle category (as published by the eirel-eval-pool renderer) → task +# type used for scoring. Tasks without a recognized category fall back +# to ``no_tool`` so they at least score the always-applicable metrics. +_CATEGORY_TO_TASK_TYPE: dict[str, str] = { + "live_lookup": "web_required", + "live_synthesis": "web_required", + # ``attached_long_doc`` ships the document in-context via + # ``metadata.attached_files``; the agent answers from context and + # no retrieval tool is needed. ``rag_required`` (below) is the + # category for genuine RAG eval where the corpus is indexed + # server-side and the agent must call ``rag.retrieve``. + "attached_long_doc": "no_tool", + "rag_required": "rag_required", + "multi_turn_agentic_memory": "memory_required", + "compute_or_orchestrate": "sandbox_required", + "abstention_probe": "no_tool", +} + + +def derive_task_type(category: str | None) -> str: + return _CATEGORY_TO_TASK_TYPE.get(str(category or "").strip(), "no_tool") + + +# -- Default weights per task type ----------------------------------------- + + +# All six dimensions are listed for every task type; non-applicable +# entries are mapped to ``None`` so the score collector knows to skip +# them. ``computation_correctness`` substitutes for ``retrieval_quality`` +# on sandbox tasks; the keys live in the same slot in the dict so the +# re-normalizer treats them as one tile. +_DEFAULT_WEIGHTS: dict[str, dict[str, float]] = { + "web_required": { + "pairwise_preference_score": 0.40, + "grounded_correctness": 0.30, + "retrieval_quality": 0.15, + "tool_routing": 0.05, + "instruction_safety": 0.05, + "latency_cost": 0.05, + }, + "rag_required": { + # retrieval_quality is intentionally absent here — the + # multi-judge LLM still scores it from the response text rather + # than from a gold chunk-id ledger. Re-normalization redistributes + # its share of weight across the remaining dimensions. + "pairwise_preference_score": 0.40, + "grounded_correctness": 0.30, + "tool_routing": 0.05, + "instruction_safety": 0.05, + "latency_cost": 0.05, + }, + "sandbox_required": { + "pairwise_preference_score": 0.40, + "grounded_correctness": 0.30, + "computation_correctness": 0.15, + "tool_routing": 0.05, + "instruction_safety": 0.05, + "latency_cost": 0.05, + }, + "memory_required": { + "pairwise_preference_score": 0.40, + "grounded_correctness": 0.30, + "tool_routing": 0.05, + "instruction_safety": 0.05, + "latency_cost": 0.05, + }, + "no_tool": { + # Per the ChatGPT-thread design — pairwise carries more weight + # when there's no tool/retrieval signal to share weight with. + "pairwise_preference_score": 0.50, + "grounded_correctness": 0.25, + "tool_routing": 0.10, + "instruction_safety": 0.10, + "latency_cost": 0.05, + }, +} + + +def default_weights(task_type: str) -> dict[str, float]: + return dict(_DEFAULT_WEIGHTS.get(task_type, _DEFAULT_WEIGHTS["no_tool"])) + + +def applicable_metrics(task_type: str) -> set[str]: + return set(default_weights(task_type).keys()) + + +# -- Deterministic per-dimension scorers ----------------------------------- + + +def score_tool_routing( + *, + task_type: str, + tools_called: list[str], + has_citations: bool = False, +) -> float: + """Did the agent pick the right tool for this task type? + + Primary signal is ``tools_called`` (parsed from ``response.tool_calls`` + in the miner payload). Some SDK runtimes don't surface tool_calls + even when web tools were used — in that case ``has_citations`` is a + reasonable proxy for ``web_required``: if the agent returned cited + URLs, it almost certainly invoked ``web_search`` or ``url_fetch``. + + For ``no_tool`` we reward correctly NOT calling anything; calling a + tool when none was needed scores 0.5 (not zero — the tool may have + been a sensible double-check, just not optimal). + """ + called = {t.strip().lower() for t in tools_called if t} + web_tools = {"web_search", "url_fetch"} + if task_type == "web_required": + if called & web_tools: + return 1.0 + if has_citations: + return 1.0 # citations present → web tool used (SDK didn't surface call) + return 0.0 + if task_type == "rag_required": + # Tool name matches the SDK ``RagTool.name`` and the + # ``rag_tool_service`` ledger entry. Older code referenced an + # imaginary ``query_attachment`` tool that never existed. + return 1.0 if "rag.retrieve" in called else 0.0 + if task_type == "sandbox_required": + return 1.0 if "sandbox_python" in called else 0.0 + if task_type == "memory_required": + # Memory tasks live in `request.turns`; no tool call needed. + return 1.0 + if task_type == "no_tool": + if not called and not has_citations: + return 1.0 + return 0.5 + return 0.5 # unknown task type — neutral + + +def score_latency_cost( + *, + miner_latency_seconds: float, + mode_budget_seconds: float | None, + proxy_cost_usd: float, + cost_budget_usd: float | None, +) -> float: + """Within latency + cost budget? + + Linear ramp from 1.0 at 0% of budget to 0.0 at 100% of budget. The + final score is the minimum of the two ramps so a miner can't trade + speed for cost or vice versa. + """ + parts: list[float] = [] + if mode_budget_seconds and mode_budget_seconds > 0: + ratio = max(0.0, miner_latency_seconds) / mode_budget_seconds + parts.append(max(0.0, 1.0 - ratio)) + if cost_budget_usd and cost_budget_usd > 0: + ratio = max(0.0, proxy_cost_usd) / cost_budget_usd + parts.append(max(0.0, 1.0 - ratio)) + if not parts: + return 1.0 + return min(parts) + + +# -- Re-normalization + final score ---------------------------------------- + + +@dataclass(slots=True) +class TaskScoreBreakdown: + """Final per-task score + bookkeeping the validator persists.""" + + task_type: str + dimension_scores: dict[str, float] # raw score per dimension (only applicable ones) + applied_weights: dict[str, float] # post-renormalization weights + final_task_score: float + applicable_metrics: list[str] + + +def renormalize( + weights: dict[str, float], + *, + applicable: set[str], +) -> dict[str, float]: + """Drop non-applicable dimensions and rescale remaining weights to sum to 1.""" + relevant = {k: v for k, v in weights.items() if k in applicable and v > 0} + total = sum(relevant.values()) + if total <= 0: + return {} + return {k: v / total for k, v in relevant.items()} + + +def assemble_task_score( + *, + task_type: str, + raw_scores: dict[str, float | None], + weights: dict[str, float] | None = None, +) -> TaskScoreBreakdown: + """Combine per-dimension scores into a single weighted task_score. + + ``raw_scores`` may contain ``None`` for dimensions that came back + as N/A (e.g. retrieval_quality on a rag_required task with no + chunk-id ledger, or grounded_correctness when the judge call + failed). Those drop out and the remaining weights re-normalize. + """ + base_weights = dict(weights or default_weights(task_type)) + # Applicable = dimension is in the default weight set AND we have a + # non-None real-numbered score. If a default-applicable dimension + # has no score, we drop it (treat as N/A) rather than imputing 0.0. + applicable: set[str] = set() + real_scores: dict[str, float] = {} + for dim, w in base_weights.items(): + score = raw_scores.get(dim) + if score is None: + continue + applicable.add(dim) + real_scores[dim] = float(score) + applied = renormalize(base_weights, applicable=applicable) + final = sum(real_scores[k] * applied.get(k, 0.0) for k in real_scores) + return TaskScoreBreakdown( + task_type=task_type, + dimension_scores=real_scores, + applied_weights=applied, + final_task_score=final, + applicable_metrics=sorted(applicable), + ) + + +__all__ = [ + "TaskScoreBreakdown", + "applicable_metrics", + "assemble_task_score", + "default_weights", + "derive_task_type", + "renormalize", + "score_latency_cost", + "score_tool_routing", +] From b2ae10f5552ec90febe5a75a475a0c6e4d9215a6 Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Thu, 7 May 2026 17:31:57 +0000 Subject: [PATCH 12/24] feat(tools): server-attested tool-call ledger; url-fetch + mcp-relay services; sandbox session persistence --- tool_platforms/_record_tool_call.py | 100 ++++ tool_platforms/mcp_relay_service/__init__.py | 0 .../mcp_relay_service/_capabilities.py | 39 ++ tool_platforms/mcp_relay_service/_ssrf.py | 94 ++++ tool_platforms/mcp_relay_service/app.py | 333 ++++++++++++ tool_platforms/mcp_relay_service/models.py | 30 ++ tool_platforms/provider_proxy/app.py | 92 +++- tool_platforms/provider_proxy/redis_store.py | 161 +++++- tool_platforms/sandbox_tool_service/app.py | 42 ++ .../sandbox_tool_service/backends.py | 483 ++++++++++++++++-- tool_platforms/sandbox_tool_service/models.py | 30 ++ .../url_fetch_tool_service/__init__.py | 12 + tool_platforms/url_fetch_tool_service/app.py | 415 +++++++++++++++ .../url_fetch_tool_service/extractor.py | 92 ++++ tool_platforms/url_fetch_tool_service/ssrf.py | 101 ++++ tool_platforms/web_search_tool_service/app.py | 39 +- 16 files changed, 1990 insertions(+), 73 deletions(-) create mode 100644 tool_platforms/_record_tool_call.py create mode 100644 tool_platforms/mcp_relay_service/__init__.py create mode 100644 tool_platforms/mcp_relay_service/_capabilities.py create mode 100644 tool_platforms/mcp_relay_service/_ssrf.py create mode 100644 tool_platforms/mcp_relay_service/app.py create mode 100644 tool_platforms/mcp_relay_service/models.py create mode 100644 tool_platforms/url_fetch_tool_service/__init__.py create mode 100644 tool_platforms/url_fetch_tool_service/app.py create mode 100644 tool_platforms/url_fetch_tool_service/extractor.py create mode 100644 tool_platforms/url_fetch_tool_service/ssrf.py diff --git a/tool_platforms/_record_tool_call.py b/tool_platforms/_record_tool_call.py new file mode 100644 index 0000000..6224d50 --- /dev/null +++ b/tool_platforms/_record_tool_call.py @@ -0,0 +1,100 @@ +"""Fire-and-forget helper for the server-attested tool-call ledger. + +Tool services (web_search, url_fetch, sandbox, mcp_relay) call this after +a successful tool invocation so the eval pipeline can compute tool-use +KPIs from the orchestrator's authoritative log — never from +miner-emitted trace frames. Like ``charge_tool_cost``, the call is +best-effort: a transient owner-api outage leaves a single ledger row +unrecorded but never breaks the tool response. The downside of a missed +write is a dropped attestation signal for one item; the upside of +fail-open is that the tool platform stays available when owner-api hiccups. + +The owner-api endpoint is ``POST /v1/internal/eval/tool_calls`` with the +``EIREL_INTERNAL_SERVICE_TOKEN`` bearer. +""" +from __future__ import annotations + +import hashlib +import json +import logging +import os +from typing import Any + +import httpx + +_logger = logging.getLogger(__name__) + +_RECORD_TIMEOUT_SECONDS = 5.0 +_DIGEST_MAX_CHARS = 512 + + +def hash_args(args: dict[str, Any]) -> str: + canonical = json.dumps(args, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +def digest_result(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + text = value + else: + try: + text = json.dumps(value, sort_keys=True, default=str) + except (TypeError, ValueError): + text = str(value) + return text[:_DIGEST_MAX_CHARS] + + +async def record_tool_call( + *, + job_id: str | None, + tool_name: str, + args: dict[str, Any] | None, + result: Any = None, + latency_ms: int = 0, + cost_usd: float = 0.0, + status_str: str = "ok", + error: str | None = None, + owner_api_url: str | None = None, + owner_api_token: str | None = None, +) -> None: + """Best-effort POST of one tool-call row to the orchestrator ledger. + + A missing ``job_id`` (validator smoke test, direct curl) or missing + owner-api config silently skips; production miner traffic always + carries ``X-Eirel-Job-Id`` so this is a fast no-op when irrelevant. + """ + if not job_id: + return + url = owner_api_url or os.getenv("EIREL_OWNER_API_URL", "") + if not url: + return + token = owner_api_token or os.getenv("EIREL_INTERNAL_SERVICE_TOKEN", "") + args_dict = args or {} + payload = { + "job_id": job_id, + "tool_name": tool_name, + "args_hash": hash_args(args_dict), + "args_json": args_dict, + "result_digest": digest_result(result), + "latency_ms": int(max(0, latency_ms)), + "cost_usd": float(max(0.0, cost_usd)), + "status": status_str, + "error": error, + } + headers: dict[str, str] = {"Content-Type": "application/json"} + if token: + headers["Authorization"] = f"Bearer {token}" + endpoint = f"{url.rstrip('/')}/v1/internal/eval/tool_calls" + try: + async with httpx.AsyncClient(timeout=_RECORD_TIMEOUT_SECONDS) as client: + await client.post(endpoint, json=payload, headers=headers) + except (httpx.HTTPError, OSError) as exc: + _logger.warning( + "record_tool_call failed: job=%s tool=%s err=%s", + job_id, tool_name, exc, + ) + + +__all__ = ["record_tool_call", "hash_args", "digest_result"] diff --git a/tool_platforms/mcp_relay_service/__init__.py b/tool_platforms/mcp_relay_service/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tool_platforms/mcp_relay_service/_capabilities.py b/tool_platforms/mcp_relay_service/_capabilities.py new file mode 100644 index 0000000..d203314 --- /dev/null +++ b/tool_platforms/mcp_relay_service/_capabilities.py @@ -0,0 +1,39 @@ +"""Helpers for canonicalizing + hashing an integration's tool surface.""" +from __future__ import annotations + +import hashlib +import json +from collections.abc import Sequence +from typing import Any + +__all__ = ["canonicalize_tools", "hash_capabilities"] + + +def canonicalize_tools(tools: Sequence[Any]) -> list[dict[str, Any]]: + """Return a sorted, normalized list of tool descriptors. + + Each entry is reduced to ``{name, description, parameters_schema}``. + Sorted by ``name`` so two equivalent surfaces hash identically + regardless of source ordering. + """ + out: list[dict[str, Any]] = [] + for item in tools: + if not isinstance(item, dict): + continue + name = item.get("name") + if not isinstance(name, str) or not name: + continue + out.append({ + "name": name, + "description": str(item.get("description") or ""), + "parameters_schema": item.get("parameters_schema") or {}, + }) + out.sort(key=lambda t: t["name"]) + return out + + +def hash_capabilities(tools: Sequence[Any]) -> str: + """Stable sha256 hex digest over the canonicalized tool list.""" + canonical = canonicalize_tools(tools) + payload = json.dumps(canonical, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() diff --git a/tool_platforms/mcp_relay_service/_ssrf.py b/tool_platforms/mcp_relay_service/_ssrf.py new file mode 100644 index 0000000..d0ac620 --- /dev/null +++ b/tool_platforms/mcp_relay_service/_ssrf.py @@ -0,0 +1,94 @@ +"""Runtime SSRF guard for the MCP relay. + +Belt-and-suspenders: an integration's ``base_url`` is also checked at +admin-registration time. The runtime check ensures that even if a +disabled integration row gets reactivated with a tampered base_url, or +DNS rebinds an originally-public host to a private IP, the relay +refuses to call. +""" +from __future__ import annotations + +import ipaddress +import socket +from urllib.parse import urlparse + +__all__ = ["MCPSSRFError", "validate_base_url"] + + +class MCPSSRFError(ValueError): + """Raised when an integration's base_url fails SSRF policy.""" + + +_BLOCKED_HOSTS = frozenset({ + "metadata.google.internal", + "metadata", + "instance-data", + "instance-data.ec2.internal", +}) + + +def _is_private(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: + return ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_multicast + or ip.is_reserved + or ip.is_unspecified + ) + + +def validate_base_url(url: str, *, allow_http: bool = False) -> tuple[str, str]: + """Validate ``url`` against SSRF policy. Returns ``(scheme, host)``. + + Production (``allow_http=False``) requires HTTPS and resolves DNS to + catch hostnames that point at private IPs (DNS rebinding). Dev / + test (``allow_http=True``) skips the DNS step so tests can use + fixture hostnames that don't resolve. The literal-IP check runs in + both modes — that's the primary defense against + ``169.254.169.254``-style attacks. + """ + parsed = urlparse(url) + scheme = (parsed.scheme or "").lower() + if scheme not in ("https",) and not (allow_http and scheme == "http"): + raise MCPSSRFError( + f"unsupported scheme {scheme!r}; production allows https only" + ) + host = (parsed.hostname or "").lower() + if not host: + raise MCPSSRFError("base_url has no hostname") + if host in _BLOCKED_HOSTS: + raise MCPSSRFError(f"hostname {host!r} is on the SSRF blocklist") + parsed_ip: ipaddress.IPv4Address | ipaddress.IPv6Address | None + try: + parsed_ip = ipaddress.ip_address(host) + except ValueError: + parsed_ip = None + if parsed_ip is not None: + if _is_private(parsed_ip): + raise MCPSSRFError( + f"IP {host!r} is in a private/reserved range" + ) + return scheme, host + if allow_http: + # Dev / test mode — skip the DNS step. The integration's actual + # endpoint is mocked in tests via httpx.MockTransport, so DNS + # would only fail against fixture hostnames that don't exist. + return scheme, host + try: + infos = socket.getaddrinfo(host, None) + except OSError as exc: + raise MCPSSRFError(f"DNS resolution failed for {host!r}: {exc}") from exc + for entry in infos: + sockaddr = entry[4] + if not sockaddr: + continue + try: + ip = ipaddress.ip_address(sockaddr[0]) + except ValueError: + continue + if _is_private(ip): + raise MCPSSRFError( + f"hostname {host!r} resolves to private IP {ip!r}" + ) + return scheme, host diff --git a/tool_platforms/mcp_relay_service/app.py b/tool_platforms/mcp_relay_service/app.py new file mode 100644 index 0000000..7f378e4 --- /dev/null +++ b/tool_platforms/mcp_relay_service/app.py @@ -0,0 +1,333 @@ +"""MCP relay service. + +Sits between :class:`MCPToolDispatcher` (orchestrator) and the +operator-vetted MCP integration. Inbound auth is the internal-service +Bearer token — only the orchestrator should reach this. Outbound auth +is the per-connection encrypted OAuth access token loaded from the DB. + +Endpoints: + + * ``POST /v1/relay/connections/{connection_id}/call`` — invoke one + tool on behalf of one consumer connection. Verifies the + integration's stored ``capabilities_hash`` matches what the caller + has (passed in the body); refuses on drift. + * ``POST /v1/relay/integrations/{integration_id}/list_tools`` — + reprobe an integration's tool surface, return the canonical hash. + Operator-only path; used by the admin reprobe route. + +The relay does NOT write the audit row. The dispatcher does, because +it owns the conversation/message context. The relay returns latency + +cost so the dispatcher can stamp them into the audit. +""" +from __future__ import annotations + +import asyncio +import logging +import os +import time +from contextlib import asynccontextmanager +from datetime import datetime +from typing import Any + +import httpx +import uvicorn +from fastapi import Depends, FastAPI, Header, HTTPException, status +from sqlalchemy import select + +from shared.common.database import Database +from shared.common.models import ( + ConsumerMcpConnection, + McpIntegration, +) +from shared.safety.token_encryption import TokenCipher, build_token_cipher +from tool_platforms.mcp_relay_service._capabilities import ( + canonicalize_tools, + hash_capabilities, +) +from tool_platforms.mcp_relay_service._ssrf import ( + MCPSSRFError, + validate_base_url, +) +from tool_platforms.mcp_relay_service.models import ( + RelayCallRequest, + RelayCallResponse, + RelayListToolsResponse, +) + + +_logger = logging.getLogger(__name__) + +DEFAULT_RELAY_TIMEOUT_SECONDS: float = 10.0 +DEFAULT_MAX_RESPONSE_BYTES: int = 256 * 1024 +DEFAULT_PER_CALL_COST_USD: float = 0.0002 + + +def create_app( + *, + database: Database, + cipher: TokenCipher | None = None, + transport: httpx.AsyncBaseTransport | None = None, + allow_http: bool | None = None, +) -> FastAPI: + @asynccontextmanager + async def lifespan(app: FastAPI): + app.state.auth_token = os.getenv("EIREL_MCP_RELAY_TOKEN", "") + app.state.timeout_seconds = float( + os.getenv( + "EIREL_MCP_RELAY_TIMEOUT_SECONDS", + str(DEFAULT_RELAY_TIMEOUT_SECONDS), + ) + ) + app.state.max_response_bytes = int( + os.getenv( + "EIREL_MCP_RELAY_MAX_RESPONSE_BYTES", + str(DEFAULT_MAX_RESPONSE_BYTES), + ) + ) + app.state.per_call_cost_usd = float( + os.getenv( + "EIREL_MCP_RELAY_PER_CALL_COST_USD", + str(DEFAULT_PER_CALL_COST_USD), + ) + ) + # ``allow_http`` defaults to True in non-production so tests + # against ``http://testserver`` work. Operators flip + # ``EIREL_MCP_RELAY_REQUIRE_HTTPS=1`` in prod. + if allow_http is None: + require_https = os.getenv( + "EIREL_MCP_RELAY_REQUIRE_HTTPS", "0" + ).strip() in {"1", "true", "True"} + app.state.allow_http = not require_https + else: + app.state.allow_http = bool(allow_http) + app.state.database = database + app.state.cipher = cipher or build_token_cipher() + app.state.transport = transport + yield + + app = FastAPI( + title="mcp-relay-service", + version="0.1.0", + lifespan=lifespan, + ) + + async def require_auth( + authorization: str | None = Header(default=None), + ) -> None: + configured: str = app.state.auth_token + if not configured: + return + if authorization == f"Bearer {configured}": + return + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="invalid mcp relay auth token", + ) + + @app.get("/healthz") + async def healthz() -> dict[str, str]: + return {"status": "ok"} + + @app.post( + "/v1/relay/integrations/{integration_id}/list_tools", + response_model=RelayListToolsResponse, + ) + async def list_tools( + integration_id: str, + _: None = Depends(require_auth), + ) -> RelayListToolsResponse: + db: Database = app.state.database + with db.sessionmaker() as session: + integration = session.get(McpIntegration, integration_id) + if integration is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="integration not found", + ) + base_url = integration.base_url + try: + validate_base_url(base_url, allow_http=app.state.allow_http) + except MCPSSRFError as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"ssrf_blocked: {exc}", + ) from exc + client_kwargs: dict[str, Any] = {"timeout": app.state.timeout_seconds} + if app.state.transport is not None: + client_kwargs["transport"] = app.state.transport + async with httpx.AsyncClient(**client_kwargs) as client: + try: + resp = await client.post( + f"{base_url.rstrip('/')}/list_tools", + json={}, + ) + resp.raise_for_status() + except httpx.HTTPError as exc: + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail=f"upstream_error: {exc}", + ) from exc + try: + payload = resp.json() + except ValueError as exc: + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail=f"upstream_invalid_json: {exc}", + ) from exc + tools_raw = payload.get("tools") if isinstance(payload, dict) else None + if not isinstance(tools_raw, list): + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail="upstream_missing_tools", + ) + canonical = canonicalize_tools(tools_raw) + h = hash_capabilities(canonical) + return RelayListToolsResponse(tools=canonical, capabilities_hash=h) + + @app.post( + "/v1/relay/connections/{connection_id}/call", + response_model=RelayCallResponse, + ) + async def call( + connection_id: str, + payload: RelayCallRequest, + _: None = Depends(require_auth), + ) -> RelayCallResponse: + db: Database = app.state.database + cipher: TokenCipher = app.state.cipher + with db.sessionmaker() as session: + connection = session.get(ConsumerMcpConnection, connection_id) + if connection is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="connection not found", + ) + if connection.status != "active": + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"connection_not_active: {connection.status}", + ) + integration = session.get(McpIntegration, connection.integration_id) + if integration is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="integration not found", + ) + if integration.status != "active": + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="integration_disabled", + ) + if ( + payload.capabilities_hash + and integration.capabilities_hash + and payload.capabilities_hash != integration.capabilities_hash + ): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="capabilities_hash_mismatch", + ) + base_url = integration.base_url + access_token: str | None = None + if connection.oauth_access_token_encrypted: + try: + access_token = cipher.decrypt( + bytes(connection.oauth_access_token_encrypted) + ).decode("utf-8") + except Exception as exc: # noqa: BLE001 — defensive + _logger.warning( + "relay token decrypt failed connection=%s err=%s", + connection_id, exc, + ) + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="token_decrypt_failed", + ) from exc + connection.last_used_at = datetime.utcnow() + session.commit() + + try: + validate_base_url(base_url, allow_http=app.state.allow_http) + except MCPSSRFError as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"ssrf_blocked: {exc}", + ) from exc + + headers: dict[str, str] = {"Content-Type": "application/json"} + if access_token: + headers["Authorization"] = f"Bearer {access_token}" + + client_kwargs: dict[str, Any] = {"timeout": app.state.timeout_seconds} + if app.state.transport is not None: + client_kwargs["transport"] = app.state.transport + + t0 = time.perf_counter() + try: + async with httpx.AsyncClient(**client_kwargs) as client: + resp = await asyncio.wait_for( + client.post( + f"{base_url.rstrip('/')}/tools/{payload.tool_name}", + json=payload.args, + headers=headers, + ), + timeout=app.state.timeout_seconds, + ) + resp.raise_for_status() + if ( + resp.headers.get("content-length") + and int(resp.headers["content-length"]) + > app.state.max_response_bytes + ): + raise httpx.HTTPError( + f"response too large: {resp.headers['content-length']}" + ) + body = resp.content + if len(body) > app.state.max_response_bytes: + raise httpx.HTTPError( + f"response too large: {len(body)} bytes" + ) + result = resp.json() + except asyncio.TimeoutError: + latency_ms = int((time.perf_counter() - t0) * 1000) + return RelayCallResponse( + ok=False, + error="upstream_timeout", + latency_ms=latency_ms, + cost_usd=0.0, + ) + except (httpx.HTTPError, ValueError) as exc: + latency_ms = int((time.perf_counter() - t0) * 1000) + return RelayCallResponse( + ok=False, + error=f"upstream_error: {exc}", + latency_ms=latency_ms, + cost_usd=0.0, + ) + + latency_ms = int((time.perf_counter() - t0) * 1000) + if not isinstance(result, dict): + return RelayCallResponse( + ok=False, + error="upstream_invalid_shape", + latency_ms=latency_ms, + cost_usd=0.0, + ) + return RelayCallResponse( + ok=True, + result=result, + latency_ms=latency_ms, + cost_usd=float(app.state.per_call_cost_usd), + ) + + return app + + +def main() -> None: + db_url = os.getenv("DATABASE_URL", "") + if not db_url: + raise RuntimeError("DATABASE_URL must be set for the mcp_relay_service") + db = Database(db_url) + app = create_app(database=db) + port = int(os.getenv("EIREL_MCP_RELAY_PORT", "8093")) + uvicorn.run(app, host="0.0.0.0", port=port, reload=False) diff --git a/tool_platforms/mcp_relay_service/models.py b/tool_platforms/mcp_relay_service/models.py new file mode 100644 index 0000000..e178290 --- /dev/null +++ b/tool_platforms/mcp_relay_service/models.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + + +class RelayCallRequest(BaseModel): + """Body for ``POST /v1/relay/connections/{connection_id}/call``.""" + + tool_name: str = Field(min_length=1, max_length=128) + args: dict[str, Any] = Field(default_factory=dict) + capabilities_hash: str | None = Field(default=None, max_length=64) + + +class RelayCallResponse(BaseModel): + """Outcome returned to the orchestrator's :class:`MCPToolDispatcher`.""" + + ok: bool + result: dict[str, Any] | None = None + error: str | None = None + latency_ms: int = 0 + cost_usd: float = 0.0 + + +class RelayListToolsResponse(BaseModel): + """Operator/admin tool — used to refresh ``capabilities_hash``.""" + + tools: list[dict[str, Any]] = Field(default_factory=list) + capabilities_hash: str = "" diff --git a/tool_platforms/provider_proxy/app.py b/tool_platforms/provider_proxy/app.py index ffb7c55..55ba665 100644 --- a/tool_platforms/provider_proxy/app.py +++ b/tool_platforms/provider_proxy/app.py @@ -41,6 +41,30 @@ class ChargeToolRequest(_BaseModel): tool_name: str amount_usd: float = _Field(ge=0.0) + span_id: str | None = None + parent_span_id: str | None = None + + +class ReserveBatchEstimate(_BaseModel): + estimated_cost: float = _Field(ge=0.0) + estimated_tokens: int = _Field(ge=0) + provider: str + model: str + span_id: str | None = None + + +class ReserveBatchRequest(_BaseModel): + """Atomic batch reservation for graph parallel-tool dispatch. + + All ``estimates`` must fit under the run budget *together* — if the + sum would exceed it, none are reserved and the call returns a 429. + This eliminates the race where N concurrent reserve_estimate calls + each pass an individually-affordable check but collectively blow + the cap. + """ + + max_usd_budget: float = _Field(ge=0.0) + estimates: list[ReserveBatchEstimate] @dataclass(slots=True) @@ -182,6 +206,13 @@ async def job_cost(job_id: str, _: str = Depends(require_auth)) -> dict[str, Any tool_cost_usd += float(value) else: llm_cost_usd += float(value) + # Aggregate per-span totals across buckets (tool/llm). + cost_by_span_totals: dict[str, float] = {} + for key, value in (usage.cost_by_span or {}).items(): + span_id, _, _bucket = key.partition("::") + cost_by_span_totals[span_id] = round( + cost_by_span_totals.get(span_id, 0.0) + float(value), 8 + ) return { "cost_usd_used": usage.cost_usd_used, "llm_cost_usd": round(llm_cost_usd, 8), @@ -189,6 +220,9 @@ async def job_cost(job_id: str, _: str = Depends(require_auth)) -> dict[str, Any "max_usd_budget": usage.max_usd_budget, "cost_rejections": usage.cost_rejections, "per_provider": dict(usage.cost_by_provider), + "per_span": cost_by_span_totals, + "per_span_buckets": dict(usage.cost_by_span or {}), + "span_parents": dict(usage.span_parents or {}), } @app.post("/v1/jobs/{job_id}/charge_tool") @@ -196,6 +230,8 @@ async def charge_tool( job_id: str, body: ChargeToolRequest, _: str = Depends(require_auth), + x_eirel_span_id: str | None = Header(default=None), + x_eirel_parent_span_id: str | None = Header(default=None), ) -> dict[str, Any]: store = app.state.services.store if not await store.exists(job_id): @@ -203,15 +239,67 @@ async def charge_tool( status_code=status.HTTP_404_NOT_FOUND, detail="job not found", ) + # Header takes precedence over body.span_id so callers that wire + # tracing once at the HTTP layer don't have to plumb span ids + # through every payload. Body fields are kept as a fallback for + # codepaths that can't set headers (e.g. retries on a queued tool + # call replay). + span_id = x_eirel_span_id or body.span_id + parent_span_id = x_eirel_parent_span_id or body.parent_span_id accepted, cost_used = await store.charge_tool( - job_id=job_id, tool_name=body.tool_name, amount_usd=body.amount_usd, + job_id=job_id, + tool_name=body.tool_name, + amount_usd=body.amount_usd, + span_id=span_id, + parent_span_id=parent_span_id, ) if not accepted: raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail="run budget exhausted", ) - return {"cost_usd_used": cost_used} + return {"cost_usd_used": cost_used, "span_id": span_id} + + @app.post("/v1/jobs/{job_id}/reserve_batch_estimate") + async def reserve_batch_estimate( + job_id: str, + body: ReserveBatchRequest, + _: str = Depends(require_auth), + ) -> dict[str, Any]: + """Atomic batch reservation for parallel-tool dispatch. + + On success returns the per-estimate ``span_id`` of every + reservation that landed alongside the new ``cost_usd_used``. + On budget exhaustion, NONE of the estimates are recorded and + the call 429s — that's the whole point of batching. + """ + store = app.state.services.store + if not body.estimates: + return {"cost_usd_used": 0.0, "reserved": []} + accepted, cost_used = await store.reserve_batch_estimate( + job_id=job_id, + max_usd_budget=body.max_usd_budget, + estimates=[ + { + "estimated_cost": e.estimated_cost, + "estimated_tokens": e.estimated_tokens, + "provider": e.provider, + "model": e.model, + "span_id": e.span_id, + } + for e in body.estimates + ], + ) + if not accepted: + app.state.services.metrics["quota_rejections_total"] += 1 + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail="run budget exhausted", + ) + return { + "cost_usd_used": cost_used, + "reserved": [e.span_id for e in body.estimates], + } @app.post("/v1/chat/completions", response_model=ProviderProxyResponse) async def chat_completions( diff --git a/tool_platforms/provider_proxy/redis_store.py b/tool_platforms/provider_proxy/redis_store.py index 00dd81b..d5f6be7 100644 --- a/tool_platforms/provider_proxy/redis_store.py +++ b/tool_platforms/provider_proxy/redis_store.py @@ -47,6 +47,10 @@ _PROVIDER_COUNTS_PREFIX = "provider_proxy:provider_counts:" _MODEL_COUNTS_PREFIX = "provider_proxy:model_counts:" _COST_BY_PROVIDER_PREFIX = "provider_proxy:cost_by_provider:" +# Per-graph-span cost attribution. Key = ``tool:`` or +# ``tool::`` so eiretes can compute per-span cost +# without parsing the cost_by_provider hash. +_COST_BY_SPAN_PREFIX = "provider_proxy:cost_by_span:" @dataclass(slots=True) @@ -63,6 +67,12 @@ class JobUsage: max_usd_budget: float = 0.0 cost_by_provider: dict[str, float] = field(default_factory=dict) cost_rejections: int = 0 + # Graph-runtime per-span cost attribution. Keys are + # ``"::"`` where bucket is ``tool:`` or + # ``llm::``. ``span_parents`` maps span_id -> + # parent_span_id so consumers can rebuild the span tree. + cost_by_span: dict[str, float] = field(default_factory=dict) + span_parents: dict[str, str] = field(default_factory=dict) # -- Lua scripts ------------------------------------------------------------- @@ -298,11 +308,25 @@ async def charge_tool( job_id: str, tool_name: str, amount_usd: float, + span_id: str | None = None, + parent_span_id: str | None = None, ) -> tuple[bool, float]: """Atomically check budget and charge a tool call. Returns ``(accepted, cost_usd_used)``. ``accepted=False`` → 429. + + When ``span_id`` is set, the same amount is also written to the + per-span cost hash so downstream consumers (eiretes trace KPIs) + can attribute cost to the graph span that emitted the call. + ``parent_span_id`` is recorded in the cost-by-span entry's + adjacent metadata key for parent linkage. """ + # The Lua script handles only the budget-and-cost-by-provider + # bookkeeping; per-span cost is layered on top via a follow-up + # HINCRBYFLOAT on the dedicated key. That stays race-safe because + # the budget invariant is enforced by the Lua atomicity *before* + # we touch the span hash — over-charging the span hash is just + # accounting, never affects rejection. if self._lua_available: try: result = await self._charge_tool_script( @@ -318,6 +342,14 @@ async def charge_tool( ) rejected = int(result[0]) == 1 cost_used = float(result[1]) + if not rejected: + await self._record_span_cost( + job_id=job_id, + span_id=span_id, + parent_span_id=parent_span_id, + amount_usd=amount_usd, + bucket=f"tool:{tool_name}", + ) return (not rejected, cost_used) except ResponseError as exc: if "unknown command" not in str(exc).lower(): @@ -325,7 +357,11 @@ async def charge_tool( self._lua_available = False _logger.info("Lua scripting unavailable; using Python fallback path") return await self._charge_tool_fallback( - job_id=job_id, tool_name=tool_name, amount_usd=amount_usd, + job_id=job_id, + tool_name=tool_name, + amount_usd=amount_usd, + span_id=span_id, + parent_span_id=parent_span_id, ) async def _charge_tool_fallback( @@ -334,6 +370,8 @@ async def _charge_tool_fallback( job_id: str, tool_name: str, amount_usd: float, + span_id: str | None = None, + parent_span_id: str | None = None, ) -> tuple[bool, float]: usage_key = _USAGE_PREFIX + job_id cost_key = _COST_BY_PROVIDER_PREFIX + job_id @@ -352,8 +390,110 @@ async def _charge_tool_fallback( pipe.expire(usage_key, self._ttl) pipe.expire(cost_key, self._ttl) await pipe.execute() + await self._record_span_cost( + job_id=job_id, + span_id=span_id, + parent_span_id=parent_span_id, + amount_usd=amount_usd, + bucket=f"tool:{tool_name}", + ) return (True, cost + amount_usd) + async def reserve_batch_estimate( + self, + *, + job_id: str, + max_usd_budget: float, + estimates: list[dict], + ) -> tuple[bool, float]: + """Atomically reserve N estimates as one operation. + + Either every estimate fits under ``max_usd_budget`` and all + reservations land, or none do and the call returns + ``(False, current_cost)``. Eliminates the over-reservation race + when graph parallel-tool nodes fire many ``reserve_estimate`` + calls concurrently against the same job_id. + + Each entry in ``estimates`` is a dict with keys + ``estimated_cost``, ``estimated_tokens``, ``provider``, ``model``, + and optional ``span_id`` for per-span cost attribution. + """ + if not estimates: + usage = await self.get(job_id) + return (True, usage.cost_usd_used if usage else 0.0) + usage_key = _USAGE_PREFIX + job_id + async with self._per_job_locks[job_id]: + cost_str = await self._client.hget(usage_key, "cost_usd_used") + current = float(cost_str) if cost_str is not None else 0.0 + sum_estimated = sum(float(e["estimated_cost"]) for e in estimates) + if current + sum_estimated > max_usd_budget: + await self._client.hincrby(usage_key, "cost_rejections", 1) + await self._client.expire(usage_key, self._ttl) + return (False, current) + # All-or-nothing commit: pipe a single transaction that + # writes every estimate. Span attribution is best-effort and + # piggybacks onto the same pipe. + async with self._client.pipeline(transaction=True) as pipe: + pipe.hsetnx(usage_key, "started_at", str(time.monotonic())) + pipe.hsetnx(usage_key, "max_usd_budget", str(max_usd_budget)) + for est in estimates: + estimated_cost = float(est["estimated_cost"]) + estimated_tokens = int(est.get("estimated_tokens") or 0) + provider = str(est["provider"]) + model = str(est["model"]) + pipe.hincrbyfloat(usage_key, "cost_usd_used", estimated_cost) + pipe.hincrby(usage_key, "request_count", 1) + pipe.hincrby(usage_key, "estimated_total_tokens", estimated_tokens) + pipe.hincrby( + _PROVIDER_COUNTS_PREFIX + job_id, provider, 1 + ) + pipe.hincrby( + _MODEL_COUNTS_PREFIX + job_id, f"{provider}:{model}", 1 + ) + pipe.expire(usage_key, self._ttl) + pipe.expire(_PROVIDER_COUNTS_PREFIX + job_id, self._ttl) + pipe.expire(_MODEL_COUNTS_PREFIX + job_id, self._ttl) + await pipe.execute() + # Per-span attribution after the budget commit — best effort. + for est in estimates: + span_id = est.get("span_id") + if not span_id: + continue + await self._record_span_cost( + job_id=job_id, + span_id=str(span_id), + parent_span_id=None, + amount_usd=float(est["estimated_cost"]), + bucket=f"llm:{est['provider']}:{est['model']}", + ) + return (True, current + sum_estimated) + + async def _record_span_cost( + self, + *, + job_id: str, + span_id: str | None, + parent_span_id: str | None, + amount_usd: float, + bucket: str, + ) -> None: + if not span_id: + return + cost_by_span_key = _COST_BY_SPAN_PREFIX + job_id + # Hash field is ":" so a span that drives both + # an LLM call and a tool call surfaces both buckets. + field = f"{span_id}::{bucket}" + async with self._client.pipeline(transaction=True) as pipe: + pipe.hincrbyfloat(cost_by_span_key, field, float(amount_usd)) + if parent_span_id: + pipe.hset( + cost_by_span_key, + f"__parent::{span_id}", + parent_span_id, + ) + pipe.expire(cost_by_span_key, self._ttl) + await pipe.execute() + # ------------------------------------------------------------------ # Reads # ------------------------------------------------------------------ @@ -366,14 +506,29 @@ async def get(self, job_id: str) -> JobUsage | None: provider_counts_key = _PROVIDER_COUNTS_PREFIX + job_id model_counts_key = _MODEL_COUNTS_PREFIX + job_id cost_key = _COST_BY_PROVIDER_PREFIX + job_id + cost_by_span_key = _COST_BY_SPAN_PREFIX + job_id async with self._client.pipeline(transaction=False) as pipe: pipe.hgetall(usage_key) pipe.hgetall(provider_counts_key) pipe.hgetall(model_counts_key) pipe.hgetall(cost_key) - usage_raw, provider_raw, model_raw, cost_raw = await pipe.execute() + pipe.hgetall(cost_by_span_key) + ( + usage_raw, + provider_raw, + model_raw, + cost_raw, + span_raw, + ) = await pipe.execute() if not usage_raw: return None + cost_by_span: dict[str, float] = {} + span_parents: dict[str, str] = {} + for k, v in (span_raw or {}).items(): + if k.startswith("__parent::"): + span_parents[k[len("__parent::"):]] = str(v) + else: + cost_by_span[k] = float(v) return JobUsage( started_at=float(usage_raw.get("started_at", 0.0) or 0.0), request_count=int(usage_raw.get("request_count", 0) or 0), @@ -385,6 +540,8 @@ async def get(self, job_id: str) -> JobUsage | None: provider_request_counts={k: int(v) for k, v in (provider_raw or {}).items()}, model_request_counts={k: int(v) for k, v in (model_raw or {}).items()}, cost_by_provider={k: float(v) for k, v in (cost_raw or {}).items()}, + cost_by_span=cost_by_span, + span_parents=span_parents, ) async def list_all(self) -> dict[str, JobUsage]: diff --git a/tool_platforms/sandbox_tool_service/app.py b/tool_platforms/sandbox_tool_service/app.py index a207170..ddcb9ee 100644 --- a/tool_platforms/sandbox_tool_service/app.py +++ b/tool_platforms/sandbox_tool_service/app.py @@ -16,6 +16,7 @@ create_job_ledger, ) from tool_platforms._charge_tool import charge_tool_cost +from tool_platforms._record_tool_call import record_tool_call from tool_platforms.sandbox_tool_service.backends import ( SandboxBackend, SubprocessBackend, @@ -23,6 +24,7 @@ from tool_platforms.sandbox_tool_service.models import ( SandboxExecuteRequest, SandboxExecuteResponse, + SandboxFileWrite, ) @@ -80,6 +82,13 @@ async def lifespan(app: FastAPI): yield finally: await app.state.job_ledger.close() + shutdown_backend = getattr(app.state, "backend", None) + close = getattr(shutdown_backend, "close", None) + if callable(close): + try: + await close() + except Exception: # pragma: no cover - best-effort teardown + pass app = FastAPI( title="sandbox-tool-service", @@ -202,10 +211,17 @@ async def execute( app.state.metrics["execute_requests_total"] += 1 retrieved_at = _utcnow() + attachments = ( + [a.model_dump() for a in payload.attachments] + if payload.attachments + else None + ) result = await app.state.backend.execute( code=payload.code, timeout_seconds=payload.timeout_seconds, memory_mb=payload.memory_mb, + session_id=payload.session_id, + attachments=attachments, ) if result.timed_out: @@ -220,7 +236,30 @@ async def execute( await charge_tool_cost( job_id=job_id, tool_name="sandbox", amount_usd=per_exec_cost, ) + await record_tool_call( + job_id=job_id, tool_name="sandbox", + args={ + "code_n_bytes": len(payload.code.encode("utf-8")), + "session_id": payload.session_id, + "n_attachments": len(payload.attachments or []), + }, + result={ + "exit_code": result.exit_code, + "duration_ms": result.duration_ms, + "stdout_n_chars": len(result.stdout or ""), + "stderr_n_chars": len(result.stderr or ""), + "n_files_written": len(getattr(result, "files", ()) or ()), + "timed_out": result.timed_out, + }, + latency_ms=int(result.duration_ms or 0), + cost_usd=per_exec_cost, + status_str="ok" if result.exit_code == 0 and not result.timed_out else "error", + ) + files_payload = [ + SandboxFileWrite(path=f.path, size=f.size, content_b64=f.content_b64) + for f in getattr(result, "files", ()) + ] return SandboxExecuteResponse( stdout=result.stdout, stderr=result.stderr, @@ -229,9 +268,12 @@ async def execute( truncated=result.truncated, retrieved_at=retrieved_at, retrieval_ledger_id=usage.ledger_id, + session_id=payload.session_id, + files=files_payload, metadata={ "backend": type(app.state.backend).__name__, "timed_out": result.timed_out, + "session_id": payload.session_id, }, ) diff --git a/tool_platforms/sandbox_tool_service/backends.py b/tool_platforms/sandbox_tool_service/backends.py index b1124d9..4c79936 100644 --- a/tool_platforms/sandbox_tool_service/backends.py +++ b/tool_platforms/sandbox_tool_service/backends.py @@ -1,12 +1,17 @@ from __future__ import annotations import asyncio +import base64 +import json import logging +import os import resource +import shutil import sys +import tempfile import time -from dataclasses import dataclass -from typing import Protocol +from dataclasses import dataclass, field +from typing import Any, Protocol _logger = logging.getLogger(__name__) @@ -60,6 +65,68 @@ def find_spec(self, name, path=None, target=None): ''' +# Worker loop for session-persistent kernels. Reads newline-delimited +# JSON commands from stdin and writes length-prefixed JSON responses +# directly to fd 1, bypassing the per-call sys.stdout swap. User code +# in exec() sees its own StringIO for stdout/stderr, so our framing +# does not collide with user prints. +_WORKER_LOOP = ''' +import io +import json +import sys +import traceback + +_globals = {"__name__": "__main__"} +_real_buf = sys.__stdout__.buffer + + +def _send(obj): + data = json.dumps(obj).encode("utf-8") + _real_buf.write(f"{len(data)}\\n".encode("ascii")) + _real_buf.write(data) + _real_buf.flush() + + +while True: + line = sys.stdin.readline() + if not line: + break + try: + cmd = json.loads(line) + except Exception: + continue + code = cmd.get("code", "") + out = io.StringIO() + err = io.StringIO() + prev_out, prev_err = sys.stdout, sys.stderr + sys.stdout, sys.stderr = out, err + exit_code = 0 + try: + exec(compile(code, "", "exec"), _globals) + except SystemExit as e: + exit_code = int(e.code) if isinstance(e.code, int) else 1 + except BaseException: + traceback.print_exc(file=err) + exit_code = 1 + finally: + sys.stdout, sys.stderr = prev_out, prev_err + _send({ + "stdout": out.getvalue(), + "stderr": err.getvalue(), + "exit_code": exit_code, + }) +''' + + +@dataclass(slots=True, frozen=True) +class SandboxFile: + """One file produced by the sandbox during execution.""" + + path: str + size: int + content_b64: str | None = None # ``None`` when the file exceeds the per-file cap. + + @dataclass(slots=True, frozen=True) class ExecutionResult: stdout: str @@ -68,6 +135,7 @@ class ExecutionResult: duration_ms: int truncated: bool = False timed_out: bool = False + files: tuple[SandboxFile, ...] = () class SandboxBackend(Protocol): @@ -75,20 +143,46 @@ async def execute( self, *, code: str, - timeout_seconds: float, - memory_mb: int, + timeout_seconds: float | None = None, + memory_mb: int | None = None, + session_id: str | None = None, + attachments: list[dict[str, Any]] | None = None, ) -> ExecutionResult: ... +@dataclass +class _SessionKernel: + """Long-lived worker subprocess for one session.""" + + proc: asyncio.subprocess.Process + workdir: str + last_access: float + lock: asyncio.Lock + + @dataclass(slots=True) class SubprocessBackend: """Execute Python code in an isolated subprocess with resource limits. + Two execution modes: + + - **One-shot** (no ``session_id``): spawn a fresh subprocess per call, + isolated mode (``-I -S``), apply RLIMIT_*, run user code, capture + stdout/stderr, return. ``attachments`` (if any) are mounted into a + throwaway temp directory; files written during execution come back + in ``ExecutionResult.files``. + - **Session** (``session_id`` set): reuse a long-lived worker + subprocess keyed by ``session_id``. The worker keeps a globals dict + across exec calls so variables, imports, and the working directory + persist between turns. Idle for more than ``session_idle_ttl`` + seconds → kernel evicted on next access. A timeout or worker crash + also evicts the kernel; the next call starts fresh. + Security model: this backend is defense-in-depth suitable for deployment **inside a hardened container**. The subprocess itself enforces: - RLIMIT_AS (memory) — hard cap per process - - RLIMIT_CPU (CPU time) — hard wall, backs up wall-clock timeout + - RLIMIT_CPU (CPU time) — hard wall (cumulative across a session) - RLIMIT_NOFILE (file descriptors) — 64 max - RLIMIT_NPROC (processes) — 8 max, blocks fork bombs - RLIMIT_FSIZE (file size) — 1 MB max per file @@ -96,21 +190,6 @@ class SubprocessBackend: - Python ``-I -S`` flags (isolated mode, no site-packages) - Import preamble blocking network / subprocess / ctypes / ssl - This does NOT provide kernel-level isolation. For production, - run this service in a Docker container hardened with:: - - docker run \\ - --user nobody \\ - --cap-drop ALL \\ - --security-opt no-new-privileges \\ - --security-opt seccomp= \\ - --network none \\ - --read-only \\ - --tmpfs /tmp:size=64m,mode=1777 \\ - --memory 512m --memory-swap 512m \\ - --pids-limit 100 \\ - sandbox-tool-service - For stronger isolation (adversarial miners, shared infrastructure), implement a ``GVisorBackend`` or ``FirecrackerBackend`` satisfying the ``SandboxBackend`` Protocol. They plug in transparently — no @@ -122,6 +201,11 @@ class SubprocessBackend: default_memory_mb: int = 128 stdout_max_bytes: int = 65_536 stderr_max_bytes: int = 16_384 + session_idle_ttl: float = 300.0 + file_size_cap: int = 256 * 1024 + files_total_cap: int = 1024 * 1024 + _sessions: dict[str, _SessionKernel] = field(default_factory=dict) + _sessions_lock: asyncio.Lock | None = None async def execute( self, @@ -129,38 +213,50 @@ async def execute( code: str, timeout_seconds: float | None = None, memory_mb: int | None = None, + session_id: str | None = None, + attachments: list[dict[str, Any]] | None = None, ) -> ExecutionResult: - wall_clock = timeout_seconds or self.default_timeout + timeout = timeout_seconds or self.default_timeout mem_mb = memory_mb or self.default_memory_mb - mem_bytes = mem_mb * 1024 * 1024 - cpu_limit = max(1, int(wall_clock) + 2) + if session_id: + return await self._execute_session( + session_id=session_id, + code=code, + timeout=timeout, + memory_mb=mem_mb, + attachments=attachments, + ) + return await self._execute_oneshot( + code=code, + timeout=timeout, + memory_mb=mem_mb, + attachments=attachments, + ) + + async def close(self) -> None: + for sid in list(self._sessions.keys()): + await self._evict(sid) + + # -- one-shot path ----------------------------------------------------- + + async def _execute_oneshot( + self, + *, + code: str, + timeout: float, + memory_mb: int, + attachments: list[dict[str, Any]] | None, + ) -> ExecutionResult: + mem_bytes = memory_mb * 1024 * 1024 + cpu_limit = max(1, int(timeout) + 2) full_code = _PREAMBLE + "\n" + code - def _set_limits() -> None: - try: - resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes)) - except (OSError, ValueError): - pass - try: - resource.setrlimit(resource.RLIMIT_CPU, (cpu_limit, cpu_limit)) - except (OSError, ValueError): - pass - try: - resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64)) - except (OSError, ValueError): - pass - try: - resource.setrlimit(resource.RLIMIT_NPROC, (8, 8)) - except (OSError, ValueError, AttributeError): - pass - try: - resource.setrlimit(resource.RLIMIT_FSIZE, (1024 * 1024, 1024 * 1024)) - except (OSError, ValueError): - pass - try: - resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) - except (OSError, ValueError): - pass + workdir: str | None = None + snap_before: dict[str, int] = {} + if attachments: + workdir = tempfile.mkdtemp(prefix="sandbox-oneshot-") + self._mount_attachments(workdir, attachments) + snap_before = self._snapshot(workdir) t0 = time.perf_counter() try: @@ -173,10 +269,13 @@ def _set_limits() -> None: stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, - preexec_fn=_set_limits, + cwd=workdir, + preexec_fn=_make_preexec(mem_bytes, cpu_limit), ) except OSError as exc: _logger.warning("sandbox spawn failed: %s", exc) + if workdir is not None: + shutil.rmtree(workdir, ignore_errors=True) return ExecutionResult( stdout="", stderr=f"sandbox_spawn_error: {exc}", @@ -189,7 +288,7 @@ def _set_limits() -> None: try: stdout_bytes, stderr_bytes = await asyncio.wait_for( proc.communicate(), - timeout=wall_clock, + timeout=timeout, ) exit_code = proc.returncode if proc.returncode is not None else 0 except asyncio.TimeoutError: @@ -214,6 +313,11 @@ def _set_limits() -> None: stderr_bytes = stderr_bytes[: self.stderr_max_bytes] truncated = True + files: tuple[SandboxFile, ...] = () + if workdir is not None: + files = self._collect_files(workdir, snap_before) + shutil.rmtree(workdir, ignore_errors=True) + return ExecutionResult( stdout=stdout_bytes.decode("utf-8", errors="replace"), stderr=stderr_bytes.decode("utf-8", errors="replace"), @@ -221,4 +325,283 @@ def _set_limits() -> None: duration_ms=duration_ms, truncated=truncated, timed_out=timed_out, + files=files, + ) + + # -- session path ------------------------------------------------------ + + async def _execute_session( + self, + *, + session_id: str, + code: str, + timeout: float, + memory_mb: int, + attachments: list[dict[str, Any]] | None, + ) -> ExecutionResult: + if self._sessions_lock is None: + self._sessions_lock = asyncio.Lock() + + async with self._sessions_lock: + kernel = self._sessions.get(session_id) + now = time.time() + if kernel is not None and ( + kernel.proc.returncode is not None + or now - kernel.last_access > self.session_idle_ttl + ): + await self._evict(session_id) + kernel = None + if kernel is None: + try: + kernel = await self._spawn_kernel(memory_mb) + except Exception as exc: + _logger.warning("sandbox kernel spawn failed: %s", exc) + return ExecutionResult( + stdout="", + stderr=f"sandbox_spawn_error: {exc}", + exit_code=-1, + duration_ms=0, + ) + self._sessions[session_id] = kernel + + async with kernel.lock: + if attachments: + self._mount_attachments(kernel.workdir, attachments) + snap_before = self._snapshot(kernel.workdir) + + assert kernel.proc.stdin is not None + assert kernel.proc.stdout is not None + t0 = time.perf_counter() + try: + msg = (json.dumps({"code": code}) + "\n").encode("utf-8") + kernel.proc.stdin.write(msg) + await kernel.proc.stdin.drain() + + length_line = await asyncio.wait_for( + kernel.proc.stdout.readline(), timeout=timeout + ) + if not length_line: + raise RuntimeError("worker terminated unexpectedly") + try: + length = int(length_line.strip()) + except ValueError as exc: + raise RuntimeError( + f"sandbox framing error: {length_line!r}" + ) from exc + data = await asyncio.wait_for( + kernel.proc.stdout.readexactly(length), timeout=timeout + ) + response = json.loads(data) + except asyncio.TimeoutError: + duration_ms = int((time.perf_counter() - t0) * 1000) + await self._evict(session_id) + return ExecutionResult( + stdout="", + stderr="sandbox session timeout", + exit_code=-9, + duration_ms=duration_ms, + timed_out=True, + ) + except Exception as exc: # noqa: BLE001 — surface to caller + _logger.warning( + "sandbox session %s exec error: %s", session_id, exc + ) + duration_ms = int((time.perf_counter() - t0) * 1000) + await self._evict(session_id) + return ExecutionResult( + stdout="", + stderr=f"sandbox_session_error: {exc}", + exit_code=-1, + duration_ms=duration_ms, + ) + + duration_ms = int((time.perf_counter() - t0) * 1000) + kernel.last_access = time.time() + + stdout = str(response.get("stdout", "")) + stderr = str(response.get("stderr", "")) + exit_code = int(response.get("exit_code", 0)) + stdout_bytes = stdout.encode("utf-8") + stderr_bytes = stderr.encode("utf-8") + truncated = False + if len(stdout_bytes) > self.stdout_max_bytes: + stdout = stdout_bytes[: self.stdout_max_bytes].decode( + "utf-8", errors="replace" + ) + truncated = True + if len(stderr_bytes) > self.stderr_max_bytes: + stderr = stderr_bytes[: self.stderr_max_bytes].decode( + "utf-8", errors="replace" + ) + truncated = True + + files = self._collect_files(kernel.workdir, snap_before) + + return ExecutionResult( + stdout=stdout, + stderr=stderr, + exit_code=exit_code, + duration_ms=duration_ms, + truncated=truncated, + timed_out=False, + files=files, + ) + + async def _spawn_kernel(self, memory_mb: int) -> _SessionKernel: + mem_bytes = memory_mb * 1024 * 1024 + # Cumulative CPU cap for the lifetime of the kernel; idle sessions + # don't consume CPU, so this caps abuse without throttling normal + # multi-turn use. + cpu_limit = max(1, int(self.session_idle_ttl)) + full_code = _PREAMBLE + "\n" + _WORKER_LOOP + workdir = tempfile.mkdtemp(prefix="sandbox-session-") + try: + proc = await asyncio.create_subprocess_exec( + self.python_path, + "-I", + "-S", + "-c", + full_code, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=workdir, + preexec_fn=_make_preexec(mem_bytes, cpu_limit), + ) + except OSError: + shutil.rmtree(workdir, ignore_errors=True) + raise + return _SessionKernel( + proc=proc, + workdir=workdir, + last_access=time.time(), + lock=asyncio.Lock(), ) + + async def _evict(self, session_id: str) -> None: + kernel = self._sessions.pop(session_id, None) + if kernel is None: + return + try: + kernel.proc.kill() + except ProcessLookupError: + pass + try: + await asyncio.wait_for(kernel.proc.wait(), timeout=2.0) + except asyncio.TimeoutError: # pragma: no cover - kill should win + pass + except Exception: # pragma: no cover + pass + shutil.rmtree(kernel.workdir, ignore_errors=True) + + # -- file helpers ------------------------------------------------------ + + def _mount_attachments( + self, + workdir: str, + attachments: list[dict[str, Any]], + ) -> None: + if not attachments: + return + real_workdir = os.path.realpath(workdir) + for att in attachments: + if isinstance(att, dict): + filename = att.get("filename") or "" + content_b64 = att.get("content_b64") or "" + else: # pragma: no cover — defensive + filename = getattr(att, "filename", "") or "" + content_b64 = getattr(att, "content_b64", "") or "" + if not isinstance(filename, str) or not isinstance(content_b64, str): + continue + safe = os.path.basename(filename) + if not safe or safe in (".", ".."): + continue + target = os.path.join(workdir, safe) + real_target = os.path.realpath(target) + try: + if os.path.commonpath([real_target, real_workdir]) != real_workdir: + continue + except ValueError: + continue + try: + data = base64.b64decode(content_b64, validate=False) + except Exception: + continue + try: + with open(target, "wb") as fh: + fh.write(data) + except OSError: + continue + + def _snapshot(self, workdir: str) -> dict[str, int]: + snap: dict[str, int] = {} + for root, _, files in os.walk(workdir): + for fname in files: + full = os.path.join(root, fname) + try: + snap[os.path.relpath(full, workdir)] = os.stat(full).st_mtime_ns + except OSError: + pass + return snap + + def _collect_files( + self, + workdir: str, + snap_before: dict[str, int], + ) -> tuple[SandboxFile, ...]: + out: list[SandboxFile] = [] + total = 0 + for root, _, files in os.walk(workdir): + for fname in files: + full = os.path.join(root, fname) + try: + st = os.stat(full) + except OSError: + continue + rel = os.path.relpath(full, workdir) + prev = snap_before.get(rel) + if prev is not None and prev == st.st_mtime_ns: + continue + size = st.st_size + content_b64: str | None = None + if size <= self.file_size_cap and total + size <= self.files_total_cap: + try: + with open(full, "rb") as fh: + content_b64 = base64.b64encode(fh.read()).decode("ascii") + total += size + except OSError: + pass + out.append(SandboxFile(path=rel, size=size, content_b64=content_b64)) + return tuple(out) + + +def _make_preexec(mem_bytes: int, cpu_limit: int): + """Return a preexec_fn that applies RLIMIT_* before exec.""" + + def _set_limits() -> None: + try: + resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes)) + except (OSError, ValueError): + pass + try: + resource.setrlimit(resource.RLIMIT_CPU, (cpu_limit, cpu_limit)) + except (OSError, ValueError): + pass + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64)) + except (OSError, ValueError): + pass + try: + resource.setrlimit(resource.RLIMIT_NPROC, (8, 8)) + except (OSError, ValueError, AttributeError): + pass + try: + resource.setrlimit(resource.RLIMIT_FSIZE, (1024 * 1024, 1024 * 1024)) + except (OSError, ValueError): + pass + try: + resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) + except (OSError, ValueError): + pass + + return _set_limits diff --git a/tool_platforms/sandbox_tool_service/models.py b/tool_platforms/sandbox_tool_service/models.py index 1e221eb..e086a3e 100644 --- a/tool_platforms/sandbox_tool_service/models.py +++ b/tool_platforms/sandbox_tool_service/models.py @@ -5,10 +5,38 @@ from pydantic import BaseModel, Field +class SandboxAttachment(BaseModel): + """Pre-existing file mounted into the sandbox working directory. + + ``content_b64`` is base64-encoded raw bytes. ``max_length`` here is a + safety bound on the request body — typical attachments come from + ``ConsumerAttachment`` rows that the orchestrator has already + size-bounded at upload. + """ + + filename: str = Field(min_length=1, max_length=255) + content_b64: str = Field(default="", max_length=14_000_000) + + +class SandboxFileWrite(BaseModel): + """One file that the sandbox produced during execution. + + ``content_b64`` is ``None`` when the file exceeds the per-file or + total cap; the path and size still come back so the agent can decide + whether to retry with a smaller output or split the work. + """ + + path: str + size: int + content_b64: str | None = None + + class SandboxExecuteRequest(BaseModel): code: str = Field(min_length=1, max_length=65_536) timeout_seconds: float | None = Field(default=None, gt=0.0, le=30.0) memory_mb: int | None = Field(default=None, gt=0, le=1024) + session_id: str | None = Field(default=None, min_length=1, max_length=128) + attachments: list[SandboxAttachment] | None = Field(default=None, max_length=16) class SandboxExecuteResponse(BaseModel): @@ -19,4 +47,6 @@ class SandboxExecuteResponse(BaseModel): truncated: bool = False retrieved_at: str | None = None retrieval_ledger_id: str | None = None + session_id: str | None = None + files: list[SandboxFileWrite] = Field(default_factory=list) metadata: dict[str, Any] = Field(default_factory=dict) diff --git a/tool_platforms/url_fetch_tool_service/__init__.py b/tool_platforms/url_fetch_tool_service/__init__.py new file mode 100644 index 0000000..66692b2 --- /dev/null +++ b/tool_platforms/url_fetch_tool_service/__init__.py @@ -0,0 +1,12 @@ +"""URL-fetch tool service. + +Subnet-owned HTTP service that miner agents call to read the contents +of a public URL. Mirrors ChatGPT's ``browser`` and Claude's ``web_fetch`` +tools at the wire level: input is a URL, output is extracted text + +metadata. +""" +from __future__ import annotations + +from tool_platforms.url_fetch_tool_service.app import create_app, main + +__all__ = ["create_app", "main"] diff --git a/tool_platforms/url_fetch_tool_service/app.py b/tool_platforms/url_fetch_tool_service/app.py new file mode 100644 index 0000000..9190283 --- /dev/null +++ b/tool_platforms/url_fetch_tool_service/app.py @@ -0,0 +1,415 @@ +"""FastAPI app for the URL-fetch tool service. + +Mirrors the auth + ledger pattern of ``sandbox_tool_service``: bearer +token (master) + per-job HMAC token, per-job request-count cap via the +shared :class:`JobLedger`, per-host rate-limit + body-size cap on the +HTTP fetch itself. +""" +from __future__ import annotations + +import asyncio +import hmac +import os +import time +from contextlib import asynccontextmanager +from datetime import UTC, datetime +from typing import Any + +import httpx +import uvicorn +from fastapi import Depends, FastAPI, Header, HTTPException, status +from fastapi.responses import PlainTextResponse +from pydantic import BaseModel, Field + +from shared.common.redis_job_ledger import ( + JobLedger, + JobUsageRecord, + create_job_ledger, +) +from tool_platforms._charge_tool import charge_tool_cost +from tool_platforms._record_tool_call import record_tool_call +from tool_platforms.url_fetch_tool_service.extractor import ( + ExtractedDocument, + extract_text, +) +from tool_platforms.url_fetch_tool_service.ssrf import ( + UrlFetchSSRFError, + validate_url, +) + + +# Hard cap on response body bytes — protects against memory blowups +# from a hostile or unbounded server. Configurable per-deploy but +# enforced per-request inside the streaming read. +_DEFAULT_MAX_RESPONSE_BYTES = 1 * 1024 * 1024 # 1 MB +_DEFAULT_FETCH_TIMEOUT_SECONDS = 15.0 +_DEFAULT_MAX_REDIRECTS = 5 +_DEFAULT_PER_HOST_RATE = 4 # max requests per window per host +_DEFAULT_PER_HOST_WINDOW_SECONDS = 10.0 + + +def generate_job_token(master_token: str, job_id: str) -> str: + return hmac.new(master_token.encode(), job_id.encode(), "sha256").hexdigest() + + +def verify_job_token(master_token: str, job_id: str, token: str) -> bool: + expected = generate_job_token(master_token, job_id) + return hmac.compare_digest(expected, token) + + +def _utcnow_iso() -> str: + return datetime.now(UTC).isoformat() + + +# -- Request / response models ---------------------------------------------- + + +class UrlFetchRequest(BaseModel): + url: str = Field(min_length=1, max_length=2048) + max_chars: int = Field(default=50_000, ge=100, le=500_000) + include_links: bool = True + + +class UrlFetchResponse(BaseModel): + url: str + final_url: str + title: str + content: str + links: list[dict[str, str]] = Field(default_factory=list) + status_code: int + content_type: str + bytes_read: int + truncated: bool + + +# -- Per-host rate limiter -------------------------------------------------- + + +class _PerHostRateLimiter: + """Token-bucket-style sliding-window rate limit, per-host. + + In-process state — fine for the single-replica deployment we ship + with. Multi-replica scale-out moves to Redis with the same window + semantics. + """ + + def __init__(self, *, max_per_window: int, window_seconds: float) -> None: + self._max = max_per_window + self._window = window_seconds + self._buckets: dict[str, list[float]] = {} + self._lock = asyncio.Lock() + + async def check(self, host: str) -> bool: + now = time.monotonic() + async with self._lock: + bucket = self._buckets.setdefault(host, []) + cutoff = now - self._window + while bucket and bucket[0] < cutoff: + bucket.pop(0) + if len(bucket) >= self._max: + return False + bucket.append(now) + return True + + +# -- App factory ------------------------------------------------------------ + + +def create_app( + *, + transport: httpx.AsyncBaseTransport | None = None, + ledger: JobLedger | None = None, +) -> FastAPI: + @asynccontextmanager + async def lifespan(app: FastAPI): + app.state.auth_token = os.getenv("EIREL_URL_FETCH_API_TOKEN", "") + app.state.default_max_requests = int( + os.getenv("EIREL_URL_FETCH_DEFAULT_MAX_REQUESTS", "12") + ) + app.state.max_response_bytes = int( + os.getenv("EIREL_URL_FETCH_MAX_RESPONSE_BYTES", str(_DEFAULT_MAX_RESPONSE_BYTES)) + ) + app.state.fetch_timeout_seconds = float( + os.getenv("EIREL_URL_FETCH_TIMEOUT_SECONDS", str(_DEFAULT_FETCH_TIMEOUT_SECONDS)) + ) + app.state.max_redirects = int( + os.getenv("EIREL_URL_FETCH_MAX_REDIRECTS", str(_DEFAULT_MAX_REDIRECTS)) + ) + per_host_rate = int( + os.getenv("EIREL_URL_FETCH_PER_HOST_RATE", str(_DEFAULT_PER_HOST_RATE)) + ) + per_host_window = float( + os.getenv( + "EIREL_URL_FETCH_PER_HOST_WINDOW_SECONDS", + str(_DEFAULT_PER_HOST_WINDOW_SECONDS), + ) + ) + app.state.rate_limiter = _PerHostRateLimiter( + max_per_window=per_host_rate, window_seconds=per_host_window, + ) + app.state.transport = transport + app.state.user_agent = os.getenv( + "EIREL_URL_FETCH_USER_AGENT", "EirelUrlFetch/0.1 (+https://eirel.ai)", + ) + if ledger is not None: + app.state.job_ledger = ledger + else: + app.state.job_ledger = create_job_ledger(os.getenv("REDIS_URL", "")) + app.state.metrics = { + "requests_total": 0, + "quota_rejections_total": 0, + "ssrf_blocks_total": 0, + "rate_limit_blocks_total": 0, + "fetch_failures_total": 0, + "size_cap_truncations_total": 0, + } + try: + yield + finally: + await app.state.job_ledger.close() + + app = FastAPI( + title="url-fetch-tool-service", + version="0.1.0", + lifespan=lifespan, + ) + + async def require_auth( + authorization: str | None = Header(default=None), + x_eirel_job_id: str | None = Header(default=None), + ) -> None: + auth_token: str = app.state.auth_token + if not auth_token: + return + if authorization == f"Bearer {auth_token}": + return + if x_eirel_job_id and authorization: + bearer = authorization.removeprefix("Bearer ").strip() + if verify_job_token(auth_token, x_eirel_job_id.strip(), bearer): + return + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="invalid url-fetch tool auth token", + ) + + async def require_job( + x_eirel_job_id: str | None = Header(default=None), + x_eirel_max_requests: str | None = Header(default=None), + ) -> tuple[str, int]: + job_id = (x_eirel_job_id or "").strip() + if not job_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="missing X-Eirel-Job-Id header", + ) + max_requests = app.state.default_max_requests + if x_eirel_max_requests: + try: + max_requests = max(1, int(x_eirel_max_requests)) + except ValueError: + pass + return job_id, max_requests + + @app.get("/healthz") + async def healthz() -> dict[str, str]: + return {"status": "ok", "now": _utcnow_iso()} + + @app.get("/metrics", response_class=PlainTextResponse) + async def metrics() -> str: + m: dict[str, Any] = app.state.metrics + lines = [ + "# HELP eirel_url_fetch_requests_total Successful URL fetch requests.", + "# TYPE eirel_url_fetch_requests_total counter", + f"eirel_url_fetch_requests_total {m['requests_total']}", + "# HELP eirel_url_fetch_quota_rejections_total Job-level quota rejections.", + "# TYPE eirel_url_fetch_quota_rejections_total counter", + f"eirel_url_fetch_quota_rejections_total {m['quota_rejections_total']}", + "# HELP eirel_url_fetch_ssrf_blocks_total SSRF policy denials.", + "# TYPE eirel_url_fetch_ssrf_blocks_total counter", + f"eirel_url_fetch_ssrf_blocks_total {m['ssrf_blocks_total']}", + "# HELP eirel_url_fetch_rate_limit_blocks_total Per-host rate-limit denials.", + "# TYPE eirel_url_fetch_rate_limit_blocks_total counter", + f"eirel_url_fetch_rate_limit_blocks_total {m['rate_limit_blocks_total']}", + "# HELP eirel_url_fetch_failures_total Upstream fetch failures.", + "# TYPE eirel_url_fetch_failures_total counter", + f"eirel_url_fetch_failures_total {m['fetch_failures_total']}", + "# HELP eirel_url_fetch_truncations_total Responses truncated by size cap.", + "# TYPE eirel_url_fetch_truncations_total counter", + f"eirel_url_fetch_truncations_total {m['size_cap_truncations_total']}", + ] + return "\n".join(lines) + "\n" + + async def check_budget( + *, job_id: str, max_requests: int, tool_name: str + ) -> JobUsageRecord: + ledger: JobLedger = app.state.job_ledger + usage = await ledger.get_or_create(job_id) + if usage.request_count >= max_requests: + app.state.metrics["quota_rejections_total"] += 1 + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail={"error": "url_fetch_quota_exhausted", "max_requests": max_requests}, + ) + usage.request_count += 1 + usage.tool_counts[tool_name] = usage.tool_counts.get(tool_name, 0) + 1 + await ledger.save(job_id, usage) + return usage + + @app.post("/v1/fetch", response_model=UrlFetchResponse) + async def fetch( + body: UrlFetchRequest, + _: None = Depends(require_auth), + job: tuple[str, int] = Depends(require_job), + ) -> UrlFetchResponse: + job_id, max_requests = job + + # SSRF guard: validate scheme + DNS + IP ranges before any + # network IO. Fails loud so attempts are visible in metrics. + try: + scheme, host = validate_url(body.url) + except UrlFetchSSRFError as exc: + app.state.metrics["ssrf_blocks_total"] += 1 + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "ssrf_blocked", "reason": str(exc)}, + ) + del scheme # unused — validation enforces http/https + + # Per-host rate limit applied AFTER SSRF (so SSRF metrics are + # accurate even under sustained probes). + if not await app.state.rate_limiter.check(host): + app.state.metrics["rate_limit_blocks_total"] += 1 + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail={"error": "per_host_rate_limit_exceeded", "host": host}, + ) + + # Per-job quota check + ledger increment, AFTER SSRF + rate-limit + # so a blocked URL never burns budget. + await check_budget( + job_id=job_id, max_requests=max_requests, tool_name="url_fetch", + ) + + client_kwargs: dict[str, Any] = { + "timeout": app.state.fetch_timeout_seconds, + "follow_redirects": True, + "max_redirects": app.state.max_redirects, + "headers": {"User-Agent": app.state.user_agent}, + } + if app.state.transport is not None: + client_kwargs["transport"] = app.state.transport + + try: + async with httpx.AsyncClient(**client_kwargs) as client: + async with client.stream("GET", body.url) as resp: + chunks: list[bytes] = [] + truncated = False + cap = app.state.max_response_bytes + kept = 0 + async for chunk in resp.aiter_bytes(): + if not chunk: + continue + if kept + len(chunk) > cap: + # Take only the head up to the cap. + remaining = max(0, cap - kept) + if remaining: + chunks.append(chunk[:remaining]) + kept += remaining + truncated = True + break + chunks.append(chunk) + kept += len(chunk) + raw = b"".join(chunks) + bytes_read = kept + final_url = str(resp.url) + status_code = resp.status_code + content_type = resp.headers.get("content-type", "") + except httpx.HTTPError as exc: + app.state.metrics["fetch_failures_total"] += 1 + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail={"error": "fetch_failed", "reason": str(exc)}, + ) from None + + # Best-effort provider-proxy charge so per-task tool cost + # surfaces in DeploymentScoreRecord.tool_cost_usd. Failures + # are logged but never break the fetch. + per_call_cost = float(os.getenv("EIREL_URL_FETCH_PER_CALL_COST_USD", "0.0005")) + await charge_tool_cost( + job_id=job_id, tool_name="url_fetch", amount_usd=per_call_cost, + ) + await record_tool_call( + job_id=job_id, tool_name="url_fetch", + args={ + "url": body.url, "max_chars": body.max_chars, + "include_links": body.include_links, + }, + result={ + "final_url": final_url, "status_code": status_code, + "bytes_read": bytes_read, "content_type": content_type, + }, + cost_usd=per_call_cost, + status_str="ok" if status_code < 400 else "error", + ) + + if truncated: + app.state.metrics["size_cap_truncations_total"] += 1 + app.state.metrics["requests_total"] += 1 + + # Heuristic: only run the HTML extractor if content-type smells + # like HTML. Otherwise return the raw text best-effort decoded. + is_html = "html" in content_type.lower() + if is_html: + try: + text = raw.decode("utf-8", errors="replace") + except Exception: # noqa: BLE001 + text = raw.decode("latin-1", errors="replace") + extracted = extract_text( + text, + base_url=final_url, + max_chars=body.max_chars, + include_links=body.include_links, + ) + return UrlFetchResponse( + url=body.url, + final_url=final_url, + title=extracted.title, + content=extracted.content, + links=extracted.links, + status_code=status_code, + content_type=content_type, + bytes_read=bytes_read, + truncated=truncated, + ) + + # Non-HTML: surface raw text up to the cap. Useful for + # plain-text/markdown/json — the LLM consumer can interpret. + try: + text = raw.decode("utf-8", errors="replace") + except Exception: # noqa: BLE001 + text = raw.decode("latin-1", errors="replace") + if len(text) > body.max_chars: + text = text[: body.max_chars] + "..." + return UrlFetchResponse( + url=body.url, + final_url=final_url, + title="", + content=text, + links=[], + status_code=status_code, + content_type=content_type, + bytes_read=bytes_read, + truncated=truncated, + ) + + return app + + +def main() -> None: + port = int(os.getenv("EIREL_URL_FETCH_PORT", "8087")) + uvicorn.run( + "tool_platforms.url_fetch_tool_service.app:create_app", + host="0.0.0.0", + port=port, + factory=True, + ) diff --git a/tool_platforms/url_fetch_tool_service/extractor.py b/tool_platforms/url_fetch_tool_service/extractor.py new file mode 100644 index 0000000..d154c44 --- /dev/null +++ b/tool_platforms/url_fetch_tool_service/extractor.py @@ -0,0 +1,92 @@ +"""HTML → readable-text extraction. + +Bare-bones BeautifulSoup pipeline: drop `` + + """ + result = extract_text(html, base_url="https://example.com/post") + assert result.title == "My Page" + assert "navigation links" not in result.content + assert "header bar" not in result.content + assert "footer junk" not in result.content + assert "console.log" not in result.content + assert "Main heading" in result.content + assert "article body" in result.content + assert any(link["href"] == "https://example.com/about" for link in result.links) + + +def test_extractor_falls_back_to_body_without_main(): + html = "

plain body content

" + result = extract_text(html) + assert "plain body content" in result.content + + +def test_extractor_caps_max_chars(): + html = "" + ("x" * 200_000) + "" + result = extract_text(html, max_chars=1_000) + assert len(result.content) <= 1_010 # 1000 + ellipsis + + +def test_extractor_skips_javascript_and_anchor_links(): + html = """ + + skip + skip too + keep + + """ + result = extract_text(html, base_url="https://example.com/x") + hrefs = [link["href"] for link in result.links] + assert "https://example.com/real" in hrefs + assert all("javascript:" not in h for h in hrefs) + assert all(not h.endswith("#section") for h in hrefs) + + +# -- SSRF guard ------------------------------------------------------------- + + +def test_ssrf_blocks_unsupported_scheme(): + with pytest.raises(UrlFetchSSRFError, match="scheme"): + validate_url("file:///etc/passwd") + with pytest.raises(UrlFetchSSRFError, match="scheme"): + validate_url("gopher://example.com/") + + +def test_ssrf_blocks_loopback_ip(): + with pytest.raises(UrlFetchSSRFError, match="private"): + validate_url("http://127.0.0.1/") + with pytest.raises(UrlFetchSSRFError, match="private"): + validate_url("http://[::1]/") + + +def test_ssrf_blocks_rfc1918(): + with pytest.raises(UrlFetchSSRFError, match="private"): + validate_url("http://10.0.0.5/") + with pytest.raises(UrlFetchSSRFError, match="private"): + validate_url("http://192.168.1.1/") + with pytest.raises(UrlFetchSSRFError, match="private"): + validate_url("http://172.16.0.1/") + + +def test_ssrf_blocks_known_metadata_hosts(): + with pytest.raises(UrlFetchSSRFError, match="blocklist"): + validate_url("http://metadata.google.internal/computeMetadata/") + + +def test_ssrf_rejects_url_with_no_host(): + with pytest.raises(UrlFetchSSRFError, match="hostname"): + validate_url("http:///") + + +# -- Service: auth ---------------------------------------------------------- + + +@pytest.fixture +def fake_redis_ledger(monkeypatch): + """Force the in-memory job ledger so tests don't need Redis.""" + monkeypatch.setenv("REDIS_URL", "") # empty → InMemoryJobLedger + + +async def test_service_rejects_missing_auth(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + app = create_app() + async with app.router.lifespan_context(app): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://testserver") as client: + resp = await client.post( + "/v1/fetch", + json={"url": "https://example.com"}, + headers={"X-Eirel-Job-Id": "job-1"}, + ) + assert resp.status_code == 401 + + +async def test_service_accepts_master_bearer(monkeypatch, fake_redis_ledger): + """Auth passes; SSRF then rejects example.com? No — example.com is public. + But we don't actually want to hit the network in tests, so we use a stub + transport. That's covered in the fetch-roundtrip tests below; here we + just confirm auth passes and a downstream-failure doesn't 401.""" + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + transport = httpx.MockTransport( + lambda req: httpx.Response(200, content=b"ok", + headers={"content-type": "text/html"}), + ) + app = create_app(transport=transport) + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + resp = await client.post( + "/v1/fetch", + json={"url": "https://example.com"}, + headers={ + "Authorization": "Bearer tok", + "X-Eirel-Job-Id": "job-1", + }, + ) + assert resp.status_code == 200, resp.text + + +# -- Service: SSRF integration --------------------------------------------- + + +async def test_service_ssrf_blocks_loopback(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + app = create_app() + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + resp = await client.post( + "/v1/fetch", + json={"url": "http://127.0.0.1/"}, + headers={"Authorization": "Bearer tok", "X-Eirel-Job-Id": "job-1"}, + ) + assert resp.status_code == 400 + detail = resp.json()["detail"] + assert detail["error"] == "ssrf_blocked" + + +# -- Service: fetch roundtrip ----------------------------------------------- + + +async def test_service_html_roundtrip(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + html = ( + b"T" + b"

Hello

World.

" + b"next
" + ) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, content=html, headers={"content-type": "text/html; charset=utf-8"}, + ) + + transport = httpx.MockTransport(handler) + app = create_app(transport=transport) + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + resp = await client.post( + "/v1/fetch", + json={"url": "https://example.com/post"}, + headers={"Authorization": "Bearer tok", "X-Eirel-Job-Id": "job-1"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["title"] == "T" + assert "Hello" in body["content"] + assert "World." in body["content"] + assert body["status_code"] == 200 + assert body["content_type"].startswith("text/html") + assert any(link["href"].endswith("/next") for link in body["links"]) + assert body["truncated"] is False + + +async def test_service_size_cap_truncates(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + monkeypatch.setenv("EIREL_URL_FETCH_MAX_RESPONSE_BYTES", "256") + big_body = b"" + (b"X" * 10_000) + b"" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, content=big_body, headers={"content-type": "text/html"}, + ) + + transport = httpx.MockTransport(handler) + app = create_app(transport=transport) + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + resp = await client.post( + "/v1/fetch", + json={"url": "https://example.com/big"}, + headers={"Authorization": "Bearer tok", "X-Eirel-Job-Id": "job-1"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["truncated"] is True + assert body["bytes_read"] <= 512 # cap + last chunk overflow + + +async def test_service_per_host_rate_limit(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + monkeypatch.setenv("EIREL_URL_FETCH_PER_HOST_RATE", "2") + monkeypatch.setenv("EIREL_URL_FETCH_PER_HOST_WINDOW_SECONDS", "60") + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=b"x", + headers={"content-type": "text/html"}) + + transport = httpx.MockTransport(handler) + app = create_app(transport=transport) + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + statuses = [] + for _ in range(4): + resp = await client.post( + "/v1/fetch", + json={"url": "https://example.com/x"}, + headers={ + "Authorization": "Bearer tok", + "X-Eirel-Job-Id": "job-rl", + }, + ) + statuses.append(resp.status_code) + if resp.status_code == 429: + body = resp.json()["detail"] + assert body["error"] == "per_host_rate_limit_exceeded" + # First two within the window pass; subsequent ones rate-limit. + assert statuses[:2] == [200, 200] + assert 429 in statuses[2:] + + +async def test_service_quota_rejects_after_limit(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + transport = httpx.MockTransport( + lambda req: httpx.Response(200, content=b"x", + headers={"content-type": "text/html"}), + ) + app = create_app(transport=transport) + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + statuses = [] + for _ in range(3): + resp = await client.post( + "/v1/fetch", + json={"url": "https://example.com/x"}, + headers={ + "Authorization": "Bearer tok", + "X-Eirel-Job-Id": "job-q", + "X-Eirel-Max-Requests": "2", + }, + ) + statuses.append(resp.status_code) + assert statuses == [200, 200, 429] + + +async def test_service_upstream_error_returns_502(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("upstream down") + + transport = httpx.MockTransport(handler) + app = create_app(transport=transport) + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + resp = await client.post( + "/v1/fetch", + json={"url": "https://example.com/dead"}, + headers={"Authorization": "Bearer tok", "X-Eirel-Job-Id": "job-d"}, + ) + assert resp.status_code == 502 + assert resp.json()["detail"]["error"] == "fetch_failed" + + +async def test_service_non_html_returns_raw_text(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + transport = httpx.MockTransport( + lambda req: httpx.Response( + 200, + content=b'{"k":"v"}', + headers={"content-type": "application/json"}, + ), + ) + app = create_app(transport=transport) + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + resp = await client.post( + "/v1/fetch", + json={"url": "https://example.com/api/data.json"}, + headers={"Authorization": "Bearer tok", "X-Eirel-Job-Id": "job-j"}, + ) + body = resp.json() + assert body["content_type"].startswith("application/json") + assert body["title"] == "" + assert "k" in body["content"] and "v" in body["content"] + assert body["links"] == [] + + +async def test_service_healthz_metrics_unauthenticated(monkeypatch, fake_redis_ledger): + monkeypatch.setenv("EIREL_URL_FETCH_API_TOKEN", "tok") + app = create_app() + async with app.router.lifespan_context(app): + asgi = ASGITransport(app=app) + async with AsyncClient(transport=asgi, base_url="http://testserver") as client: + health = await client.get("/healthz") + assert health.status_code == 200 + assert health.json()["status"] == "ok" + metrics = await client.get("/metrics") + assert metrics.status_code == 200 + assert "eirel_url_fetch_requests_total" in metrics.text diff --git a/tests/services/test_web_search_tool_service.py b/tests/services/test_web_search_tool_service.py index a3ca77f..a64872c 100644 --- a/tests/services/test_web_search_tool_service.py +++ b/tests/services/test_web_search_tool_service.py @@ -13,8 +13,8 @@ ) -async def test_research_tool_service_enforces_auth_and_request_budgets(monkeypatch): - monkeypatch.setenv("EIREL_RESEARCH_TOOL_API_TOKEN", "tool-token") +async def test_web_search_tool_service_enforces_auth_and_request_budgets(monkeypatch): + monkeypatch.setenv("EIREL_WEB_SEARCH_TOOL_API_TOKEN", "tool-token") app = create_app( ResearchCatalogStore( documents={ @@ -66,8 +66,8 @@ async def test_research_tool_service_enforces_auth_and_request_budgets(monkeypat assert usage.json()["retrieval_ledger_id"] == "ledger:job-1" -async def test_research_tool_service_exposes_retrieval_ledger(monkeypatch): - monkeypatch.setenv("EIREL_RESEARCH_TOOL_API_TOKEN", "tool-token") +async def test_web_search_tool_service_exposes_retrieval_ledger(monkeypatch): + monkeypatch.setenv("EIREL_WEB_SEARCH_TOOL_API_TOKEN", "tool-token") app = create_app( ResearchCatalogStore( documents={ @@ -106,8 +106,8 @@ async def test_research_tool_service_exposes_retrieval_ledger(monkeypatch): assert payload["find_on_page_events"][0]["pattern"] == "coverage" -async def test_research_tool_service_brave_search_reranks_preferred_domains(monkeypatch): - monkeypatch.setenv("EIREL_RESEARCH_TOOL_API_TOKEN", "tool-token") +async def test_web_search_tool_service_brave_search_reranks_preferred_domains(monkeypatch): + monkeypatch.setenv("EIREL_WEB_SEARCH_TOOL_API_TOKEN", "tool-token") monkeypatch.setenv("EIREL_BRAVE_SEARCH_API_KEY", "brave-token") def _search_handler(request: httpx.Request) -> httpx.Response: @@ -162,8 +162,8 @@ def _search_handler(request: httpx.Request) -> httpx.Response: assert payload["documents"][0]["document_id"].startswith("web-") -async def test_research_tool_service_live_open_page_and_find_on_page_use_opened_text(monkeypatch): - monkeypatch.setenv("EIREL_RESEARCH_TOOL_API_TOKEN", "tool-token") +async def test_web_search_tool_service_live_open_page_and_find_on_page_use_opened_text(monkeypatch): + monkeypatch.setenv("EIREL_WEB_SEARCH_TOOL_API_TOKEN", "tool-token") monkeypatch.setenv("EIREL_BRAVE_SEARCH_API_KEY", "brave-token") def _search_handler(request: httpx.Request) -> httpx.Response: @@ -266,8 +266,8 @@ def _fetch_handler(request: httpx.Request) -> httpx.Response: ] -async def test_research_tool_service_live_open_page_marks_unsupported_content(monkeypatch): - monkeypatch.setenv("EIREL_RESEARCH_TOOL_API_TOKEN", "tool-token") +async def test_web_search_tool_service_live_open_page_marks_unsupported_content(monkeypatch): + monkeypatch.setenv("EIREL_WEB_SEARCH_TOOL_API_TOKEN", "tool-token") monkeypatch.setenv("EIREL_BRAVE_SEARCH_API_KEY", "brave-token") def _search_handler(request: httpx.Request) -> httpx.Response: @@ -327,7 +327,7 @@ def _fetch_handler(request: httpx.Request) -> httpx.Response: async def test_per_job_scoped_token_auth(monkeypatch): - monkeypatch.setenv("EIREL_RESEARCH_TOOL_API_TOKEN", "master-token") + monkeypatch.setenv("EIREL_WEB_SEARCH_TOOL_API_TOKEN", "master-token") app = create_app( ResearchCatalogStore( documents={ diff --git a/tests/staging_validation.py b/tests/staging_validation.py index 35f42dd..0512455 100644 --- a/tests/staging_validation.py +++ b/tests/staging_validation.py @@ -13,12 +13,6 @@ DIRECT_ANALYSIS_PROMPT = "Analyze the request directly and return a verified response." RESEARCH_TO_BUILD_PROMPT = "Research the task, implement the solution, and verify the result." CONCEPT_TO_MEDIA_PROMPT = "Turn the concept into media, then review the generated artifact." -DEPRECATED_OWNER_ENDPOINTS = ( - "/v1/workflow-serving/coalitions", - "/v1/workflow-serving/releases/current", - "/v1/operators/workflow-serving/freezes", - "/v1/families/builder/coalition-scores", -) def _status_from_checks(checks: dict[str, bool]) -> str: @@ -386,10 +380,6 @@ async def _validate_serving_cutover( ) for family_id in REQUIRED_FAMILIES } - deprecated_endpoint_status = { - path: await _endpoint_unreachable(client=client, url=path) - for path in DEPRECATED_OWNER_ENDPOINTS - } latest_run = runs[0] if isinstance(runs, list) and runs else current_run latest_family_results = dict(latest_run.get("family_results") or {}) release_metadata = dict((current_release.get("release") or {}).get("metadata") or {}) @@ -431,11 +421,9 @@ async def _validate_serving_cutover( bool(family_id and deployment_id and deployment_id == selected_by_family.get(family_id)) ) checks["workflow_composition_matches_family_winners"] = bool(composition_checks) and all(composition_checks) - checks["deprecated_endpoints_absent"] = all(deprecated_endpoint_status.values()) return { "status": _status_from_checks(checks), "checks": checks, - "deprecated_endpoint_status": deprecated_endpoint_status, "workflow_composition_registry": composition, } diff --git a/tests/test_staging_validation.py b/tests/test_staging_validation.py index 2f476d3..c817bb5 100644 --- a/tests/test_staging_validation.py +++ b/tests/test_staging_validation.py @@ -247,7 +247,6 @@ async def fake_fetch_json( "benchmark_version": f"{family_id}_family_v2", "rubric_version": f"{family_id}_family_rubric_v2", "retrieval_environment": {"mode": "live_web"} if family_id == "analyst" else {}, - "judge_config": {"model": "judge-model"} if family_id == "analyst" else {}, "evaluation_bundle": { "kind": "family_evaluation_bundle", "tasks": ( @@ -451,7 +450,6 @@ async def fake_endpoint_unreachable(**kwargs: Any) -> bool: assert report["serving_cutover"]["checks"]["builder_registry_matches_release"] is True assert report["serving_cutover"]["checks"]["builder_winner_matches_serving_release"] is True assert report["serving_cutover"]["checks"]["workflow_composition_matches_family_winners"] is True - assert report["serving_cutover"]["checks"]["deprecated_endpoints_absent"] is True @pytest.mark.asyncio diff --git a/tests/validator/calibration/__init__.py b/tests/validator/calibration/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/validator/calibration/test_json_parse_rate_gate.py b/tests/validator/calibration/test_json_parse_rate_gate.py new file mode 100644 index 0000000..9c144af --- /dev/null +++ b/tests/validator/calibration/test_json_parse_rate_gate.py @@ -0,0 +1,199 @@ +"""JSON parse-rate gate tests. + +Verifies the harness logic — pass/marginal/fail thresholds, error +classification, empty-fixture guard. Operators run the same harness +with real provider keys at deploy time. +""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from validation.validator.calibration.gate_result import GateResult +from validation.validator.calibration.json_parse_rate import ( + DEFAULT_PARSE_RATE_THRESHOLD, + SWAP_THRESHOLD, + JsonParseRateFixture, + JsonParseRateGate, + measure_json_parse_rate, +) +from validation.validator.providers.types import ( + ProviderError, + ProviderResponse, +) + + +pytestmark = pytest.mark.asyncio + + +class _RotatingClient: + """Provider-client stub returning a fixed sequence of responses. + + Each ``complete_structured`` call advances through the list, + raising/returning whatever is at the current index. Lets us + script a "37 successes, 3 failures" run for threshold tests. + """ + + def __init__( + self, + responses: list[str | Exception], + ) -> None: + self._responses = list(responses) + self._idx = 0 + + async def complete_structured( + self, + *, + system: str, + user: str, + response_schema: dict[str, Any], + temperature: float = 0.0, + max_tokens: int | None = None, + schema_name: str = "response", + ) -> ProviderResponse: + item = self._responses[self._idx % len(self._responses)] + self._idx += 1 + if isinstance(item, Exception): + raise item + return ProviderResponse( + text=item, latency_ms=10, usage_usd=0.0, finish_reason="stop", + ) + + +def _fixture(n: int = 1) -> list[JsonParseRateFixture]: + return [ + JsonParseRateFixture( + system="be terse", user=f"prompt {i}", + response_schema={"type": "object"}, + ) + for i in range(n) + ] + + +# -- pass / marginal / fail thresholds ----------------------------------- + + +async def test_all_parsed_returns_pass(): + """All fixtures return valid JSON → status=pass.""" + client = _RotatingClient([json.dumps({"ok": True})]) + result = await measure_json_parse_rate(client, _fixture(20)) + assert result.status == "pass" + assert result.measured_rate == 1.0 + assert result.n_samples == 20 + assert result.details["n_parsed"] == 20 + assert result.details["n_failed"] == 0 + + +async def test_just_above_threshold_passes(): + """98 valid + 2 invalid = 98% parse rate, exactly at default + threshold → pass.""" + responses: list[str | Exception] = [json.dumps({"ok": True})] * 98 + ["not json"] * 2 + client = _RotatingClient(responses) + result = await measure_json_parse_rate(client, _fixture(100)) + assert result.measured_rate == 0.98 + assert result.status == "pass" + + +async def test_marginal_zone_returns_marginal(): + """Parse rate ∈ [SWAP_THRESHOLD, default_threshold) → marginal. + Operator should wrap with JSON-repair retry before promoting.""" + # 92% parse rate: 92 valid + 8 malformed + responses: list[str | Exception] = ( + [json.dumps({"x": 1})] * 92 + ["malformed"] * 8 + ) + client = _RotatingClient(responses) + result = await measure_json_parse_rate(client, _fixture(100)) + assert result.measured_rate == 0.92 + assert SWAP_THRESHOLD <= result.measured_rate < DEFAULT_PARSE_RATE_THRESHOLD + assert result.status == "marginal" + + +async def test_below_swap_threshold_returns_fail(): + """Parse rate < SWAP_THRESHOLD (~90%) → fail. Model not ready; + JSON-repair retry can't recover this regime.""" + # 50% parse rate + responses: list[str | Exception] = ( + [json.dumps({"x": 1})] * 50 + ["bad"] * 50 + ) + client = _RotatingClient(responses) + result = await measure_json_parse_rate(client, _fixture(100)) + assert result.measured_rate == 0.5 + assert result.status == "fail" + + +async def test_custom_threshold_overrides_default(): + """Operator-supplied threshold takes precedence over the default.""" + client = _RotatingClient([json.dumps({"ok": True})] * 9 + ["bad"]) + # Default: 0.98 → 0.9 fails. Custom: 0.85 → 0.9 passes. + result_default = await measure_json_parse_rate(client, _fixture(10)) + assert result_default.status in ("marginal", "fail") + client2 = _RotatingClient([json.dumps({"ok": True})] * 9 + ["bad"]) + result_custom = await measure_json_parse_rate( + client2, _fixture(10), threshold=0.85, + ) + assert result_custom.status == "pass" + + +# -- error classification ------------------------------------------------ + + +async def test_provider_call_errors_count_as_parse_failures(): + """A ProviderError raised by the wrapped client is counted as a + failure, not propagated. The gate harness keeps going so a single + flaky call doesn't tank the run.""" + responses: list[str | Exception] = [ + json.dumps({"ok": True}), + ProviderError("boom"), + json.dumps({"ok": True}), + ] + client = _RotatingClient(responses) + result = await measure_json_parse_rate(client, _fixture(3)) + assert result.measured_rate == pytest.approx(2 / 3) + failure_details = result.details["failures"] + assert any("call_error" in (f.get("error") or "") for f in failure_details) + + +async def test_failed_calls_recorded_in_details(): + responses: list[str | Exception] = [ + json.dumps({"ok": True}), + "totally not json", + json.dumps({"ok": True}), + ] + client = _RotatingClient(responses) + result = await measure_json_parse_rate(client, _fixture(3)) + failures = result.details["failures"] + assert len(failures) == 1 + assert failures[0]["fixture_index"] == 1 + assert "parse_error" in failures[0]["error"] + + +# -- edge cases ---------------------------------------------------------- + + +async def test_empty_fixtures_returns_fail(): + client = _RotatingClient([json.dumps({"ok": True})]) + result = await measure_json_parse_rate(client, []) + assert result.status == "fail" + assert result.n_samples == 0 + assert result.details.get("reason") == "no_fixtures_provided" + + +async def test_gate_helper_runs_fixtures(): + """Operator-facing wrapper: hold fixtures + threshold, call + ``run(client)``.""" + client = _RotatingClient([json.dumps({"ok": True})] * 10) + gate = JsonParseRateGate(_fixture(10), threshold=0.95) + assert gate.threshold == 0.95 + assert gate.n_fixtures == 10 + result = await gate.run(client) + assert result.status == "pass" + + +async def test_gate_result_carries_passed_property(): + client = _RotatingClient([json.dumps({"ok": True})]) + result = await measure_json_parse_rate(client, _fixture(5)) + assert isinstance(result, GateResult) + assert result.passed is True diff --git a/tests/validator/calibration/test_rank_parity.py b/tests/validator/calibration/test_rank_parity.py new file mode 100644 index 0000000..3c8abd2 --- /dev/null +++ b/tests/validator/calibration/test_rank_parity.py @@ -0,0 +1,112 @@ +"""Validator rank parity (Spearman correlation) helper tests.""" + +from __future__ import annotations + +import pytest + +from validation.validator.calibration.rank_parity import rank_parity_spearman + + +# -- Identical / reversed / uncorrelated -------------------------------- + + +def test_identical_ranks_returns_one(): + """Same ordering → Spearman ρ = 1.0 (perfect rank agreement). + The actual numeric values can differ — Spearman is rank-based.""" + weighted = [0.1, 0.5, 0.9] + composite = [0.05, 0.40, 0.95] # different scale, same order + assert rank_parity_spearman(weighted, composite) == pytest.approx(1.0) + + +def test_reversed_ranks_returns_minus_one(): + """Opposite ordering → ρ = -1.0.""" + weighted = [0.1, 0.5, 0.9] + composite = [0.9, 0.5, 0.1] + assert rank_parity_spearman(weighted, composite) == pytest.approx(-1.0) + + +def test_uncorrelated_ranks_returns_zero_or_close(): + """Two random rank orderings have ρ near 0.0 over many samples. + For specific small samples, the value can be exact-zero (no + monotonic relationship).""" + weighted = [1, 2, 3, 4] + composite = [2, 4, 1, 3] + rho = rank_parity_spearman(weighted, composite) + # |ρ| ≤ 1; for this specific permutation, ρ = 0 + assert -0.5 <= rho <= 0.5 + + +# -- Production rollout scenarios --------------------------------------- + + +def test_high_correlation_above_rollout_threshold(): + """Real shadow-mode scenario: 10 miners ranked the same by both + paths with one slight reordering. ρ should be > 0.85 (the + rollout threshold from the plan).""" + weighted = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] + composite = [0.1, 0.2, 0.3, 0.5, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0] # 4 ↔ 5 swap + rho = rank_parity_spearman(weighted, composite) + assert rho > 0.85, f"ρ={rho} should clear the rollout threshold" + + +def test_modest_correlation_below_rollout_threshold(): + """A 4-of-10 reorder breaks correlation enough to fail the + rollout gate. Operator should investigate before flipping the + default scoring formula.""" + weighted = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] + # Aggressive reorder + composite = [0.5, 0.1, 0.7, 0.2, 0.9, 0.3, 0.6, 0.4, 0.8, 1.0] + rho = rank_parity_spearman(weighted, composite) + assert rho < 0.85, f"ρ={rho} should NOT clear the rollout threshold" + + +# -- Tie handling ------------------------------------------------------- + + +def test_ties_in_weighted_sum(): + """Average-rank for ties means a single tie doesn't produce a + spurious rank disagreement when the composite breaks the tie.""" + weighted = [0.5, 0.5, 0.9] # 1 and 2 are tied + composite = [0.4, 0.6, 0.9] # composite breaks the tie + rho = rank_parity_spearman(weighted, composite) + # Two miners with tied weighted scores get rank 1.5; composite + # ranks them 1 and 2. The third miner is rank 3 in both. + # Correlation is positive but not 1.0. + assert 0.5 < rho < 1.0 + + +def test_all_ties_returns_zero(): + """If every miner has the same score in one path, there's no + rank variance to correlate against → ρ = 0.""" + weighted = [0.5, 0.5, 0.5, 0.5] + composite = [0.1, 0.2, 0.3, 0.4] + assert rank_parity_spearman(weighted, composite) == 0.0 + + +# -- Edge cases --------------------------------------------------------- + + +def test_empty_inputs_return_zero(): + assert rank_parity_spearman([], []) == 0.0 + + +def test_single_element_returns_zero(): + """One miner — no rank correlation possible.""" + assert rank_parity_spearman([0.5], [0.7]) == 0.0 + + +def test_mismatched_lengths_return_zero(): + assert rank_parity_spearman([0.1, 0.2], [0.1, 0.2, 0.3]) == 0.0 + + +def test_average_rank_for_ties(): + """Verify the rank assignment internals: [10, 5, 5, 1] should + rank as [4.0, 2.5, 2.5, 1.0]. We test this indirectly via the + correlation result with a known counterpart.""" + weighted = [10, 5, 5, 1] + # Same ranks → ρ = 1 + composite = [10, 5, 5, 1] + assert rank_parity_spearman(weighted, composite) == pytest.approx(1.0) + # Reverse → ρ = -1 + composite_rev = [1, 5, 5, 10] + assert rank_parity_spearman(weighted, composite_rev) == pytest.approx(-1.0) diff --git a/tests/validator/calibration/test_reconciler_agreement_gate.py b/tests/validator/calibration/test_reconciler_agreement_gate.py new file mode 100644 index 0000000..db51117 --- /dev/null +++ b/tests/validator/calibration/test_reconciler_agreement_gate.py @@ -0,0 +1,203 @@ +"""Reconciler claim-set agreement gate tests.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import pytest + +from validation.validator.calibration.reconciler_agreement import ( + DEFAULT_AGREEMENT_THRESHOLD, + ReconcilerAgreementFixture, + ReconcilerAgreementGate, + claim_set_jaccard, + measure_reconciler_agreement, +) +from validation.validator.oracles.base import OracleGrounding +from validation.validator.reconciler import ReconciledOracle + + +pytestmark = pytest.mark.asyncio + + +class _ScriptedReconciler: + """Reconciler-protocol stub returning canned ReconciledOracles.""" + + def __init__(self, expected_claims_per_call: list[list[str]]) -> None: + self._queue = list(expected_claims_per_call) + + async def reconcile( + self, + *, + prompt: str, + groundings: list[OracleGrounding], + must_not_claim_floor: Iterable[str] = (), + ) -> ReconciledOracle: + if not self._queue: + raise RuntimeError("scripted reconciler exhausted") + claims = self._queue.pop(0) + return ReconciledOracle( + expected_claims=claims, + must_not_claim=list(must_not_claim_floor), + oracle_status="consensus", + consensus_claims=claims, + ) + + +def _grounding(vendor: str, text: str) -> OracleGrounding: + return OracleGrounding(vendor=vendor, status="ok", raw_text=text) + + +def _fixture(prompt: str, golden: list[str]) -> ReconcilerAgreementFixture: + return ReconcilerAgreementFixture( + prompt=prompt, + groundings=[ + _grounding("openai", "stub-a"), + _grounding("gemini", "stub-b"), + _grounding("grok", "stub-c"), + ], + golden_claims=golden, + ) + + +# -- Jaccard helper ----------------------------------------------------- + + +def test_jaccard_identical_sets(): + assert claim_set_jaccard(["Paris", "France"], ["Paris", "France"]) == 1.0 + + +def test_jaccard_disjoint_sets(): + assert claim_set_jaccard(["Paris"], ["Berlin"]) == 0.0 + + +def test_jaccard_partial_overlap(): + # 1 intersect / 3 union = 0.333... + j = claim_set_jaccard(["a", "b"], ["b", "c"]) + assert j == pytest.approx(1 / 3) + + +def test_jaccard_normalizes_case_and_whitespace(): + j = claim_set_jaccard(["Paris is in France"], ["paris is in france"]) + assert j == 1.0 + j2 = claim_set_jaccard([" hello world "], ["hello world"]) + assert j2 == 1.0 + + +def test_jaccard_both_empty_returns_one(): + """Both empty → no disagreement (trivially identical).""" + assert claim_set_jaccard([], []) == 1.0 + + +def test_jaccard_one_empty_returns_zero(): + assert claim_set_jaccard([], ["X"]) == 0.0 + assert claim_set_jaccard(["X"], []) == 0.0 + + +def test_jaccard_filters_empty_strings(): + """Empty/whitespace-only strings are ignored.""" + assert claim_set_jaccard(["X", "", " "], ["X"]) == 1.0 + + +# -- Gate runner -------------------------------------------------------- + + +async def test_perfect_agreement_passes(): + rec = _ScriptedReconciler( + [["Paris is the capital"], ["1969 was Apollo 11"], ["DNA has four bases"]] + ) + fixtures = [ + _fixture("capital of france?", ["Paris is the capital"]), + _fixture("when apollo 11?", ["1969 was Apollo 11"]), + _fixture("DNA bases?", ["DNA has four bases"]), + ] + result = await measure_reconciler_agreement(rec, fixtures) + assert result.status == "pass" + assert result.measured_rate == 1.0 + assert result.n_samples == 3 + assert result.details["min_jaccard"] == 1.0 + assert result.details["n_below_threshold"] == 0 + + +async def test_below_threshold_fails(): + """Reconciler outputs disagree with golden on most fixtures → + mean Jaccard < 0.85 → fail.""" + rec = _ScriptedReconciler( + [["wrong"], ["also wrong"], ["right answer"]] + ) + fixtures = [ + _fixture("Q1", ["right answer"]), + _fixture("Q2", ["right answer"]), + _fixture("Q3", ["right answer"]), + ] + # 2 of 3 disagree (Jaccard=0); 1 agrees (Jaccard=1) → mean ≈ 0.333 + result = await measure_reconciler_agreement(rec, fixtures) + assert result.status == "fail" + assert result.measured_rate == pytest.approx(1 / 3) + assert result.details["n_below_threshold"] == 2 + + +async def test_worst_fixtures_surfaced_in_details(): + """When the gate fails, the details include the worst fixtures so + the operator can debug curation vs reconciler model.""" + rec = _ScriptedReconciler( + [["A"], ["WRONG"], ["A"], ["A"], ["A"], ["WRONG"]] + ) + fixtures = [ + _fixture(f"Q{i}", ["A"]) for i in range(6) + ] + result = await measure_reconciler_agreement(rec, fixtures) + worst = result.details["worst_fixtures"] + assert len(worst) >= 1 + assert all(w["jaccard"] < DEFAULT_AGREEMENT_THRESHOLD for w in worst) + assert all("WRONG" in (w["expected_claims"] or [""])[0] for w in worst) + + +async def test_reconciler_error_counts_as_zero(): + """A reconciler call that raises an exception counts as zero + Jaccard for that fixture (worst case) and is recorded with the + error message.""" + + class _RaisingReconciler: + async def reconcile(self, **kwargs): + raise RuntimeError("simulated reconciler crash") + + fixtures = [_fixture("Q1", ["A"])] + result = await measure_reconciler_agreement(_RaisingReconciler(), fixtures) + assert result.measured_rate == 0.0 + assert result.status == "fail" + worst = result.details["worst_fixtures"] + assert len(worst) == 1 + assert "reconcile_error" in (worst[0]["error"] or "") + + +async def test_empty_fixtures_returns_fail(): + rec = _ScriptedReconciler([]) + result = await measure_reconciler_agreement(rec, []) + assert result.status == "fail" + assert result.n_samples == 0 + + +async def test_custom_threshold_overrides_default(): + """Operator can lower the bar for less-strict scenarios.""" + rec = _ScriptedReconciler([["A", "B"], ["A"]]) + fixtures = [ + _fixture("Q1", ["A", "B"]), + _fixture("Q2", ["A", "B"]), # Jaccard = 1/2 = 0.5 + ] + # mean = (1.0 + 0.5) / 2 = 0.75 + result_strict = await measure_reconciler_agreement(rec, fixtures) + assert result_strict.status == "fail" # 0.75 < 0.85 + rec2 = _ScriptedReconciler([["A", "B"], ["A"]]) + result_loose = await measure_reconciler_agreement( + rec2, fixtures, threshold=0.7, + ) + assert result_loose.status == "pass" + + +async def test_gate_helper_runs_fixtures(): + rec = _ScriptedReconciler([["A"]]) + gate = ReconcilerAgreementGate([_fixture("Q1", ["A"])]) + assert gate.n_fixtures == 1 + result = await gate.run(rec) + assert result.status == "pass" diff --git a/tests/validator/oracles/__init__.py b/tests/validator/oracles/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/validator/oracles/test_fanout.py b/tests/validator/oracles/test_fanout.py new file mode 100644 index 0000000..5dc67fe --- /dev/null +++ b/tests/validator/oracles/test_fanout.py @@ -0,0 +1,200 @@ +"""Oracle fanout tests. + +Covers: + * All-3-OK / 1-down / 2-down / all-down combinations → vendor_status_map + reflects the right shape; successful_groundings drops errors. + * Vendor order in results matches client construction order. + * Sequential mode produces same outputs as parallel. + * Last-resort safety net: an oracle client that *raises* (instead + of returning ``status="error"``) gets caught. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from validation.validator.oracles.base import ( + OracleClient, + OracleContext, + OracleGrounding, +) +from validation.validator.oracles.fanout import ( + OracleFanout, + successful_groundings, + vendor_status_map, +) + + +pytestmark = pytest.mark.asyncio + + +class _StubOracle(OracleClient): + """OracleClient that returns a pre-configured grounding.""" + + def __init__(self, vendor: str, grounding: OracleGrounding) -> None: + self._vendor = vendor + self._grounding = grounding + self.calls = 0 + + @property + def vendor(self) -> str: + return self._vendor + + async def produce_grounding(self, context: OracleContext) -> OracleGrounding: + self.calls += 1 + return self._grounding + + async def aclose(self) -> None: + pass + + +class _RaisingOracle(OracleClient): + """OracleClient that raises an unexpected exception (safety-net path).""" + + def __init__(self, vendor: str) -> None: + self._vendor = vendor + + @property + def vendor(self) -> str: + return self._vendor + + async def produce_grounding(self, context: OracleContext) -> OracleGrounding: + raise RuntimeError("simulated client bug") + + async def aclose(self) -> None: + pass + + +def _ctx() -> OracleContext: + return OracleContext(task_id="t1", prompt="What's 2+2?") + + +async def test_all_three_ok(): + a = _StubOracle("openai", OracleGrounding(vendor="openai", status="ok", raw_text="4")) + b = _StubOracle("gemini", OracleGrounding(vendor="gemini", status="ok", raw_text="4")) + c = _StubOracle("grok", OracleGrounding(vendor="grok", status="ok", raw_text="4")) + fan = OracleFanout([a, b, c]) + out = await fan.run(_ctx()) + + assert [g.vendor for g in out] == ["openai", "gemini", "grok"] + assert all(g.status == "ok" for g in out) + assert vendor_status_map(out) == {"openai": "ok", "gemini": "ok", "grok": "ok"} + assert len(successful_groundings(out)) == 3 + + +async def test_one_oracle_down_returns_two_ok_one_error(): + a = _StubOracle("openai", OracleGrounding(vendor="openai", status="ok", raw_text="4")) + b = _StubOracle("gemini", OracleGrounding(vendor="gemini", status="ok", raw_text="4")) + c = _StubOracle( + "grok", + OracleGrounding(vendor="grok", status="error", error_msg="timeout: vendor down"), + ) + fan = OracleFanout([a, b, c]) + out = await fan.run(_ctx()) + + assert vendor_status_map(out) == {"openai": "ok", "gemini": "ok", "grok": "error"} + ok = successful_groundings(out) + assert len(ok) == 2 + assert {g.vendor for g in ok} == {"openai", "gemini"} + + +async def test_two_oracles_down_returns_one_ok(): + a = _StubOracle("openai", OracleGrounding(vendor="openai", status="ok", raw_text="4")) + b = _StubOracle("gemini", OracleGrounding(vendor="gemini", status="error", error_msg="boom")) + c = _StubOracle("grok", OracleGrounding(vendor="grok", status="error", error_msg="boom")) + fan = OracleFanout([a, b, c]) + out = await fan.run(_ctx()) + + ok = successful_groundings(out) + assert len(ok) == 1 + assert ok[0].vendor == "openai" + assert vendor_status_map(out)["gemini"] == "error" + + +async def test_all_three_down_returns_no_ok(): + a = _StubOracle("openai", OracleGrounding(vendor="openai", status="error", error_msg="x")) + b = _StubOracle("gemini", OracleGrounding(vendor="gemini", status="error", error_msg="x")) + c = _StubOracle("grok", OracleGrounding(vendor="grok", status="error", error_msg="x")) + fan = OracleFanout([a, b, c]) + out = await fan.run(_ctx()) + + assert successful_groundings(out) == [] + assert vendor_status_map(out) == {"openai": "error", "gemini": "error", "grok": "error"} + + +async def test_blocked_status_isolated_from_error(): + """Gemini's safety-block surface as status=blocked, not error.""" + a = _StubOracle("openai", OracleGrounding(vendor="openai", status="ok", raw_text="x")) + b = _StubOracle( + "gemini", OracleGrounding(vendor="gemini", status="blocked", error_msg="SAFETY"), + ) + c = _StubOracle("grok", OracleGrounding(vendor="grok", status="ok", raw_text="x")) + fan = OracleFanout([a, b, c]) + out = await fan.run(_ctx()) + + statuses = vendor_status_map(out) + assert statuses == {"openai": "ok", "gemini": "blocked", "grok": "ok"} + # blocked is NOT an "ok" — successful_groundings drops it. + assert {g.vendor for g in successful_groundings(out)} == {"openai", "grok"} + + +async def test_safety_net_catches_unexpected_exception(): + """A buggy oracle client that raises (instead of returning + status=error) should be caught by the fanout layer's last-resort + try/except, surfacing as status=error rather than crashing the run.""" + a = _RaisingOracle("openai") + b = _StubOracle("gemini", OracleGrounding(vendor="gemini", status="ok", raw_text="4")) + fan = OracleFanout([a, b]) + out = await fan.run(_ctx()) + + a_grounding = next(g for g in out if g.vendor == "openai") + assert a_grounding.status == "error" + assert "unexpected_exception" in (a_grounding.error_msg or "") + b_grounding = next(g for g in out if g.vendor == "gemini") + assert b_grounding.status == "ok" + + +async def test_parallel_mode_runs_concurrently(): + """Sanity check: parallel mode kicks off all 3 calls before + any returns. Stub oracles each sleep briefly; total wall-clock + in parallel ≈ max(per-call), in sequential ≈ sum.""" + + class _SlowStub(OracleClient): + def __init__(self, vendor: str, sleep_s: float) -> None: + self._vendor = vendor + self._sleep_s = sleep_s + + @property + def vendor(self) -> str: + return self._vendor + + async def produce_grounding(self, context: OracleContext) -> OracleGrounding: + await asyncio.sleep(self._sleep_s) + return OracleGrounding(vendor=self._vendor, status="ok", raw_text="x") + + async def aclose(self) -> None: + pass + + clients = [_SlowStub(v, 0.05) for v in ("openai", "gemini", "grok")] + parallel_fan = OracleFanout(clients, parallel=True) + sequential_fan = OracleFanout(clients, parallel=False) + + t0 = asyncio.get_event_loop().time() + await parallel_fan.run(_ctx()) + t_par = asyncio.get_event_loop().time() - t0 + + t0 = asyncio.get_event_loop().time() + await sequential_fan.run(_ctx()) + t_seq = asyncio.get_event_loop().time() - t0 + + # Parallel should finish in ~0.05s vs ~0.15s sequential. Allow + # generous slack for CI scheduler variance. + assert t_par < 0.12, f"parallel took {t_par:.3f}s (expected ~0.05s)" + assert t_seq > 0.12, f"sequential took {t_seq:.3f}s (expected ~0.15s)" + + +async def test_empty_clients_rejected(): + with pytest.raises(ValueError): + OracleFanout([]) diff --git a/tests/validator/oracles/test_vendor_clients.py b/tests/validator/oracles/test_vendor_clients.py new file mode 100644 index 0000000..794cd5d --- /dev/null +++ b/tests/validator/oracles/test_vendor_clients.py @@ -0,0 +1,220 @@ +"""Per-vendor oracle client tests. + +Exercises the OpenAI / Gemini / Grok oracle wrappers against an +httpx-mocked provider client. Verifies: + * Successful structured response → OracleGrounding(status="ok"). + * Provider errors / timeouts → status="error" with diagnostic msg. + * Gemini safety-block surfaces as status="blocked". + * Malformed JSON in the answer field → status="error". +""" + +from __future__ import annotations + +import json + +import httpx +import pytest + +from validation.validator.eval_config import ProviderConfig +from validation.validator.oracles.base import OracleContext +from validation.validator.oracles.gemini_oracle import GeminiOracle +from validation.validator.oracles.grok_oracle import GrokOracle +from validation.validator.oracles.openai_oracle import OpenAIOracle +from validation.validator.providers.gemini import ( + DEFAULT_GEMINI_BASE_URL, + GeminiClient, +) +from validation.validator.providers.openai_compatible import ( + OpenAICompatibleClient, +) + + +pytestmark = pytest.mark.asyncio + + +def _openai_cfg() -> ProviderConfig: + return ProviderConfig( + base_url="http://openai.test", + api_key="tok", + model="gpt-5.4", + timeout_seconds=5.0, + max_tokens=512, + ) + + +def _grok_cfg() -> ProviderConfig: + return ProviderConfig( + base_url="http://grok.test", + api_key="tok", + model="grok-4.3", + timeout_seconds=5.0, + max_tokens=512, + ) + + +def _gemini_cfg() -> ProviderConfig: + return ProviderConfig( + base_url=DEFAULT_GEMINI_BASE_URL, + api_key="tok", + model="gemini-3.1-pro-preview", + timeout_seconds=5.0, + max_tokens=512, + ) + + +def _ok_openai(answer: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "choices": [{ + "message": {"content": json.dumps({"answer": answer})}, + "finish_reason": "stop", + }], + "usage": {"total_cost_usd": 0.001}, + }, + ) + + +def _ok_gemini(answer: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "candidates": [{ + "content": {"parts": [{"text": json.dumps({"answer": answer})}]}, + "finishReason": "STOP", + }], + }, + ) + + +def _ctx() -> OracleContext: + return OracleContext( + task_id="t1", + prompt="What is the capital of France?", + ) + + +# -- OpenAI oracle --------------------------------------------------------- + + +async def test_openai_oracle_ok(): + transport = httpx.MockTransport(lambda req: _ok_openai("Paris")) + client = OpenAICompatibleClient(_openai_cfg(), transport=transport) + oracle = OpenAIOracle(client=client) + + g = await oracle.produce_grounding(_ctx()) + await oracle.aclose() + + assert g.vendor == "openai" + assert g.status == "ok" + assert g.raw_text == "Paris" + assert g.cost_usd == pytest.approx(0.001) + assert g.finish_reason == "stop" + + +async def test_openai_oracle_timeout_surfaces_error(): + def handler(req: httpx.Request) -> httpx.Response: + raise httpx.TimeoutException("simulated") + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient( + _openai_cfg(), transport=transport, + max_retries=0, backoff_base_seconds=0.001, + ) + oracle = OpenAIOracle(client=client) + g = await oracle.produce_grounding(_ctx()) + await oracle.aclose() + + assert g.vendor == "openai" + assert g.status == "error" + assert "timeout" in (g.error_msg or "") + + +async def test_openai_oracle_malformed_answer_surfaces_error(): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "choices": [ + {"message": {"content": "not even close to JSON"}} + ], + }, + ) + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient(_openai_cfg(), transport=transport) + oracle = OpenAIOracle(client=client) + g = await oracle.produce_grounding(_ctx()) + await oracle.aclose() + + assert g.status == "error" + assert "malformed_response" in (g.error_msg or "") + + +async def test_openai_oracle_missing_answer_field_surfaces_error(): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={"choices": [{"message": {"content": json.dumps({"foo": "bar"})}}]}, + ) + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient(_openai_cfg(), transport=transport) + oracle = OpenAIOracle(client=client) + g = await oracle.produce_grounding(_ctx()) + await oracle.aclose() + + assert g.status == "error" + + +# -- Grok oracle (same shape as OpenAI) ------------------------------------ + + +async def test_grok_oracle_ok(): + transport = httpx.MockTransport(lambda req: _ok_openai("Paris")) + client = OpenAICompatibleClient(_grok_cfg(), transport=transport) + oracle = GrokOracle(client=client) + g = await oracle.produce_grounding(_ctx()) + await oracle.aclose() + + assert g.vendor == "grok" + assert g.status == "ok" + assert g.raw_text == "Paris" + + +# -- Gemini oracle --------------------------------------------------------- + + +async def test_gemini_oracle_ok(): + transport = httpx.MockTransport(lambda req: _ok_gemini("Paris")) + client = GeminiClient(_gemini_cfg(), transport=transport) + oracle = GeminiOracle(client=client) + + g = await oracle.produce_grounding(_ctx()) + await oracle.aclose() + + assert g.vendor == "gemini" + assert g.status == "ok" + assert g.raw_text == "Paris" + assert g.finish_reason == "STOP" + + +async def test_gemini_oracle_safety_block_surfaces_blocked_status(): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "candidates": [], + "promptFeedback": {"blockReason": "SAFETY"}, + }, + ) + + transport = httpx.MockTransport(handler) + client = GeminiClient(_gemini_cfg(), transport=transport) + oracle = GeminiOracle(client=client) + g = await oracle.produce_grounding(_ctx()) + await oracle.aclose() + + assert g.vendor == "gemini" + assert g.status == "blocked" + assert "SAFETY" in (g.error_msg or "") diff --git a/tests/validator/providers/__init__.py b/tests/validator/providers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/validator/providers/test_gemini.py b/tests/validator/providers/test_gemini.py new file mode 100644 index 0000000..8db18eb --- /dev/null +++ b/tests/validator/providers/test_gemini.py @@ -0,0 +1,212 @@ +"""HTTP-level tests for the validator's Gemini client. + +Exercises: + * Successful generateContent call → ProviderResponse with text and + finish_reason populated. + * Retry on 429/502/503/504; non-retryable 4xx surfaces immediately. + * Timeout → ProviderTimeout. + * Safety-block (no candidates + promptFeedback.blockReason) surfaces + a ProviderError with the block reason. + * Request shape: ``contents`` + ``systemInstruction`` + camelCase + ``responseSchema`` + ``responseMimeType: application/json``. +""" + +from __future__ import annotations + +import json + +import httpx +import pytest + +from validation.validator.eval_config import ProviderConfig +from validation.validator.providers.gemini import ( + DEFAULT_GEMINI_BASE_URL, + GeminiClient, +) +from validation.validator.providers.types import ( + ProviderError, + ProviderResponse, + ProviderTimeout, +) + + +pytestmark = pytest.mark.asyncio + + +def _cfg(**overrides) -> ProviderConfig: + base = dict( + base_url=DEFAULT_GEMINI_BASE_URL, + api_key="tok", + model="gemini-3.1-pro-preview", + timeout_seconds=5.0, + max_tokens=512, + ) + base.update(overrides) + return ProviderConfig(**base) + + +def _ok_response(text: str, finish_reason: str = "STOP") -> httpx.Response: + return httpx.Response( + 200, + json={ + "candidates": [ + { + "content": {"parts": [{"text": text}]}, + "finishReason": finish_reason, + } + ], + "usageMetadata": { + "promptTokenCount": 100, "candidatesTokenCount": 50, + }, + }, + ) + + +async def test_complete_structured_request_shape(): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["body"] = json.loads(request.content.decode("utf-8")) + return _ok_response(json.dumps({"answer": "42"})) + + transport = httpx.MockTransport(handler) + client = GeminiClient(_cfg(), transport=transport) + + resp = await client.complete_structured( + system="you are a judge", + user='{"prompt": "foo"}', + response_schema={ + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + }, + ) + await client.aclose() + + assert isinstance(resp, ProviderResponse) + assert json.loads(resp.text) == {"answer": "42"} + assert resp.finish_reason == "STOP" + assert resp.usage_usd is None # Google doesn't surface USD + + # URL: ?key=... + :generateContent suffix + assert "/models/gemini-3.1-pro-preview:generateContent" in captured["url"] + assert "key=tok" in captured["url"] + body = captured["body"] + # System prompt → top-level systemInstruction (NOT a role) + assert body["systemInstruction"]["parts"][0]["text"] == "you are a judge" + # User content nested in contents[].parts[] + assert body["contents"][0]["role"] == "user" + assert body["contents"][0]["parts"][0]["text"] == '{"prompt": "foo"}' + # JSON-mode lives in generationConfig with camelCase keys + cfg = body["generationConfig"] + assert cfg["responseMimeType"] == "application/json" + assert cfg["responseSchema"]["type"] == "object" + assert cfg["temperature"] == 0.0 + assert cfg["maxOutputTokens"] == 512 + + +async def test_retry_succeeds_after_503(): + calls = {"n": 0} + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + if calls["n"] == 1: + return httpx.Response(503, json={"error": {"message": "transient"}}) + return _ok_response('{"ok":true}') + + transport = httpx.MockTransport(handler) + client = GeminiClient( + _cfg(), transport=transport, + max_retries=2, backoff_base_seconds=0.001, + ) + resp = await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + assert calls["n"] == 2 + assert json.loads(resp.text) == {"ok": True} + + +async def test_retry_exhausted_raises_provider_error(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": {"message": "always 503"}}) + + transport = httpx.MockTransport(handler) + client = GeminiClient( + _cfg(), transport=transport, + max_retries=2, backoff_base_seconds=0.001, + ) + with pytest.raises(ProviderError) as exc: + await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + assert "503" in str(exc.value) + + +async def test_non_retryable_4xx_raises_immediately(): + calls = {"n": 0} + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + return httpx.Response(400, json={"error": {"message": "bad request"}}) + + transport = httpx.MockTransport(handler) + client = GeminiClient( + _cfg(), transport=transport, + max_retries=3, backoff_base_seconds=0.001, + ) + with pytest.raises(ProviderError) as exc: + await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + assert calls["n"] == 1 + assert "400" in str(exc.value) + + +async def test_timeout_raises_provider_timeout(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.TimeoutException("simulated") + + transport = httpx.MockTransport(handler) + client = GeminiClient( + _cfg(), transport=transport, + max_retries=1, backoff_base_seconds=0.001, + ) + with pytest.raises(ProviderTimeout): + await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + + +async def test_safety_block_surfaces_block_reason(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "candidates": [], + "promptFeedback": {"blockReason": "SAFETY"}, + }, + ) + + transport = httpx.MockTransport(handler) + client = GeminiClient(_cfg(), transport=transport) + with pytest.raises(ProviderError) as exc: + await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + assert "SAFETY" in str(exc.value) + + +async def test_unconfigured_raises_at_init(): + with pytest.raises(ProviderError): + GeminiClient(_cfg(api_key="")) diff --git a/tests/validator/providers/test_json_repair.py b/tests/validator/providers/test_json_repair.py new file mode 100644 index 0000000..3bd16b7 --- /dev/null +++ b/tests/validator/providers/test_json_repair.py @@ -0,0 +1,184 @@ +"""JSON-repair retry wrapper tests. + +The wrapper sits in front of the provider client (OpenAI / Gemini / +Chutes). On the 90-98% parse-rate band, a 2-retry repair loop turns +flaky models into usable ones. Below 90%, the model isn't ready and +should be swapped; the wrapper is for the recoverable zone. +""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from validation.validator.providers.json_repair import ( + JsonRepairClient, + with_json_repair, +) +from validation.validator.providers.types import ( + ProviderError, + ProviderResponse, +) + + +pytestmark = pytest.mark.asyncio + + +class _ScriptedClient: + """Provider-client stub that returns scripted responses in order.""" + + def __init__(self, responses: list[str]) -> None: + self._responses = list(responses) + self.calls: list[dict[str, Any]] = [] + + async def complete_structured( + self, + *, + system: str, + user: str, + response_schema: dict[str, Any], + temperature: float = 0.0, + max_tokens: int | None = None, + schema_name: str = "response", + ) -> ProviderResponse: + self.calls.append( + { + "system": system, + "user": user, + "schema": response_schema, + "schema_name": schema_name, + } + ) + if not self._responses: + raise RuntimeError("scripted client exhausted") + text = self._responses.pop(0) + return ProviderResponse( + text=text, latency_ms=10, usage_usd=0.0, finish_reason="stop", + ) + + +# -- Happy path ---------------------------------------------------------- + + +async def test_first_call_succeeds_no_retry(): + """Valid JSON on the first call → no retries, no repair prompt.""" + valid = json.dumps({"answer": "ok"}) + inner = _ScriptedClient([valid]) + wrapper = with_json_repair(inner, max_retries=2) + resp = await wrapper.complete_structured( + system="s", user="u", response_schema={"type": "object"}, + ) + assert json.loads(resp.text) == {"answer": "ok"} + assert len(inner.calls) == 1 + + +# -- Recovery zone (90-98%) ---------------------------------------------- + + +async def test_first_call_malformed_second_succeeds(): + """First response is invalid JSON; wrapper retries once and + succeeds. Caller sees the second response.""" + inner = _ScriptedClient(["not json", json.dumps({"answer": "ok"})]) + wrapper = with_json_repair(inner, max_retries=2) + resp = await wrapper.complete_structured( + system="s", user="please answer", response_schema={"type": "object"}, + ) + assert json.loads(resp.text) == {"answer": "ok"} + assert len(inner.calls) == 2 + # Repair instruction is appended to the user prompt on retry. + repair_user = inner.calls[1]["user"] + assert "please answer" in repair_user # original ask preserved + assert "malformed JSON" in repair_user + assert "STRICT JSON" in repair_user + + +async def test_two_retries_consumed_then_succeeds(): + inner = _ScriptedClient(["malformed", "still bad", json.dumps({"a": 1})]) + wrapper = with_json_repair(inner, max_retries=2) + resp = await wrapper.complete_structured( + system="s", user="u", response_schema={"type": "object"}, + ) + assert json.loads(resp.text) == {"a": 1} + assert len(inner.calls) == 3 + + +# -- Below recovery zone ------------------------------------------------- + + +async def test_all_retries_fail_raises_provider_error(): + """3 attempts (initial + 2 retries) all malformed → ProviderError. + The caller (judge / reconciler) treats this as a failed call and + surfaces 'disputed' or equivalent fallback.""" + inner = _ScriptedClient(["bad", "still bad", "no really bad"]) + wrapper = with_json_repair(inner, max_retries=2) + with pytest.raises(ProviderError, match="JSON parse failed after 3"): + await wrapper.complete_structured( + system="s", user="u", response_schema={"type": "object"}, + ) + assert len(inner.calls) == 3 + + +# -- Edge cases ---------------------------------------------------------- + + +async def test_zero_retries_means_one_attempt_total(): + """max_retries=0 → no retries, original call only. Useful when + the caller wants to know about parse failures immediately + without paying for repairs.""" + inner = _ScriptedClient(["malformed"]) + wrapper = with_json_repair(inner, max_retries=0) + with pytest.raises(ProviderError, match="JSON parse failed after 1"): + await wrapper.complete_structured( + system="s", user="u", response_schema={"type": "object"}, + ) + assert len(inner.calls) == 1 + + +async def test_repair_prompt_includes_parse_error_message(): + """The repair instruction includes the parse error so the model + sees what failed. Caps the error text at 200 chars to keep the + prompt small.""" + inner = _ScriptedClient(["malformed", json.dumps({"x": 1})]) + wrapper = with_json_repair(inner, max_retries=2) + await wrapper.complete_structured( + system="s", user="u", response_schema={"type": "object"}, + ) + repair_user = inner.calls[1]["user"] + # Some recognizable token from the JSONDecodeError message. + assert any( + marker in repair_user + for marker in ("Expecting", "Extra data", "char", "line") + ) + + +async def test_kwargs_pass_through_to_inner_client(): + """temperature / max_tokens / schema_name forwarded as-is.""" + inner = _ScriptedClient([json.dumps({"a": 1})]) + wrapper = with_json_repair(inner, max_retries=2) + schema = {"type": "object", "properties": {"a": {"type": "integer"}}} + await wrapper.complete_structured( + system="sys", + user="usr", + response_schema=schema, + temperature=0.7, + max_tokens=512, + schema_name="my_schema", + ) + call = inner.calls[0] + assert call["system"] == "sys" + assert call["user"] == "usr" + assert call["schema"] == schema + assert call["schema_name"] == "my_schema" + + +async def test_empty_string_treated_as_malformed(): + """Empty response is also malformed JSON — triggers repair.""" + inner = _ScriptedClient(["", json.dumps({"x": 1})]) + wrapper = with_json_repair(inner, max_retries=2) + resp = await wrapper.complete_structured( + system="s", user="u", response_schema={"type": "object"}, + ) + assert json.loads(resp.text) == {"x": 1} + assert len(inner.calls) == 2 diff --git a/tests/validator/providers/test_openai_compatible.py b/tests/validator/providers/test_openai_compatible.py new file mode 100644 index 0000000..46f5415 --- /dev/null +++ b/tests/validator/providers/test_openai_compatible.py @@ -0,0 +1,199 @@ +"""HTTP-level tests for the validator's OpenAI-compatible client. + +Exercises: + * Successful structured-output call → ProviderResponse with text + + latency populated. + * Retry on 429/502/503/504 (bounded by max_retries). + * Non-retryable 4xx surfaces immediately as ProviderError. + * Timeout paths surface ProviderTimeout. + * ``content`` returned as a list of parts joins correctly (some + OpenAI-compatible providers emit this shape). +""" + +from __future__ import annotations + +import json + +import httpx +import pytest + +from validation.validator.eval_config import ProviderConfig +from validation.validator.providers.openai_compatible import ( + OpenAICompatibleClient, +) +from validation.validator.providers.types import ( + ProviderError, + ProviderResponse, + ProviderTimeout, +) + + +pytestmark = pytest.mark.asyncio + + +def _cfg(**overrides) -> ProviderConfig: + base = dict( + base_url="http://provider.test", + api_key="tok", + model="test-model", + timeout_seconds=5.0, + max_tokens=512, + ) + base.update(overrides) + return ProviderConfig(**base) + + +def _ok_response(content: str | list) -> httpx.Response: + return httpx.Response( + 200, + json={ + "choices": [ + { + "message": {"content": content}, + "finish_reason": "stop", + } + ], + "usage": {"total_cost_usd": 0.0123}, + }, + ) + + +async def test_complete_structured_returns_text_and_latency(): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["request"] = request + captured["body"] = json.loads(request.content.decode("utf-8")) + return _ok_response(json.dumps({"answer": "ok"})) + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient(_cfg(), transport=transport) + + resp = await client.complete_structured( + system="you are a judge", + user=json.dumps({"prompt": "foo"}), + response_schema={"type": "object", "properties": {"answer": {"type": "string"}}}, + ) + await client.aclose() + + assert isinstance(resp, ProviderResponse) + assert json.loads(resp.text) == {"answer": "ok"} + assert resp.latency_ms >= 0 + assert resp.usage_usd == pytest.approx(0.0123) + assert resp.finish_reason == "stop" + + body = captured["body"] + assert body["model"] == "test-model" + assert body["temperature"] == 0.0 + assert body["max_tokens"] == 512 + assert body["response_format"]["type"] == "json_schema" + assert body["response_format"]["json_schema"]["strict"] is True + assert body["response_format"]["json_schema"]["schema"]["type"] == "object" + assert body["messages"][0]["role"] == "system" + assert body["messages"][1]["role"] == "user" + + +async def test_retry_succeeds_after_503(): + calls = {"n": 0} + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + if calls["n"] == 1: + return httpx.Response(503, json={"error": "transient"}) + return _ok_response(json.dumps({"ok": True})) + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient( + _cfg(), transport=transport, + max_retries=2, backoff_base_seconds=0.001, + ) + resp = await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + assert calls["n"] == 2 + assert json.loads(resp.text) == {"ok": True} + + +async def test_retry_exhausted_raises_provider_error(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": "always 503"}) + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient( + _cfg(), transport=transport, + max_retries=2, backoff_base_seconds=0.001, + ) + with pytest.raises(ProviderError) as exc: + await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + assert "503" in str(exc.value) + + +async def test_non_retryable_4xx_raises_immediately(): + calls = {"n": 0} + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + return httpx.Response(400, json={"error": "bad request"}) + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient( + _cfg(), transport=transport, + max_retries=3, backoff_base_seconds=0.001, + ) + with pytest.raises(ProviderError) as exc: + await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + assert calls["n"] == 1 # NOT retried + assert "400" in str(exc.value) + + +async def test_timeout_raises_provider_timeout(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.TimeoutException("simulated") + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient( + _cfg(), transport=transport, + max_retries=1, backoff_base_seconds=0.001, + ) + with pytest.raises(ProviderTimeout): + await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + + +async def test_content_as_parts_list_joins(): + """Some OpenAI-compatible providers return content as + [{type: 'text', text: '...'}, ...]. Client must concatenate text.""" + parts = [ + {"type": "text", "text": '{"a":'}, + {"type": "text", "text": ' "b"}'}, + ] + + def handler(request: httpx.Request) -> httpx.Response: + return _ok_response(parts) + + transport = httpx.MockTransport(handler) + client = OpenAICompatibleClient(_cfg(), transport=transport) + resp = await client.complete_structured( + system="s", user="u", + response_schema={"type": "object"}, + ) + await client.aclose() + assert json.loads(resp.text) == {"a": "b"} + + +async def test_unconfigured_raises_at_init(): + with pytest.raises(ProviderError): + OpenAICompatibleClient(_cfg(api_key="")) diff --git a/tests/validator/test_engine_oracle_helpers.py b/tests/validator/test_engine_oracle_helpers.py new file mode 100644 index 0000000..f9a0d5b --- /dev/null +++ b/tests/validator/test_engine_oracle_helpers.py @@ -0,0 +1,312 @@ +"""Engine.py oracle-enrichment integration helpers. + +Covers the building blocks ``_enrich_task_oracle``, +``_fetch_ledger_tools``, and ``_build_oracle_layer``. These are the +units the validator's ``_evaluate_task`` and ``_judge_miner`` consume +to wire oracle enrichment + composite scoring; the full closure +chain is exercised end-to-end in production smoke — mocking the +closure's many captured dependencies isn't worth the test friction. +""" + +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import patch + +import httpx +import pytest + +from validation.validator.engine import ( + _build_oracle_layer, + _enrich_task_oracle, + _fetch_ledger_tools, +) +from validation.validator.oracles.base import OracleGrounding +from validation.validator.reconciler import ReconciledOracle, Reconciler + + +pytestmark = pytest.mark.asyncio + + +def _stub_signer(token: str = "hk-test-validator") -> SimpleNamespace: + """Minimal signer stub. Returns a fixed Authorization header so the + test handlers can confirm the request was hotkey-signed.""" + return SimpleNamespace( + signed_headers=lambda method, path, body_hash: { + "Authorization": f"Hotkey {token}", + "X-Eirel-Hotkey": token, + }, + hotkey=token, + ) + + +# -- _enrich_task_oracle -------------------------------------------------- + + +def _task(*, oracle_source: str | None = None, **expected_output) -> SimpleNamespace: + """Build the duck-typed task_obj the validator engine creates.""" + return SimpleNamespace( + task_id="t-1", + prompt="What is 2+2?", + turns=None, + category="factual", + expected_output=expected_output, + oracle_source=oracle_source, + ) + + +async def test_deterministic_path_wraps_pool_answer(): + """``oracle_source=deterministic`` (or unset) builds a + ReconciledOracle from the task's pre-baked expected_output — + no LLM calls, no fanout, no reconciler needed.""" + task = _task(answer="4", must_not_claim=["five"]) + rec = await _enrich_task_oracle(task, fanout=None, reconciler=None) + assert rec.oracle_status == "deterministic" + assert rec.expected_claims == ["4"] + assert rec.must_not_claim == ["five"] + + +async def test_unset_oracle_source_treated_as_deterministic(): + task = _task(oracle_source=None, answer="4") + rec = await _enrich_task_oracle(task, fanout=None, reconciler=None) + assert rec.oracle_status == "deterministic" + + +async def test_three_oracle_without_layer_falls_back_to_disputed(): + """Operator forgot to configure oracle creds — three_oracle items + surface as ``disputed`` with the template floor preserved instead + of crashing the run.""" + task = _task(oracle_source="three_oracle", must_not_claim=["never X"]) + rec = await _enrich_task_oracle(task, fanout=None, reconciler=None) + assert rec.oracle_status == "disputed" + assert rec.expected_claims == [] + assert rec.must_not_claim == ["never X"] + assert rec.disagreement_note == "oracle_layer_not_configured" + + +async def test_three_oracle_with_mocked_layer_runs_fanout_and_reconciler(): + """Happy path: fanout returns 3 OK groundings, reconciler emits + consensus claims; expected_claims surface on the result.""" + + class _StubFanout: + async def run(self, ctx): + return [ + OracleGrounding(vendor="openai", status="ok", raw_text="Paris"), + OracleGrounding(vendor="gemini", status="ok", raw_text="Paris"), + OracleGrounding(vendor="grok", status="ok", raw_text="Paris"), + ] + + class _StubReconciler: + async def reconcile(self, *, prompt, groundings, must_not_claim_floor): + return ReconciledOracle( + expected_claims=["Paris is the capital"], + must_not_claim=list(must_not_claim_floor), + oracle_status="consensus", + consensus_claims=["Paris is the capital"], + ) + + task = _task(oracle_source="three_oracle", must_not_claim=["never London"]) + rec = await _enrich_task_oracle( + task, fanout=_StubFanout(), reconciler=_StubReconciler(), + ) + assert rec.oracle_status == "consensus" + assert rec.expected_claims == ["Paris is the capital"] + assert rec.must_not_claim == ["never London"] + + +# -- _fetch_ledger_tools -------------------------------------------------- + + +def _patch_engine_async_client(transport: httpx.MockTransport): + """Replace ``engine.httpx.AsyncClient`` with one that always uses + the given mock transport. ``patch.object`` on the class itself + keeps the constructor's other kwargs (e.g. ``timeout``) intact.""" + + real = httpx.AsyncClient + + class _Fake(real): + def __init__(self, *args, **kwargs): + kwargs["transport"] = transport + super().__init__(*args, **kwargs) + + from validation.validator import engine as _engine + return patch.object(_engine.httpx, "AsyncClient", _Fake) + + +async def test_fetch_ledger_tools_happy_path(): + """Owner-api returns a ledger; helper extracts unique tool names + in arrival order. Request is hotkey-signed.""" + captured_headers: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured_headers.update(dict(request.headers)) + return httpx.Response( + 200, + json={ + "job_id": "job-abc", + "n_calls": 3, + "tool_calls": [ + {"tool_name": "web_search"}, + {"tool_name": "url_fetch"}, + {"tool_name": "web_search"}, # duplicate — dedup + ], + }, + ) + + transport = httpx.MockTransport(handler) + with _patch_engine_async_client(transport): + tools = await _fetch_ledger_tools( + "job-abc", owner_url="http://owner.test", signer=_stub_signer(), + ) + assert captured_headers.get("authorization") == "Hotkey hk-test-validator" + assert tools == ["web_search", "url_fetch"] + + +async def test_fetch_ledger_tools_no_job_id_returns_empty(): + tools = await _fetch_ledger_tools( + "", owner_url="http://owner.test", signer=_stub_signer(), + ) + assert tools == [] + + +async def test_fetch_ledger_tools_5xx_returns_empty(): + """Owner-api error → empty list (composite gets 0 + tool_attestation_factor for tasks where required_tool is set — + fail-safe behavior).""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": "transient"}) + + transport = httpx.MockTransport(handler) + with _patch_engine_async_client(transport): + tools = await _fetch_ledger_tools( + "job-abc", owner_url="http://owner.test", signer=_stub_signer(), + ) + assert tools == [] + + +async def test_fetch_ledger_tools_network_error_returns_empty(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("simulated") + + transport = httpx.MockTransport(handler) + with _patch_engine_async_client(transport): + tools = await _fetch_ledger_tools( + "job-abc", owner_url="http://owner.test", signer=_stub_signer(), + ) + assert tools == [] + + +async def test_fetch_ledger_tools_empty_ledger(): + """Job had no tool calls — empty list, no error.""" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={"job_id": "job-x", "n_calls": 0, "tool_calls": []}, + ) + + transport = httpx.MockTransport(handler) + with _patch_engine_async_client(transport): + tools = await _fetch_ledger_tools( + "job-x", owner_url="http://owner.test", signer=_stub_signer(), + ) + assert tools == [] + + +# -- _build_oracle_layer -------------------------------------------------- + + +def test_build_oracle_layer_returns_none_when_no_creds(monkeypatch): + """Without any provider creds configured, the layer returns + (None, None). All three_oracle items will fall back to disputed + in ``_enrich_task_oracle``.""" + for prefix in ( + "EIREL_VALIDATOR_ORACLE_OPENAI_", + "EIREL_VALIDATOR_ORACLE_GEMINI_", + "EIREL_VALIDATOR_ORACLE_GROK_", + "EIREL_VALIDATOR_RECONCILER_", + ): + for suffix in ("BASE_URL", "API_KEY", "MODEL"): + monkeypatch.delenv(f"{prefix}{suffix}", raising=False) + fanout, reconciler = _build_oracle_layer() + assert fanout is None + assert reconciler is None + + +def test_build_oracle_layer_returns_none_when_only_one_oracle_configured(monkeypatch): + """Only OpenAI configured, no Gemini/Grok → fewer than 2 oracles, + layer returns None (plurality voting not possible with 1 vote).""" + for prefix in ( + "EIREL_VALIDATOR_ORACLE_GEMINI_", + "EIREL_VALIDATOR_ORACLE_GROK_", + ): + for suffix in ("BASE_URL", "API_KEY", "MODEL"): + monkeypatch.delenv(f"{prefix}{suffix}", raising=False) + monkeypatch.setenv( + "EIREL_VALIDATOR_ORACLE_OPENAI_BASE_URL", "http://openai.test", + ) + monkeypatch.setenv("EIREL_VALIDATOR_ORACLE_OPENAI_API_KEY", "tok") + monkeypatch.setenv("EIREL_VALIDATOR_ORACLE_OPENAI_MODEL", "gpt-5.4") + monkeypatch.setenv( + "EIREL_VALIDATOR_RECONCILER_BASE_URL", "http://chutes.test", + ) + monkeypatch.setenv("EIREL_VALIDATOR_RECONCILER_API_KEY", "tok") + monkeypatch.setenv( + "EIREL_VALIDATOR_RECONCILER_MODEL", "zai-org/GLM-5.1-TEE", + ) + fanout, reconciler = _build_oracle_layer() + assert fanout is None + assert reconciler is None + + +def test_build_oracle_layer_succeeds_with_two_oracles_plus_reconciler(monkeypatch): + """Two oracles + reconciler all configured → layer comes up. + Grok missing is the realistic case (Grok had more downtime + historically); fanout still works with 2-of-3.""" + monkeypatch.delenv("EIREL_VALIDATOR_ORACLE_GROK_BASE_URL", raising=False) + monkeypatch.delenv("EIREL_VALIDATOR_ORACLE_GROK_API_KEY", raising=False) + monkeypatch.setenv( + "EIREL_VALIDATOR_ORACLE_OPENAI_BASE_URL", "http://openai.test", + ) + monkeypatch.setenv("EIREL_VALIDATOR_ORACLE_OPENAI_API_KEY", "tok") + monkeypatch.setenv("EIREL_VALIDATOR_ORACLE_OPENAI_MODEL", "gpt-5.4") + monkeypatch.setenv( + "EIREL_VALIDATOR_ORACLE_GEMINI_BASE_URL", "http://gemini.test", + ) + monkeypatch.setenv("EIREL_VALIDATOR_ORACLE_GEMINI_API_KEY", "tok") + monkeypatch.setenv( + "EIREL_VALIDATOR_ORACLE_GEMINI_MODEL", "gemini-3.1-pro-preview", + ) + monkeypatch.setenv( + "EIREL_VALIDATOR_RECONCILER_BASE_URL", "http://chutes.test", + ) + monkeypatch.setenv("EIREL_VALIDATOR_RECONCILER_API_KEY", "tok") + monkeypatch.setenv( + "EIREL_VALIDATOR_RECONCILER_MODEL", "zai-org/GLM-5.1-TEE", + ) + fanout, reconciler = _build_oracle_layer() + assert fanout is not None + assert reconciler is not None + assert fanout.vendors == ["openai", "gemini"] + + +def test_build_oracle_layer_returns_none_when_reconciler_missing(monkeypatch): + """All 3 oracles configured but no reconciler creds → no layer. + Reconciler is the synthesis step; without it the oracles are + unusable.""" + for prefix in ( + "EIREL_VALIDATOR_ORACLE_OPENAI_", + "EIREL_VALIDATOR_ORACLE_GEMINI_", + "EIREL_VALIDATOR_ORACLE_GROK_", + ): + monkeypatch.setenv(f"{prefix}BASE_URL", "http://test") + monkeypatch.setenv(f"{prefix}API_KEY", "tok") + monkeypatch.setenv(f"{prefix}MODEL", "test-model") + monkeypatch.delenv("EIREL_VALIDATOR_RECONCILER_BASE_URL", raising=False) + monkeypatch.delenv("EIREL_VALIDATOR_RECONCILER_API_KEY", raising=False) + fanout, reconciler = _build_oracle_layer() + assert fanout is None + assert reconciler is None diff --git a/tests/validator/test_judge_client_eval.py b/tests/validator/test_judge_client_eval.py new file mode 100644 index 0000000..16317e2 --- /dev/null +++ b/tests/validator/test_judge_client_eval.py @@ -0,0 +1,128 @@ +"""JudgeServiceClient ``judge_eval`` + ``judge_eval_composite`` adaptation. + +Covers: + * Request serialization to ``/v1/judge/eval`` + * Optional prompt vs. turns dispatch + * Composite endpoint shape +""" +from __future__ import annotations + +import json + +import httpx + +from shared.core.judge_client import JudgeServiceClient + + +def _make_client(handler) -> JudgeServiceClient: + transport = httpx.MockTransport(handler) + return JudgeServiceClient(base_url="http://mock", transport=transport) + + +def test_judge_eval_posts_full_payload_for_single_turn(): + captured: dict = {} + + def _handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + json={ + "outcome": "correct", + "failure_mode": None, + "guidance": "", + }, + ) + + client = _make_client(_handler) + result = client.judge_eval( + bundle={ + "question": "What is 2+2?", + "answers": ["The answer is 4."], + }, + expected_answer="4", + must_not_claim=["five", "six"], + oracle_source="three_oracle", + ) + assert captured["url"] == "http://mock/v1/judge/eval" + body = captured["body"] + assert body["bundle"]["question"] == "What is 2+2?" + assert body["bundle"]["answers"] == ["The answer is 4."] + assert body["expected_answer"] == "4" + assert body["must_not_claim"] == ["five", "six"] + assert body["oracle_source"] == "three_oracle" + assert result["outcome"] == "correct" + + +def test_judge_eval_passes_turns_for_multi_turn_item(): + captured: dict = {} + + def _handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + json={"outcome": "correct", "guidance": ""}, + ) + + client = _make_client(_handler) + turns = [ + {"role": "user", "content": "I work in Python."}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": "What language do I work in?"}, + ] + client.judge_eval( + bundle={ + "question": "What language do I work in?", + "conversation_recent": turns, + "answers": ["You work in Python."], + }, + expected_answer="Python", + oracle_source="deterministic", + ) + body = captured["body"] + assert body["bundle"]["conversation_recent"] == turns + assert body["bundle"]["answers"] == ["You work in Python."] + assert body["oracle_source"] == "deterministic" + + +def test_judge_eval_composite_posts_to_composite_endpoint(): + captured: dict = {} + + def _handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + json={ + "composite": 1.0, + "outcome_score": 1.0, + "tool_attestation_factor": 1.0, + "efficiency_factor": 1.0, + "hallucination_knockout": 1.0, + "cost_attestation_knockout": 1.0, + "knockout_reason": None, + }, + ) + + client = _make_client(_handler) + result = client.judge_eval_composite( + outcome="correct", + candidate_response="ok", + must_not_claim=[], + required_tool="web_search", + ledger_tools=["web_search"], + latency_ms=200, + cost_usd=0.001, + latency_budget_ms=5000, + cost_budget_usd=0.01, + cost_floor_usd=0.00005, + ) + assert captured["url"] == "http://mock/v1/judge/eval/composite" + body = captured["body"] + assert body["outcome"] == "correct" + assert body["required_tool"] == "web_search" + assert body["ledger_tools"] == ["web_search"] + assert body["latency_budget_ms"] == 5000 + assert body["cost_budget_usd"] == 0.01 + assert body["cost_floor_usd"] == 0.00005 + assert result["composite"] == 1.0 diff --git a/tests/validator/test_oracle_cache.py b/tests/validator/test_oracle_cache.py new file mode 100644 index 0000000..0cc0abd --- /dev/null +++ b/tests/validator/test_oracle_cache.py @@ -0,0 +1,94 @@ +"""TaskOracleCache tests. + +Cache is intentionally a plain dict — these tests exist to lock the +contract (set/get/require/clear) so future refactors don't silently +change behavior the validator engine depends on. +""" + +from __future__ import annotations + +import pytest + +from validation.validator.oracle_cache import TaskOracleCache +from validation.validator.reconciler import ReconciledOracle + + +def _make_reconciled(claims: list[str]) -> ReconciledOracle: + return ReconciledOracle( + expected_claims=claims, + oracle_status="consensus", + ) + + +def test_set_and_get(): + cache = TaskOracleCache() + rec = _make_reconciled(["X"]) + cache.set("task-1", rec) + assert cache.get("task-1") is rec + + +def test_get_missing_returns_none(): + cache = TaskOracleCache() + assert cache.get("absent") is None + + +def test_require_raises_on_missing(): + cache = TaskOracleCache() + with pytest.raises(KeyError): + cache.require("absent") + + +def test_require_returns_value_on_hit(): + cache = TaskOracleCache() + rec = _make_reconciled(["X"]) + cache.set("t", rec) + assert cache.require("t") is rec + + +def test_set_overwrites(): + cache = TaskOracleCache() + cache.set("t", _make_reconciled(["X"])) + cache.set("t", _make_reconciled(["Y"])) + assert cache.get("t").expected_claims == ["Y"] + + +def test_clear_drops_everything(): + cache = TaskOracleCache() + cache.set("a", _make_reconciled(["1"])) + cache.set("b", _make_reconciled(["2"])) + assert len(cache) == 2 + cache.clear() + assert len(cache) == 0 + assert cache.get("a") is None + assert cache.get("b") is None + + +def test_contains(): + cache = TaskOracleCache() + cache.set("present", _make_reconciled(["X"])) + assert "present" in cache + assert "absent" not in cache + + +def test_iter_yields_task_ids(): + cache = TaskOracleCache() + cache.set("a", _make_reconciled(["1"])) + cache.set("b", _make_reconciled(["2"])) + assert sorted(cache) == ["a", "b"] + + +def test_cache_lifetime_simulates_batch_processing(): + """Simulate the validator's batch flow: claim phase populates, + judge phase consumes (possibly N times for N miners), batch end + clears.""" + cache = TaskOracleCache() + # Claim phase: 3 tasks. + for tid in ("t1", "t2", "t3"): + cache.set(tid, _make_reconciled([f"answer-{tid}"])) + # Judge phase: each task consumed multiple times (per miner). + for _ in range(5): + for tid in ("t1", "t2", "t3"): + assert cache.require(tid).expected_claims == [f"answer-{tid}"] + # Batch end. + cache.clear() + assert len(cache) == 0 diff --git a/tests/validator/test_reconciler.py b/tests/validator/test_reconciler.py new file mode 100644 index 0000000..d7bbf60 --- /dev/null +++ b/tests/validator/test_reconciler.py @@ -0,0 +1,349 @@ +"""Reconciler tests. + +Exercises: + * 3 successful groundings → reconciler call → ReconciledOracle with + expected_claims = consensus + majority. + * 2 successful (1 down) → reconciler still runs. + * <2 successful → skip reconciler call, return disputed. + * Reconciler error → disputed with template floor preserved. + * Reconciler returns malformed JSON → disputed. + * Additive must_not_claim: floor + extras, deduplicated. + * from_deterministic constructor for non-three_oracle items. + * Telemetry fields populated (vendor_status, latency, cost). +""" + +from __future__ import annotations + +import json + +import httpx +import pytest + +from validation.validator.eval_config import ProviderConfig +from validation.validator.oracles.base import OracleGrounding +from validation.validator.providers.openai_compatible import ( + OpenAICompatibleClient, +) +from validation.validator.reconciler import ( + ReconciledOracle, + Reconciler, +) + + +pytestmark = pytest.mark.asyncio + + +def _cfg() -> ProviderConfig: + return ProviderConfig( + base_url="http://chutes.test", + api_key="tok", + model="zai-org/GLM-5.1-TEE", + timeout_seconds=5.0, + max_tokens=512, + ) + + +def _ok_response(parsed: dict) -> httpx.Response: + return httpx.Response( + 200, + json={ + "choices": [{"message": {"content": json.dumps(parsed)}}], + "usage": {"total_cost_usd": 0.0007}, + }, + ) + + +def _three_groundings() -> list[OracleGrounding]: + return [ + OracleGrounding(vendor="openai", status="ok", raw_text="Paris is the capital."), + OracleGrounding(vendor="gemini", status="ok", raw_text="The capital of France is Paris."), + OracleGrounding(vendor="grok", status="ok", raw_text="Paris."), + ] + + +async def test_reconciler_consensus_path(): + captured: dict = {} + + def handler(req: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(req.content.decode("utf-8")) + return _ok_response({ + "consensus_claims": ["Paris is the capital of France"], + "majority_claims": [], + "minority_claims": [], + "must_not_claim_extras": [], + "oracle_status": "consensus", + "disagreement_note": None, + }) + + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(handler)) + rec = Reconciler(client=client) + result = await rec.reconcile( + prompt="What is the capital of France?", + groundings=_three_groundings(), + ) + await rec.aclose() + + assert result.oracle_status == "consensus" + assert result.expected_claims == ["Paris is the capital of France"] + assert result.must_not_claim == [] + assert result.vendor_status == {"openai": "ok", "gemini": "ok", "grok": "ok"} + assert result.reconciler_cost_usd == pytest.approx(0.0007) + + # User prompt sent to reconciler must include both the user + # question AND each oracle answer. + user_msg = next( + m for m in captured["body"]["messages"] if m["role"] == "user" + ) + payload = json.loads(user_msg["content"]) + assert payload["user_prompt"] == "What is the capital of France?" + assert {a["vendor"] for a in payload["oracle_answers"]} == {"openai", "gemini", "grok"} + + +async def test_reconciler_majority_path(): + """2-of-3 oracles agree on a fact; reconciler emits majority_claim + + correctly maps it into expected_claims.""" + + def handler(req: httpx.Request) -> httpx.Response: + return _ok_response({ + "consensus_claims": ["Eiffel Tower is in Paris"], + "majority_claims": [ + {"claim": "Built in 1889", "supporting_oracles": ["openai", "gemini"]}, + ], + "minority_claims": [ + {"claim": "Tower is 312m tall", "supporting_oracle": "grok"}, + ], + "must_not_claim_extras": [], + "oracle_status": "majority", + "disagreement_note": "grok mentioned tower height, others didn't", + }) + + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(handler)) + rec = Reconciler(client=client) + result = await rec.reconcile( + prompt="Tell me about the Eiffel Tower", + groundings=_three_groundings(), + ) + await rec.aclose() + + assert result.oracle_status == "majority" + # consensus + majority → expected_claims (minority dropped from scoring) + assert result.expected_claims == [ + "Eiffel Tower is in Paris", + "Built in 1889", + ] + # minority claims preserved for telemetry + assert len(result.minority_claims) == 1 + assert result.minority_claims[0]["claim"] == "Tower is 312m tall" + assert result.minority_claims[0]["supporting_oracle"] == "grok" + assert result.disagreement_note == "grok mentioned tower height, others didn't" + + +async def test_reconciler_disputed_path(): + def handler(req: httpx.Request) -> httpx.Response: + return _ok_response({ + "consensus_claims": [], + "majority_claims": [], + "minority_claims": [ + {"claim": "Yes, X is true", "supporting_oracle": "openai"}, + {"claim": "Actually X is false", "supporting_oracle": "gemini"}, + ], + "must_not_claim_extras": [], + "oracle_status": "disputed", + "disagreement_note": "openai and gemini disagree on X", + }) + + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(handler)) + rec = Reconciler(client=client) + result = await rec.reconcile(prompt="Is X true?", groundings=_three_groundings()) + await rec.aclose() + + assert result.oracle_status == "disputed" + assert result.expected_claims == [] + assert result.disagreement_note == "openai and gemini disagree on X" + + +async def test_must_not_claim_is_additive(): + """Reconciler extras + template floor are unioned, dedup preserves + floor order so the floor ends up first in the final list.""" + + def handler(req: httpx.Request) -> httpx.Response: + return _ok_response({ + "consensus_claims": ["X is the answer"], + "majority_claims": [], + "minority_claims": [], + "must_not_claim_extras": [ + "claim about Y is wrong", + "do not assert Z", + ], + "oracle_status": "consensus", + "disagreement_note": None, + }) + + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(handler)) + rec = Reconciler(client=client) + result = await rec.reconcile( + prompt="What's the answer?", + groundings=_three_groundings(), + must_not_claim_floor=["never claim A", "do not assert Z"], # 'do not assert Z' overlaps + ) + await rec.aclose() + + # Floor entries first (preserving order), then non-overlapping extras. + assert result.must_not_claim == [ + "never claim A", + "do not assert Z", + "claim about Y is wrong", + ] + + +async def test_reconciler_skipped_when_fewer_than_two_oracles_succeeded(): + """1 ok + 2 errors → skip reconciler call, return disputed with + template floor.""" + calls = {"n": 0} + + def handler(req: httpx.Request) -> httpx.Response: + calls["n"] += 1 + return _ok_response({}) # would succeed if called + + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(handler)) + rec = Reconciler(client=client) + result = await rec.reconcile( + prompt="...", + groundings=[ + OracleGrounding(vendor="openai", status="ok", raw_text="Paris"), + OracleGrounding(vendor="gemini", status="error", error_msg="boom"), + OracleGrounding(vendor="grok", status="error", error_msg="boom"), + ], + must_not_claim_floor=["never X"], + ) + await rec.aclose() + + assert calls["n"] == 0 # reconciler call NOT made + assert result.oracle_status == "disputed" + assert result.expected_claims == [] + assert result.must_not_claim == ["never X"] + assert result.vendor_status == {"openai": "ok", "gemini": "error", "grok": "error"} + assert "1/3 oracles" in (result.disagreement_note or "") + + +async def test_reconciler_all_oracles_failed_returns_disputed(): + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(lambda r: _ok_response({}))) + rec = Reconciler(client=client) + result = await rec.reconcile( + prompt="...", + groundings=[ + OracleGrounding(vendor="openai", status="error", error_msg="x"), + OracleGrounding(vendor="gemini", status="error", error_msg="x"), + OracleGrounding(vendor="grok", status="error", error_msg="x"), + ], + ) + await rec.aclose() + + assert result.oracle_status == "disputed" + assert result.expected_claims == [] + + +async def test_reconciler_error_falls_back_to_disputed_with_floor(): + """Provider 5xx after retries → reconciler returns disputed with + expected_claims=[] and must_not_claim=template_floor only.""" + + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": "always 503"}) + + client = OpenAICompatibleClient( + _cfg(), transport=httpx.MockTransport(handler), + max_retries=1, backoff_base_seconds=0.001, + ) + rec = Reconciler(client=client) + result = await rec.reconcile( + prompt="...", + groundings=_three_groundings(), + must_not_claim_floor=["floor item"], + ) + await rec.aclose() + + assert result.oracle_status == "disputed" + assert result.expected_claims == [] + assert result.must_not_claim == ["floor item"] + assert "reconciler_error" in (result.disagreement_note or "") + + +async def test_reconciler_malformed_json_falls_back_to_disputed(): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={"choices": [{"message": {"content": "not json"}}]}, + ) + + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(handler)) + rec = Reconciler(client=client) + result = await rec.reconcile( + prompt="...", + groundings=_three_groundings(), + must_not_claim_floor=["floor"], + ) + await rec.aclose() + + assert result.oracle_status == "disputed" + assert result.must_not_claim == ["floor"] + assert "malformed_json" in (result.disagreement_note or "") + + +async def test_reconciler_invalid_status_downgrades_to_disputed(): + def handler(req: httpx.Request) -> httpx.Response: + return _ok_response({ + "consensus_claims": ["X"], + "majority_claims": [], + "minority_claims": [], + "must_not_claim_extras": [], + "oracle_status": "totally-bogus-status", + "disagreement_note": None, + }) + + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(handler)) + rec = Reconciler(client=client) + result = await rec.reconcile(prompt="...", groundings=_three_groundings()) + await rec.aclose() + + assert result.oracle_status == "disputed" + + +async def test_from_deterministic_factory(): + """Pool kinds with built-in graders skip the reconciler entirely.""" + rec = ReconciledOracle.from_deterministic( + answer="Paris", + must_not_claim_floor=["never claim London"], + ) + assert rec.oracle_status == "deterministic" + assert rec.expected_claims == ["Paris"] + assert rec.must_not_claim == ["never claim London"] + assert rec.consensus_claims == ["Paris"] + assert rec.vendor_status == {} + + +async def test_from_deterministic_with_empty_answer(): + rec = ReconciledOracle.from_deterministic(answer="") + assert rec.expected_claims == [] + assert rec.consensus_claims == [] + assert rec.oracle_status == "deterministic" + + +async def test_telemetry_populated_on_success(): + def handler(req: httpx.Request) -> httpx.Response: + return _ok_response({ + "consensus_claims": ["X"], + "majority_claims": [], + "minority_claims": [], + "must_not_claim_extras": [], + "oracle_status": "consensus", + "disagreement_note": None, + }) + + client = OpenAICompatibleClient(_cfg(), transport=httpx.MockTransport(handler)) + rec = Reconciler(client=client) + result = await rec.reconcile(prompt="...", groundings=_three_groundings()) + await rec.aclose() + + assert result.reconciler_latency_ms >= 0 + assert result.reconciler_cost_usd == pytest.approx(0.0007) + assert result.vendor_status == {"openai": "ok", "gemini": "ok", "grok": "ok"} diff --git a/tests/validator/test_validator_metrics.py b/tests/validator/test_validator_metrics.py new file mode 100644 index 0000000..2190492 --- /dev/null +++ b/tests/validator/test_validator_metrics.py @@ -0,0 +1,151 @@ +"""Validator observability metric helpers — emit counters / gauges / histograms.""" + +from __future__ import annotations + +import pytest + +from validation.validator import metrics + + +def _label_value(counter, **labels) -> float: + """Pull the numeric value out of a labelled prometheus counter/gauge.""" + return counter.labels(**labels)._value.get() # type: ignore[attr-defined] + + +def test_record_oracle_grounding_bumps_counter(): + before = _label_value( + metrics.oracle_grounding_outcomes_total, + vendor="openai", status="ok", + ) + metrics.record_oracle_grounding("openai", "ok") + after = _label_value( + metrics.oracle_grounding_outcomes_total, + vendor="openai", status="ok", + ) + assert after == before + 1.0 + + +def test_record_oracle_grounding_distinguishes_vendors(): + before_gemini = _label_value( + metrics.oracle_grounding_outcomes_total, + vendor="gemini", status="ok", + ) + before_openai = _label_value( + metrics.oracle_grounding_outcomes_total, + vendor="openai", status="ok", + ) + metrics.record_oracle_grounding("gemini", "ok") + after_gemini = _label_value( + metrics.oracle_grounding_outcomes_total, + vendor="gemini", status="ok", + ) + after_openai = _label_value( + metrics.oracle_grounding_outcomes_total, + vendor="openai", status="ok", + ) + assert after_gemini == before_gemini + 1.0 + assert after_openai == before_openai # not bumped + + +def test_record_oracle_status_per_status_bucket(): + for status in ("consensus", "majority", "disputed", "deterministic"): + before = _label_value( + metrics.oracle_status_per_run_total, + family="general_chat", oracle_status=status, + ) + metrics.record_oracle_status("general_chat", status) + after = _label_value( + metrics.oracle_status_per_run_total, + family="general_chat", oracle_status=status, + ) + assert after == before + 1.0, f"status={status} didn't bump" + + +def test_record_eval_outcome_normalizes_case(): + """Inputs may be uppercase; the metric label normalizes.""" + before = _label_value( + metrics.eval_outcome_total, + family="general_chat", outcome="correct", + ) + metrics.record_eval_outcome("general_chat", "CORRECT") + after = _label_value( + metrics.eval_outcome_total, + family="general_chat", outcome="correct", + ) + assert after == before + 1.0 + + +def test_record_composite_score_clamps_to_unit_range(): + """Defensive: out-of-range inputs don't crash. They get clamped.""" + metrics.record_composite_score("general_chat", -0.5) # → 0.0 + metrics.record_composite_score("general_chat", 1.5) # → 1.0 + metrics.record_composite_score("general_chat", 0.42) + # Histogram doesn't expose a clean per-bucket count; just verify + # we didn't crash. + + +def test_record_composite_knockout_per_factor(): + factors = ( + "tool_attestation", + "hallucination_knockout", + "cost_attestation_knockout", + "outcome_zero", + ) + for factor in factors: + before = _label_value( + metrics.composite_knockouts_total, + family="general_chat", knockout_factor=factor, + ) + metrics.record_composite_knockout("general_chat", factor) + after = _label_value( + metrics.composite_knockouts_total, + family="general_chat", knockout_factor=factor, + ) + assert after == before + 1.0 + + +def test_record_judge_json_malformation_per_role(): + for role in ("pairwise", "multi", "eval"): + before = _label_value( + metrics.judge_json_malformations_total, + judge_role=role, + ) + metrics.record_judge_json_malformation(role) + after = _label_value( + metrics.judge_json_malformations_total, + judge_role=role, + ) + assert after == before + 1.0 + + +def test_record_oracle_call_supports_batch_increment(): + before = _label_value( + metrics.oracle_layer_calls_total, vendor="openai", + ) + metrics.record_oracle_call("openai", n=5) + after = _label_value( + metrics.oracle_layer_calls_total, vendor="openai", + ) + assert after == before + 5.0 + + +def test_record_disputed_rate_clamps(): + metrics.record_disputed_rate("general_chat", 1.5) # → 1.0 + val_high = metrics.reconciler_disputed_rate.labels(family="general_chat")._value.get() + assert val_high == 1.0 + metrics.record_disputed_rate("general_chat", -0.1) # → 0.0 + val_low = metrics.reconciler_disputed_rate.labels(family="general_chat")._value.get() + assert val_low == 0.0 + metrics.record_disputed_rate("general_chat", 0.42) + val_mid = metrics.reconciler_disputed_rate.labels(family="general_chat")._value.get() + assert val_mid == pytest.approx(0.42) + + +def test_record_rank_parity_clamps_to_minus_one_to_one(): + """Spearman ρ ∈ [-1, 1]; clamp out-of-range defensively.""" + metrics.record_rank_parity("general_chat", 1.5) + assert metrics.rank_parity_spearman_gauge.labels(family="general_chat")._value.get() == 1.0 + metrics.record_rank_parity("general_chat", -2.0) + assert metrics.rank_parity_spearman_gauge.labels(family="general_chat")._value.get() == -1.0 + metrics.record_rank_parity("general_chat", 0.87) + assert metrics.rank_parity_spearman_gauge.labels(family="general_chat")._value.get() == pytest.approx(0.87) From f36690771cf7ba27b6973caaff6c62c8eb9d9521 Mon Sep 17 00:00:00 2001 From: kyron <266216617+kyron1112567@users.noreply.github.com> Date: Thu, 7 May 2026 18:31:20 +0000 Subject: [PATCH 24/24] refactor(tests): update OpenAI/Grok response handling in tests for consistency --- pyproject.toml | 10 ++++ .../validator/oracles/test_vendor_clients.py | 58 ++++++++++++++++--- 2 files changed, 61 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 91bc8d6..2ff6723 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,11 +36,21 @@ tracing = [ "opentelemetry-sdk>=1.25.0,<2", "opentelemetry-exporter-otlp-proto-grpc>=1.25.0,<2", ] +documents = [ + # Used by the orchestrator's document_extractor for PDF / DOCX + # ingestion. Optional at runtime — the extractor fails soft with + # ``status="unsupported"`` when these aren't importable. Pinned in + # ``dev`` so the happy-path tests can exercise real extraction. + "pypdf>=4.0,<6", + "python-docx>=1.1,<2", +] dev = [ "pytest>=9.0.2,<10", "pytest-asyncio>=1.3.0,<2", "moto[s3]>=5.0.0,<6", "fakeredis>=2.26.0,<3", + "pypdf>=4.0,<6", + "python-docx>=1.1,<2", ] [project.scripts] diff --git a/tests/validator/oracles/test_vendor_clients.py b/tests/validator/oracles/test_vendor_clients.py index 794cd5d..a6cfbbc 100644 --- a/tests/validator/oracles/test_vendor_clients.py +++ b/tests/validator/oracles/test_vendor_clients.py @@ -63,13 +63,29 @@ def _gemini_cfg() -> ProviderConfig: def _ok_openai(answer: str) -> httpx.Response: + """OpenAI/Grok ``/v1/responses`` payload shape (Responses API). + + Both vendors share the same wire format here — a top-level + ``output`` array of ``{type: message, status, content: [...]}`` + items, where each content part of type ``output_text`` carries the + text. See ``openai_compatible.py:_parse_responses_api_response``. + """ return httpx.Response( 200, json={ - "choices": [{ - "message": {"content": json.dumps({"answer": answer})}, - "finish_reason": "stop", - }], + "output": [ + { + "type": "message", + "status": "stop", + "content": [ + { + "type": "output_text", + "text": json.dumps({"answer": answer}), + "annotations": [], + } + ], + } + ], "usage": {"total_cost_usd": 0.001}, }, ) @@ -131,12 +147,25 @@ def handler(req: httpx.Request) -> httpx.Response: async def test_openai_oracle_malformed_answer_surfaces_error(): + """Server returned a valid Responses-API envelope but the + ``output_text`` is non-JSON garbage — oracle parses ``raw_text`` + successfully but fails the JSON-shape extraction step.""" def handler(req: httpx.Request) -> httpx.Response: return httpx.Response( 200, json={ - "choices": [ - {"message": {"content": "not even close to JSON"}} + "output": [ + { + "type": "message", + "status": "stop", + "content": [ + { + "type": "output_text", + "text": "not even close to JSON", + "annotations": [], + } + ], + } ], }, ) @@ -152,10 +181,25 @@ def handler(req: httpx.Request) -> httpx.Response: async def test_openai_oracle_missing_answer_field_surfaces_error(): + """JSON parses fine but the required ``answer`` field is absent.""" def handler(req: httpx.Request) -> httpx.Response: return httpx.Response( 200, - json={"choices": [{"message": {"content": json.dumps({"foo": "bar"})}}]}, + json={ + "output": [ + { + "type": "message", + "status": "stop", + "content": [ + { + "type": "output_text", + "text": json.dumps({"foo": "bar"}), + "annotations": [], + } + ], + } + ], + }, ) transport = httpx.MockTransport(handler)