"""Harness boundary: the ONLY module that changes when the real harness contract lands. Reads tasks from `/input/tasks.json` (default), with fallback to `--input`/ `--output` CLI args or stdin JSON. Writes `{id, output}` pairs to `/output/results.json` (or `--output`) and logs a run summary (tokens by tier, route distribution) to stderr — never into the results file. All actual solving logic lives in `router.py`; this module only handles I/O, schema tolerance, and summary logging. """ from __future__ import annotations import argparse import json import logging import sys from collections import Counter from pathlib import Path from typing import Any, TextIO from routing_agent.client import FireworksClient, TokenLedger from routing_agent.config import Settings, load_policy from routing_agent.models import Task from routing_agent.registry import KNOWN_MODELS, resolve_allowed from routing_agent.router import route logger = logging.getLogger(__name__) DEFAULT_INPUT_PATH = "/input/tasks.json" DEFAULT_OUTPUT_PATH = "/output/results.json" # Guards against hostile/corrupt task files (SECURITY.md M-1): a runaway input # would otherwise mean unbounded memory and unbounded token spend downstream. MAX_INPUT_BYTES = 50 * 1024 * 1024 MAX_TASKS = 10_000 def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( prog="routing_agent.adapter", description="Run the routing agent over a tasks.json file.", ) parser.add_argument( "--input", default=None, help=f"Path to tasks.json (default: {DEFAULT_INPUT_PATH}, or stdin if absent).", ) parser.add_argument( "--output", default=None, help=f"Path to write results.json (default: {DEFAULT_OUTPUT_PATH}).", ) parser.add_argument( "--policy", default=None, help="Path to a routing policy YAML (default: built-in policy defaults).", ) return parser.parse_args(argv) def _read_tasks(input_path: str | None, stdin: TextIO) -> list[dict[str, Any]]: """Load raw task dicts from a file path, the default input path, or stdin. Resolution order: explicit --input path -> DEFAULT_INPUT_PATH (if it exists) -> stdin. Raises ValueError on malformed JSON or a non-list top-level payload, so failures surface immediately rather than producing a silently empty results.json. """ if input_path is not None: raw_text = Path(input_path).read_text(encoding="utf-8") elif Path(DEFAULT_INPUT_PATH).exists(): raw_text = Path(DEFAULT_INPUT_PATH).read_text(encoding="utf-8") else: raw_text = stdin.read() if not raw_text.strip(): raise ValueError("no task input provided (empty file/stdin)") if len(raw_text.encode("utf-8")) > MAX_INPUT_BYTES: raise ValueError(f"task input exceeds {MAX_INPUT_BYTES} bytes") parsed = json.loads(raw_text) if not isinstance(parsed, list): raise ValueError("tasks.json must contain a top-level JSON array") if len(parsed) > MAX_TASKS: raise ValueError(f"task count {len(parsed)} exceeds limit of {MAX_TASKS}") return parsed def _parse_tasks(raw_tasks: list[dict[str, Any]]) -> list[Task]: tasks: list[Task] = [] for index, raw_task in enumerate(raw_tasks): if "id" not in raw_task: raw_task = {**raw_task, "id": str(index)} tasks.append(Task.model_validate(raw_task)) return tasks def run( tasks: list[Task], settings: Settings, policy_path: str | None, ) -> tuple[list[dict[str, str]], dict[str, Any]]: """Route every task and return (results, summary). `results` is the harness-facing payload: only `id` and `output` per task. `summary` is diagnostic-only (route distribution, token totals by tier) and must never be written into the results file. """ policy = load_policy(policy_path) allowed_models = resolve_allowed(settings.allowed_models) ledger = TokenLedger() client = FireworksClient( api_key=settings.fireworks_api_key, base_url=settings.fireworks_base_url, ledger=ledger, ) results: list[dict[str, str]] = [] tier_counts: Counter[int] = Counter() task_type_counts: Counter[str] = Counter() for task in tasks: outcome = route(task, client, allowed_models, policy) results.append({"id": task.id, "output": outcome.output}) tier_counts[outcome.route.tier] += 1 task_type_counts[outcome.route.task_type] += 1 price_models = {m.id: m for m in allowed_models} | KNOWN_MODELS summary = { "total_tasks": len(tasks), "tier_distribution": dict(tier_counts), "task_type_distribution": dict(task_type_counts), "total_raw_tokens": ledger.total_raw_tokens, "total_price_weighted_usd": round(ledger.total_price_weighted(price_models), 6), "total_calls": len(ledger.records), "retried_calls": sum(1 for r in ledger.records if r.retry), } return results, summary def main(argv: list[str] | None = None) -> int: logging.basicConfig(level=logging.INFO, stream=sys.stderr) args = parse_args(argv) try: raw_tasks = _read_tasks(args.input, sys.stdin) tasks = _parse_tasks(raw_tasks) settings = Settings.from_env() except (ValueError, OSError, json.JSONDecodeError) as exc: logger.error("adapter startup failed: %s", exc) return 1 results, summary = run(tasks, settings, args.policy) output_path = Path(args.output or DEFAULT_OUTPUT_PATH) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text(json.dumps(results, indent=2), encoding="utf-8") logger.info("run summary: %s", json.dumps(summary)) return 0 if __name__ == "__main__": raise SystemExit(main())