diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 000000000..d71a36355 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,36 @@ +# Changelog + +All notable changes to this project are documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +### Fixed +- NIM discovery catalog body bound + dry-run call budget uses max_steps; + offline cost-quality rejects malformed scripted answers and zeros failed-cell usage. + +### Added +- Offline NIM capability probe plan + fixture classification (issue #86). +- Offline cost-quality `--use-mock-orchestrator` path: Fugu `route_once` and + Conductor/TRINITY `conduct` via `mock://` agents (issue #86 paper-path exercise). + +### Added +- Offline NIM cost-quality comparison harness (`nim_cost_quality` + + `nim-cost-quality-offline` CLI) for issue #86 post-discovery: locked task + manifest scorers, honest unknown actual/hypothetical cost, policy summaries, + and quality-latency / quality-cost Pareto frontiers without live egress. +- Offline NIM capability inventory + dry-run benchmark plan (issue #86). +- `discover-nim-models` CLI and `nim_discovery` module (issue #86): allowlisted + NVIDIA HTTPS `/v1/models` only; offline fixture status; unique agent ids on + slug collision; live tests require `RUN_LIVE_NIM_TESTS=1`. +- Role-differentiated sampling temperatures for paper-role ablation. + +### Security +- Semgrep nosemgrep hygiene for audited SQL placeholders / TLS opt-out / urllib. + +## [0.1.0] - 2026-07-13 + +### Added +- Initial OpenAI-compatible orchestration gateway. diff --git a/README.md b/README.md index 65f57dd4c..0afa34044 100644 --- a/README.md +++ b/README.md @@ -285,3 +285,35 @@ python tests/test_commercial_purchase_approval_packet.py python tests/test_commercial_due_diligence_room.py python tests/test_commercial_investment_committee_memo.py ``` + +### NIM model discovery + +```bash +python -m contextual_orchestrator discover-nim-models +python -m contextual_orchestrator discover-nim-models --as-agent-pool +python -m contextual_orchestrator discover-nim-models --benchmark-dry-run +python -m contextual_orchestrator discover-nim-models --capability-probe-plan +python -m contextual_orchestrator discover-nim-models --capability-probe-dry-run examples/nim_capability_probe_fixtures.json +``` + +Requires `NVIDIA_NIM_API_KEY` in the KV (`register-credential`) for live catalog +listing. Offline dry-run / capability inventory paths stay secret-free. + +### Offline NIM cost-quality (issue #86) + +```bash +python -m contextual_orchestrator nim-cost-quality-offline \ + --task-manifest examples/nim_task_manifest_offline.json +python -m contextual_orchestrator nim-cost-quality-offline \ + --task-manifest examples/nim_task_manifest_offline.json \ + --pricing-scenario examples/nim_pricing_scenario_offline.json \ + --markdown +python -m contextual_orchestrator nim-cost-quality-offline \ + --use-mock-orchestrator --agents examples/agents.mock.json --markdown +python tests/test_nim_cost_quality.py +``` + +Compares route/conduct/direct policies on a locked task set with honest +`unknown` costs unless a pricing scenario covers every model. Optional +`--use-mock-orchestrator` exercises Fugu `route_once` and Conductor/TRINITY +`conduct` on `mock://` agents. Never uses `COPILOT_GITHUB_TOKEN` for model calls. diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 5f68c3b74..422375b38 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -55,11 +55,238 @@ def _register_credential_command(argv: list[str]) -> None: print(json.dumps({"registered": args.name, "backend": "kv"}, ensure_ascii=False)) + +def _discover_nim_models_command(argv: list[str]) -> None: + """List NIM model IDs via KV credential and print agent-pool JSON candidates.""" + from .nim_discovery import ( + DEFAULT_NIM_MODELS_URL, + NimDiscoveryError, + build_benchmark_plan_dry_run, + build_capability_inventory, + build_capability_probe_plan, + discover_nim_models, + models_to_agent_pool_entries, + run_capability_probes_dry_run, + validate_nim_models_url, + ) + + parser = argparse.ArgumentParser( + prog="python -m contextual_orchestrator discover-nim-models", + description="Discover NVIDIA NIM model IDs using the KV credential NVIDIA_NIM_API_KEY.", + ) + parser.add_argument( + "--models-url", + default=DEFAULT_NIM_MODELS_URL, + help=( + "HTTPS NVIDIA catalog URL (default: integrate.api.nvidia.com/v1/models). " + "Only allowlisted NVIDIA hosts with path /v1/models are accepted; " + "the API key is never sent to other origins." + ), + ) + parser.add_argument( + "--as-agent-pool", + action="store_true", + help="Emit agent-pool JSON entries instead of the discovery report.", + ) + parser.add_argument( + "--capability-inventory", + action="store_true", + help="Emit offline capability-hint inventory for discovered model ids (issue #86 dry path).", + ) + parser.add_argument( + "--benchmark-dry-run", + action="store_true", + help="Emit a fail-closed dry-run benchmark plan with unknown costs (issue #86).", + ) + parser.add_argument( + "--hard-request-budget", + type=int, + default=100, + help="Hard call budget for dry-run admission (default: 100).", + ) + parser.add_argument( + "--capability-probe-plan", + action="store_true", + help="Emit offline capability probe plan (models x probe kinds) without network.", + ) + parser.add_argument( + "--capability-probe-dry-run", + metavar="FIXTURE_JSON", + default=None, + help=( + "Classify offline probe fixtures from JSON list of " + "{model_id, probe_kind, status_code|error_class, body?} rows." + ), + ) + args = parser.parse_args(argv) + try: + models_url = validate_nim_models_url(args.models_url) + except NimDiscoveryError as exc: + parser.error(str(exc)) + if args.capability_probe_dry_run: + try: + with open(args.capability_probe_dry_run, encoding="utf-8") as handle: + fixtures = json.load(handle) + plan = run_capability_probes_dry_run( + fixtures if isinstance(fixtures, list) else fixtures.get("probe_rows") or fixtures.get("fixtures") or [], + hard_request_budget=args.hard_request_budget, + ) + except (NimDiscoveryError, OSError, ValueError, TypeError) as exc: + parser.error(str(exc)) + print(json.dumps(plan, ensure_ascii=False, indent=2)) + return + + report = discover_nim_models(models_url=models_url) + model_ids = report.get("model_ids") or [] + if args.capability_probe_plan: + try: + plan = build_capability_probe_plan( + model_ids, hard_request_budget=args.hard_request_budget + ) + except NimDiscoveryError as exc: + parser.error(str(exc)) + print(json.dumps(plan, ensure_ascii=False, indent=2)) + return + if args.benchmark_dry_run: + try: + plan = build_benchmark_plan_dry_run( + model_ids, hard_request_budget=args.hard_request_budget + ) + except NimDiscoveryError as exc: + parser.error(str(exc)) + print(json.dumps(plan, ensure_ascii=False, indent=2)) + elif args.capability_inventory: + print(json.dumps(build_capability_inventory(model_ids), ensure_ascii=False, indent=2)) + elif args.as_agent_pool: + print(json.dumps(models_to_agent_pool_entries(model_ids), ensure_ascii=False, indent=2)) + else: + print(json.dumps(report, ensure_ascii=False, indent=2)) + + +def _nim_cost_quality_offline_command(argv: list[str]) -> None: + """Run the offline cost-quality harness against a locked task manifest (issue #86).""" + from .nim_cost_quality import ( + CostQualityContractError, + build_orchestrator_policy_runners, + build_scripted_policy_runners, + load_pricing_scenario, + load_task_manifest, + locked_evaluation_tasks, + render_cost_quality_markdown, + run_offline_cost_quality, + validate_scripted_answers, + ) + from .orchestrator import TaskOrchestrator, load_agents + + parser = argparse.ArgumentParser( + prog="python -m contextual_orchestrator nim-cost-quality-offline", + description=( + "Offline cost-quality comparison for issue #86 (post-discovery). " + "Uses scripted answers by default so CI never needs NVIDIA_NIM_API_KEY. " + "Pass --use-mock-orchestrator to drive Fugu route_once vs Conductor " + "conduct through mock:// agents. Never invents prices." + ), + ) + parser.add_argument( + "--task-manifest", + default="examples/nim_task_manifest_offline.json", + help="Path to the versioned task manifest (locked split only).", + ) + parser.add_argument( + "--pricing-scenario", + default=None, + help="Optional USD-per-million-token scenario JSON; omit to keep costs unknown.", + ) + parser.add_argument( + "--scripted-answers", + default=None, + help=( + "Optional JSON map {task_id: {policy_name: answer}}. " + "When omitted, answers are empty (scores zero) for structural dry-run only." + ), + ) + parser.add_argument( + "--use-mock-orchestrator", + action="store_true", + help=( + "Run policies via TaskOrchestrator route_once/conduct on --agents " + "(default examples/agents.mock.json). Mutually exclusive with " + "--scripted-answers." + ), + ) + parser.add_argument( + "--agents", + default="examples/agents.mock.json", + help="Agent pool JSON for --use-mock-orchestrator (mock:// recommended).", + ) + parser.add_argument( + "--model-id", + default="mock-scripted", + help="Model id recorded on cells and used for pricing lookups (default: mock-scripted).", + ) + parser.add_argument( + "--markdown", + action="store_true", + help="Emit a short markdown summary instead of the full JSON report.", + ) + args = parser.parse_args(argv) + if args.use_mock_orchestrator and args.scripted_answers: + parser.error("--use-mock-orchestrator cannot be combined with --scripted-answers") + try: + manifest = load_task_manifest(args.task_manifest) + tasks = locked_evaluation_tasks(manifest) + pricing = load_pricing_scenario(args.pricing_scenario) + if args.use_mock_orchestrator: + agents = load_agents(args.agents) + if not agents: + parser.error("--agents pool is empty") + non_mock = [a.id for a in agents if not str(a.base_url).startswith("mock://")] + if non_mock: + parser.error( + "--use-mock-orchestrator requires mock:// agents only; " + f"non-mock: {non_mock}" + ) + orchestrator = TaskOrchestrator(agents) + runners = build_orchestrator_policy_runners(orchestrator) + model_id = args.model_id if args.model_id != "mock-scripted" else "mock-orchestrator" + else: + answers: dict = {} + if args.scripted_answers: + with open(args.scripted_answers, encoding="utf-8") as handle: + raw_answers = json.load(handle) + answers = validate_scripted_answers(raw_answers) + runners = build_scripted_policy_runners(answers, model_id=args.model_id) + model_id = args.model_id + report = run_offline_cost_quality( + tasks=tasks, + policy_runners=runners, + model_id=model_id, + pricing_scenario=pricing, + ) + if args.use_mock_orchestrator: + report["runner_backend"] = "mock_orchestrator" + report["agent_pool_path"] = args.agents + else: + report["runner_backend"] = "scripted_answers" + except (CostQualityContractError, OSError, ValueError) as exc: + parser.error(str(exc)) + if args.markdown: + print(render_cost_quality_markdown(report)) + else: + print(json.dumps(report, ensure_ascii=False, indent=2)) + + def main() -> None: """Parse CLI options and run bootstrap, prompt completion, or the HTTP server.""" if len(sys.argv) > 1 and sys.argv[1] == "register-credential": _register_credential_command(sys.argv[2:]) return + if len(sys.argv) > 1 and sys.argv[1] == "discover-nim-models": + _discover_nim_models_command(sys.argv[2:]) + return + if len(sys.argv) > 1 and sys.argv[1] == "nim-cost-quality-offline": + _nim_cost_quality_offline_command(sys.argv[2:]) + return parser = argparse.ArgumentParser(description="Route or conduct chat requests across model agents.") parser.add_argument("prompt", nargs="?", help="User prompt for CLI mode.") diff --git a/contextual_orchestrator/cost_ledger.py b/contextual_orchestrator/cost_ledger.py index d3943c5be..8ca3c6dfd 100644 --- a/contextual_orchestrator/cost_ledger.py +++ b/contextual_orchestrator/cost_ledger.py @@ -583,12 +583,12 @@ def _seed_dimension_catalog(self) -> None: ph = self._placeholder() cur = self._conn.cursor() for order, (name, label, _column) in enumerate(ATTRIBUTION_DIMENSION_CATALOG): - cur.execute( + cur.execute( # nosemgrep -- sqlalchemy-execute-raw-query FP: only the DB-API placeholder char is interpolated; the value is bound. f"SELECT 1 FROM cost_attribution_dimensions WHERE dimension_name = {ph}", # nosec B608 - ph is a DB-API placeholder. (name,), ) if cur.fetchone() is None: - cur.execute( + cur.execute( # nosemgrep -- sqlalchemy-execute-raw-query FP: only DB-API placeholder chars are interpolated; values are bound. "INSERT INTO cost_attribution_dimensions " f"(dimension_name, dimension_label, dimension_order) VALUES ({ph}, {ph}, {ph})", # nosec B608 - ph is a DB-API placeholder. (name, label, order), @@ -602,7 +602,7 @@ def append(self, record: UsageRecord) -> None: placeholders = ", ".join(ph for _ in _USAGE_COLUMNS) columns = ", ".join(_USAGE_COLUMNS) cur = self._conn.cursor() - cur.execute( + cur.execute( # nosemgrep -- sqlalchemy-execute-raw-query FP: columns are the fixed _USAGE_COLUMNS constant; values are bound. f"INSERT INTO llm_usage_records ({columns}) VALUES ({placeholders})", # nosec B608 - columns are fixed _USAGE_COLUMNS. tuple(row.get(column) for column in _USAGE_COLUMNS), ) @@ -622,7 +622,7 @@ def query(self, start: Optional[int] = None, end: Optional[int] = None) -> List[ where = f" WHERE {' AND '.join(clauses)}" if clauses else "" columns = ", ".join(_USAGE_COLUMNS) cur = self._conn.cursor() - cur.execute(f"SELECT {columns} FROM llm_usage_records{where}", tuple(params)) # nosec B608 - columns and clauses are fixed. + cur.execute(f"SELECT {columns} FROM llm_usage_records{where}", tuple(params)) # nosec B608 - columns and clauses are fixed. # nosemgrep -- sqlalchemy-execute-raw-query FP: fixed columns and clause templates; all values are bound. return [dict(zip(_USAGE_COLUMNS, values)) for values in cur.fetchall()] diff --git a/contextual_orchestrator/nim_cost_quality.py b/contextual_orchestrator/nim_cost_quality.py new file mode 100644 index 000000000..090af62aa --- /dev/null +++ b/contextual_orchestrator/nim_cost_quality.py @@ -0,0 +1,646 @@ +"""Offline cost-quality comparison harness for issue #86 (post-discovery). + +Builds on :mod:`nim_discovery` dry-run plans: given a locked task manifest and +mock (or scripted) policy runners, compare Fugu-style single-route, Conductor- +style bounded conduct, and per-worker direct baselines with honest cost fields. + +Optional :func:`build_orchestrator_policy_runners` drives the same comparison +through a live ``TaskOrchestrator`` (typically ``mock://`` agents) so route vs +conduct paper paths are exercised offline without NIM credentials. + +Never invents prices: hypothetical paid cost stays ``\"unknown\"`` until a +versioned pricing scenario covers every model used in a cell. Live NIM egress +requires ``RUN_LIVE_NIM_TESTS=1`` and is out of scope for this offline module. + +References +---------- +Chen, L., Zaharia, M., & Zou, J. (2023). *FrugalGPT: How to use large language +models while reducing cost and improving performance* (arXiv:2305.05176). + +Ong, I., et al. (2024). *RouteLLM: Learning to route LLMs with preference data* +(arXiv:2406.18665). + +Ding, D., et al. (2024). *Hybrid LLM: Cost-efficient and quality-aware query +routing* (arXiv:2404.14618). +""" + +from __future__ import annotations + +import hashlib +import json +import math +import re +import time +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +from .conventions import is_two_word_snake_case +from .nim_discovery import ( + CAPABILITY_CHAT, + CAPABILITY_UNKNOWN, + NimDiscoveryError, + classify_model_capability_hint, +) + +PolicyRunner = Callable[[str], dict[str, Any]] + + +class CostQualityContractError(ValueError): + """Raised when a cost-quality manifest, scenario, or report contract fails.""" + + +def score_exact_number_match(expected: Mapping[str, Any], answer_text: str) -> float: + """Return 1.0 when the expected number appears as a standalone token.""" + pattern = rf"(? float: + """Return 1.0 when the expected substring appears case-insensitively.""" + return 1.0 if str(expected["substring"]).lower() in answer_text.lower() else 0.0 + + +SCORER_REGISTRY: dict[tuple[str, str], Callable[[Mapping[str, Any], str], float]] = { + ("exact_number_match", "1"): score_exact_number_match, + ("substring_match", "1"): score_substring_match, +} + +_VALID_TASK_SPLITS = frozenset({"locked", "exploratory"}) +_POLICY_NAMES = ( + "direct_worker", + "route_once", + "bounded_conduct", + "hindsight_best_single", +) + + +def load_task_manifest(path: str) -> dict[str, Any]: + """Load and validate a versioned task manifest (no prompt leakage).""" + with open(path, encoding="utf-8") as handle: + try: + manifest = json.load(handle) + except ValueError as exc: + raise CostQualityContractError(f"task manifest is not valid JSON: {exc}") from exc + if not isinstance(manifest, dict) or not isinstance(manifest.get("manifest_version"), str): + raise CostQualityContractError( + "task manifest must be an object with a string 'manifest_version'" + ) + tasks = manifest.get("tasks") + if not isinstance(tasks, list) or not tasks: + raise CostQualityContractError("task manifest must carry a non-empty 'tasks' list") + seen_task_ids: set[str] = set() + for task in tasks: + if not isinstance(task, dict): + raise CostQualityContractError("every task manifest entry must be an object") + task_id = task.get("task_id") + if not isinstance(task_id, str) or not is_two_word_snake_case(task_id): + raise CostQualityContractError( + f"task_id must be two-plus-word snake_case: {task_id!r}" + ) + if task_id in seen_task_ids: + raise CostQualityContractError(f"duplicate task_id in manifest: {task_id!r}") + seen_task_ids.add(task_id) + if task.get("split") not in _VALID_TASK_SPLITS: + raise CostQualityContractError( + f"task {task_id!r} split must be 'locked' or 'exploratory'" + ) + prompt = task.get("prompt") + if not isinstance(prompt, str) or not prompt.strip(): + raise CostQualityContractError(f"task {task_id!r} must carry a non-empty prompt") + scorer = task.get("scorer") + if not isinstance(scorer, dict): + raise CostQualityContractError(f"task {task_id!r} must carry a scorer object") + scorer_key = (str(scorer.get("name")), str(scorer.get("version"))) + if scorer_key not in SCORER_REGISTRY: + raise CostQualityContractError( + f"task {task_id!r} names an unregistered scorer: {scorer_key}" + ) + expected = task.get("expected") + if not isinstance(expected, dict) or not expected: + raise CostQualityContractError( + f"task {task_id!r} must carry a non-empty expected object" + ) + if SCORER_REGISTRY[scorer_key](expected, prompt) != 0.0: + raise CostQualityContractError( + f"task {task_id!r} leaks its expected answer into the prompt" + ) + return manifest + + +def locked_evaluation_tasks(manifest: Mapping[str, Any]) -> list[dict[str, Any]]: + """Return only the locked evaluation split, in manifest order.""" + return [task for task in manifest["tasks"] if task["split"] == "locked"] + + +def score_task_answer(task: Mapping[str, Any], answer_text: str) -> dict[str, Any]: + """Score one answer with the task's registered scorer identity.""" + scorer = task["scorer"] + key = (str(scorer["name"]), str(scorer["version"])) + score = float(SCORER_REGISTRY[key](task["expected"], answer_text)) + if not math.isfinite(score) or score < 0.0 or score > 1.0: + raise CostQualityContractError( + f"scorer {key} returned non-finite or out-of-range score for {task['task_id']!r}" + ) + return { + "task_id": task["task_id"], + "scorer_name": key[0], + "scorer_version": key[1], + "score": score, + } + + +def _require_finite_rate(value: Any, label: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0: + raise CostQualityContractError( + f"pricing scenario rate {label} must be a finite non-negative number" + ) + return float(value) + + +def load_pricing_scenario(path: str | None) -> dict[str, Any] | None: + """Load an optional USD-per-million-token scenario; ``None`` keeps costs unknown.""" + if path is None: + return None + with open(path, encoding="utf-8") as handle: + try: + scenario = json.load(handle) + except ValueError as exc: + raise CostQualityContractError(f"pricing scenario is not valid JSON: {exc}") from exc + if not isinstance(scenario, dict) or not isinstance(scenario.get("scenario_version"), str): + raise CostQualityContractError( + "pricing scenario must be an object with a string 'scenario_version'" + ) + if scenario.get("scenario_status") not in ("example_unreviewed", "reviewed"): + raise CostQualityContractError( + "pricing scenario_status must be 'example_unreviewed' or 'reviewed'" + ) + rates = scenario.get("usd_per_million_tokens") + if not isinstance(rates, dict): + raise CostQualityContractError( + "pricing scenario must carry a 'usd_per_million_tokens' object" + ) + for model_id, rate in rates.items(): + if not isinstance(model_id, str) or not model_id.strip(): + raise CostQualityContractError("pricing model id must be a non-empty string") + if not isinstance(rate, dict): + raise CostQualityContractError(f"pricing entry for {model_id!r} must be an object") + _require_finite_rate(rate.get("input"), f"{model_id}.input") + _require_finite_rate(rate.get("output"), f"{model_id}.output") + return scenario + + +def hypothetical_cost_usd( + pricing_scenario: Mapping[str, Any] | None, + usage_by_model: Mapping[str, Mapping[str, int]], +) -> float | str: + """Return scenario cost, or ``\"unknown\"`` when any used model lacks a rate.""" + if pricing_scenario is None: + return "unknown" + rates = pricing_scenario["usd_per_million_tokens"] + total = 0.0 + for model_id, usage in usage_by_model.items(): + rate = rates.get(model_id) + if rate is None: + return "unknown" + prompt_tokens = int(usage.get("prompt_tokens", 0)) + completion_tokens = int(usage.get("completion_tokens", 0)) + if prompt_tokens < 0 or completion_tokens < 0: + raise CostQualityContractError("token counts must be non-negative") + total += prompt_tokens * float(rate["input"]) / 1_000_000 + total += completion_tokens * float(rate["output"]) / 1_000_000 + if not math.isfinite(total): + raise CostQualityContractError("hypothetical cost must be finite") + return round(total, 10) + + +def estimate_token_counts(text: str) -> int: + """Heuristic ~4 chars/token estimate; never claimed as provider-reported usage.""" + if not text: + return 0 + return max(1, (len(text) + 3) // 4) + + +def chat_eligible_model_ids(model_ids: Sequence[str]) -> list[str]: + """Filter catalog ids to chat/unknown capability hints for offline plans.""" + return [ + mid + for mid in sorted({m.strip() for m in model_ids if isinstance(m, str) and m.strip()}) + if classify_model_capability_hint(mid) in {CAPABILITY_CHAT, CAPABILITY_UNKNOWN} + ] + + +def validate_scripted_answers(answers_by_task_id: Any) -> dict[str, dict[str, str]]: + """Validate ``{task_id: {policy_name: answer_text}}`` maps for offline runners.""" + if not isinstance(answers_by_task_id, Mapping): + raise CostQualityContractError("scripted answers must be a JSON object") + normalized: dict[str, dict[str, str]] = {} + for task_id, policy_map in answers_by_task_id.items(): + if not isinstance(task_id, str) or not task_id.strip(): + raise CostQualityContractError("scripted answer task_id must be a non-empty string") + if not isinstance(policy_map, Mapping): + raise CostQualityContractError( + f"scripted answers for {task_id!r} must be an object of policy_name -> answer" + ) + row: dict[str, str] = {} + for policy_name, answer in policy_map.items(): + if not isinstance(policy_name, str) or not policy_name.strip(): + raise CostQualityContractError( + f"scripted answer policy_name under {task_id!r} must be a non-empty string" + ) + if not isinstance(answer, str): + raise CostQualityContractError( + f"scripted answer for {task_id!r}/{policy_name!r} must be a string" + ) + row[policy_name] = answer + normalized[task_id] = row + return normalized + + +def build_scripted_policy_runners( + answers_by_task_id: Mapping[str, Mapping[str, str]] | None = None, + *, + model_id: str = "mock-scripted", +) -> dict[str, PolicyRunner]: + """Build offline policy runners that return scripted answers per task id. + + ``answers_by_task_id[task_id][policy_name]`` supplies the answer text. Missing + cells yield an empty answer (score 0). Used by tests and dry-run demos so CI + never needs ``NVIDIA_NIM_API_KEY`` or network egress. + """ + answers = validate_scripted_answers(answers_by_task_id or {}) + + def _runner(policy_name: str) -> PolicyRunner: + def run(prompt: str) -> dict[str, Any]: + # Prompt carries an embedded task marker when callers use format_task_prompt. + task_id = _extract_task_id_marker(prompt) + answer = "" + if task_id and task_id in answers: + answer = answers[task_id].get(policy_name, "") + return { + "mode": policy_name, + "answer": answer, + "model_id": model_id, + "trace": [{"role": "worker", "agent_id": "scripted_worker", "output": answer}], + "verification": {"accepted": bool(answer)}, + } + + return run + + return {name: _runner(name) for name in _POLICY_NAMES if name != "hindsight_best_single"} + + +def build_orchestrator_policy_runners(orchestrator: Any) -> dict[str, PolicyRunner]: + """Build policy runners backed by a ``TaskOrchestrator`` instance. + + Intended for offline mock pools (``mock://`` agents) and for hermetic CI. + Does not read provider secrets: non-mock agents still resolve keys via KV + ``get_credential`` inside the orchestrator. + + - ``direct_worker`` / ``route_once``: single-worker ``route_once`` path (Fugu). + - ``bounded_conduct``: multi-step ``conduct`` path (Conductor / TRINITY roles). + """ + if orchestrator is None or not hasattr(orchestrator, "route_once") or not hasattr( + orchestrator, "conduct" + ): + raise CostQualityContractError( + "orchestrator must provide route_once and conduct callables" + ) + + def _messages(prompt: str) -> list[dict[str, str]]: + return [{"role": "user", "content": prompt}] + + def direct_worker(prompt: str) -> dict[str, Any]: + result = orchestrator.route_once(_messages(prompt)) + result = dict(result) + result["mode"] = "direct_worker" + return result + + def route_once(prompt: str) -> dict[str, Any]: + result = orchestrator.route_once(_messages(prompt)) + result = dict(result) + result["mode"] = "route_once" + return result + + def bounded_conduct(prompt: str) -> dict[str, Any]: + result = orchestrator.conduct(_messages(prompt)) + result = dict(result) + result["mode"] = "bounded_conduct" + return result + + return { + "direct_worker": direct_worker, + "route_once": route_once, + "bounded_conduct": bounded_conduct, + } + + +def format_task_prompt(task: Mapping[str, Any]) -> str: + """Return the scorable user prompt with a non-scoring task marker line.""" + # Marker is structural only; scorers ignore it and leakage checks use bare prompt. + return f"[task_id={task['task_id']}]\n{task['prompt']}" + + +def _extract_task_id_marker(prompt: str) -> str | None: + match = re.match(r"\[task_id=([a-z0-9_]+)\]\n", prompt) + return match.group(1) if match else None + + +def _usage_from_result(prompt: str, result: Mapping[str, Any], model_id: str) -> dict[str, dict[str, int]]: + answer = str(result.get("answer") or "") + return { + model_id: { + "prompt_tokens": estimate_token_counts(prompt), + "completion_tokens": estimate_token_counts(answer), + } + } + + +def run_policy_cell( + *, + task: Mapping[str, Any], + policy_name: str, + runner: PolicyRunner, + model_id: str, + pricing_scenario: Mapping[str, Any] | None, +) -> dict[str, Any]: + """Execute one policy/task cell and return scored, secret-free evidence.""" + if policy_name not in _POLICY_NAMES: + raise CostQualityContractError(f"unknown policy_name: {policy_name!r}") + prompt = format_task_prompt(task) + started = time.perf_counter() + try: + result = runner(prompt) + outcome = "success" + error_class = None + except Exception as exc: # noqa: BLE001 - classify for evidence only + result = {"mode": policy_name, "answer": "", "trace": [], "verification": {"accepted": False}} + outcome = "failed" + error_class = type(exc).__name__ + latency_ms = round((time.perf_counter() - started) * 1000, 3) + answer = str(result.get("answer") or "") + score_block = score_task_answer(task, answer) + # Failed cells did not complete a provider call — do not invent token/cost usage. + if outcome == "success": + usage = _usage_from_result(prompt, result, model_id) + hyp_cost = hypothetical_cost_usd(pricing_scenario, usage) + call_count = max(1, len(result.get("trace") or [])) + prompt_tokens = usage[model_id]["prompt_tokens"] + completion_tokens = usage[model_id]["completion_tokens"] + usage_source = "estimated" + else: + hyp_cost = "unknown" + call_count = 0 + prompt_tokens = 0 + completion_tokens = 0 + usage_source = "none" + content_hash = hashlib.sha256(answer.encode("utf-8")).hexdigest() if answer else None + return { + "policy_name": policy_name, + "task_id": task["task_id"], + "model_id": model_id, + "outcome": outcome, + "error_class": error_class, + "score": score_block["score"], + "scorer_name": score_block["scorer_name"], + "scorer_version": score_block["scorer_version"], + "latency_ms": latency_ms, + "call_count": call_count, + "workflow_depth": call_count, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + "actual_api_cost": "unknown", + "hypothetical_paid_cost": hyp_cost, + "pricing_scenario_id": ( + pricing_scenario.get("scenario_version") if pricing_scenario is not None else None + ), + "usage_source": usage_source, + "answer_content_hash": content_hash, + "verification_accepted": bool((result.get("verification") or {}).get("accepted")), + } + + +def run_offline_cost_quality( + *, + tasks: Sequence[Mapping[str, Any]], + policy_runners: Mapping[str, PolicyRunner], + model_id: str = "mock-scripted", + pricing_scenario: Mapping[str, Any] | None = None, + include_hindsight_best_single: bool = True, +) -> dict[str, Any]: + """Run fair offline policy comparisons and summarize quality/cost evidence. + + Policies present in ``policy_runners`` are executed for every locked task. + When multiple direct-style runners share the ``direct_worker`` key only one + direct path is used; hindsight best-single reuses per-task direct scores when + ``include_hindsight_best_single`` is true. + """ + if not tasks: + raise CostQualityContractError("tasks must be non-empty") + required = {"direct_worker", "route_once", "bounded_conduct"} + missing = required - set(policy_runners) + if missing: + raise CostQualityContractError(f"policy_runners missing required policies: {sorted(missing)}") + + cells: list[dict[str, Any]] = [] + for task in tasks: + for policy_name in ("direct_worker", "route_once", "bounded_conduct"): + cells.append( + run_policy_cell( + task=task, + policy_name=policy_name, + runner=policy_runners[policy_name], + model_id=model_id, + pricing_scenario=pricing_scenario, + ) + ) + if include_hindsight_best_single: + direct_for_task = [c for c in cells if c["task_id"] == task["task_id"] and c["policy_name"] == "direct_worker"] + best = max(direct_for_task, key=lambda row: row["score"]) if direct_for_task else None + if best is not None: + hindsight = dict(best) + hindsight["policy_name"] = "hindsight_best_single" + cells.append(hindsight) + + summaries = summarize_policy_cells(cells) + frontiers = build_pareto_frontiers(summaries) + return { + "measurement_status": "offline_cost_quality", + "model_id": model_id, + "task_count": len(tasks), + "cell_count": len(cells), + "cells": cells, + "policy_summaries": summaries, + "pareto_frontiers": frontiers, + "pricing_scenario_id": ( + pricing_scenario.get("scenario_version") if pricing_scenario is not None else None + ), + "cost_honesty": ( + "actual_api_cost is unknown offline; hypothetical_paid_cost is unknown " + "unless a pricing scenario prices every model used in the cell" + ), + "quality_proxy": ( + "strict scorer registry on locked tasks; mock/scripted answers only — " + "not a live NIM quality claim" + ), + } + + +def summarize_policy_cells(cells: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]: + """Aggregate mean score/latency and cost honesty per policy.""" + by_policy: dict[str, list[Mapping[str, Any]]] = {} + for cell in cells: + by_policy.setdefault(str(cell["policy_name"]), []).append(cell) + summaries: list[dict[str, Any]] = [] + for policy_name in sorted(by_policy): + rows = by_policy[policy_name] + scores = [float(r["score"]) for r in rows] + latencies = [float(r["latency_ms"]) for r in rows] + hyp_costs = [r["hypothetical_paid_cost"] for r in rows] + numeric_costs = [c for c in hyp_costs if isinstance(c, (int, float))] + summaries.append( + { + "policy_name": policy_name, + "cell_count": len(rows), + "mean_score": round(sum(scores) / len(scores), 6) if scores else 0.0, + "mean_latency_ms": round(sum(latencies) / len(latencies), 3) if latencies else 0.0, + "hypothetical_paid_cost_mean": ( + round(sum(float(c) for c in numeric_costs) / len(numeric_costs), 10) + if numeric_costs and len(numeric_costs) == len(hyp_costs) + else "unknown" + ), + "success_rate": round( + sum(1 for r in rows if r.get("outcome") == "success") / len(rows), 6 + ), + } + ) + return summaries + + +def pareto_frontier( + points: Sequence[Mapping[str, Any]], + *, + quality_key: str = "mean_score", + cost_key: str = "hypothetical_paid_cost_mean", + higher_quality_better: bool = True, + lower_cost_better: bool = True, +) -> list[dict[str, Any]]: + """Return undominated points on a quality-vs-cost frontier (numeric costs only).""" + usable: list[Mapping[str, Any]] = [] + for point in points: + quality = point.get(quality_key) + cost = point.get(cost_key) + if not isinstance(quality, (int, float)) or not math.isfinite(float(quality)): + continue + if not isinstance(cost, (int, float)) or not math.isfinite(float(cost)): + continue + usable.append(point) + frontier: list[dict[str, Any]] = [] + for candidate in usable: + dominated = False + for other in usable: + if other is candidate: + continue + better_or_equal_quality = ( + float(other[quality_key]) >= float(candidate[quality_key]) + if higher_quality_better + else float(other[quality_key]) <= float(candidate[quality_key]) + ) + better_or_equal_cost = ( + float(other[cost_key]) <= float(candidate[cost_key]) + if lower_cost_better + else float(other[cost_key]) >= float(candidate[cost_key]) + ) + strictly_better = ( + float(other[quality_key]) != float(candidate[quality_key]) + or float(other[cost_key]) != float(candidate[cost_key]) + ) + if better_or_equal_quality and better_or_equal_cost and strictly_better: + dominated = True + break + if not dominated: + frontier.append(dict(candidate)) + frontier.sort(key=lambda row: (-float(row[quality_key]), float(row[cost_key]))) + return frontier + + +def build_pareto_frontiers(summaries: Sequence[Mapping[str, Any]]) -> dict[str, Any]: + """Build quality-latency and quality-hypothetical-cost frontiers.""" + quality_latency = pareto_frontier( + [ + { + "policy_name": s["policy_name"], + "mean_score": s["mean_score"], + "hypothetical_paid_cost_mean": s["mean_latency_ms"], + } + for s in summaries + ], + quality_key="mean_score", + cost_key="hypothetical_paid_cost_mean", + ) + # Rename latency key for clarity in the latency frontier payload. + quality_latency = [ + { + "policy_name": row["policy_name"], + "mean_score": row["mean_score"], + "mean_latency_ms": row["hypothetical_paid_cost_mean"], + } + for row in quality_latency + ] + quality_cost = pareto_frontier(list(summaries)) + return { + "quality_latency": quality_latency, + "quality_hypothetical_cost": quality_cost, + "notes": ( + "Points with unknown hypothetical cost are excluded from the cost frontier; " + "latency frontier uses mean_latency_ms as the cost axis." + ), + } + + +def plan_from_discovery_models( + model_ids: Sequence[str], + *, + task_manifest_id: str = "locked_eval_v1", + hard_request_budget: int = 100, +) -> dict[str, Any]: + """Admit a dry-run comparison plan for chat-eligible discovered models. + + Thin wrapper that reuses discovery capability hints so the cost-quality path + stays aligned with ``build_benchmark_plan_dry_run`` admission rules. + """ + from .nim_discovery import build_benchmark_plan_dry_run + + try: + return build_benchmark_plan_dry_run( + list(model_ids), + task_manifest_id=task_manifest_id, + hard_request_budget=hard_request_budget, + ) + except NimDiscoveryError as exc: + raise CostQualityContractError(str(exc)) from exc + + +def render_cost_quality_markdown(report: Mapping[str, Any]) -> str: + """Render a short operator-facing markdown summary (no secrets).""" + lines = [ + "# Offline cost-quality report", + "", + f"- measurement_status: `{report.get('measurement_status')}`", + f"- model_id: `{report.get('model_id')}`", + f"- task_count: {report.get('task_count')}", + f"- cell_count: {report.get('cell_count')}", + f"- pricing_scenario_id: `{report.get('pricing_scenario_id')}`", + "", + "## Policy summaries", + "", + ] + for summary in report.get("policy_summaries") or []: + lines.append( + f"- **{summary['policy_name']}**: mean_score={summary['mean_score']}, " + f"mean_latency_ms={summary['mean_latency_ms']}, " + f"hypothetical_paid_cost_mean={summary['hypothetical_paid_cost_mean']}, " + f"success_rate={summary['success_rate']}" + ) + lines.extend(["", "## Cost honesty", "", str(report.get("cost_honesty") or ""), ""]) + return "\n".join(lines) diff --git a/contextual_orchestrator/nim_discovery.py b/contextual_orchestrator/nim_discovery.py new file mode 100644 index 000000000..1f86d7ddd --- /dev/null +++ b/contextual_orchestrator/nim_discovery.py @@ -0,0 +1,621 @@ +"""Evidence-grade NVIDIA NIM model discovery for agent-pool population. + +Discovers OpenAI-compatible model IDs from a NIM-compatible ``/models`` endpoint +using the KV credential ``NVIDIA_NIM_API_KEY`` (never ``COPILOT_GITHUB_TOKEN``). +Operators convert discovered models into agent pool entries; routing still uses +the deterministic route/conduct policies grounded in Fugu / Conductor / TRINITY +paper contracts. + +References +---------- +Touvron, H., et al. (2023). *Llama 2: Open foundation and fine-tuned chat models* +(arXiv:2307.09288) — open weights commonly hosted on NIM for gateway evaluation. + +Live discovery is optional: tests use offline fixtures so CI stays hermetic. +""" + +from __future__ import annotations + +import json +import ssl +import urllib.error +import urllib.request +from typing import Any +from urllib.parse import urlparse + +from .credentials import get_credential + +DEFAULT_NIM_MODELS_URL = "https://integrate.api.nvidia.com/v1/models" +NIM_CREDENTIAL_NAME = "NVIDIA_NIM_API_KEY" +# Bound catalog body consumption so a pathological allowlisted response cannot OOM the CLI. +NIM_CATALOG_MAX_BYTES = 8 * 1024 * 1024 +# Hosts that may receive the NVIDIA_NIM_API_KEY on authenticated catalog requests. +ALLOWED_NIM_MODELS_HOSTS = frozenset( + { + "integrate.api.nvidia.com", + "api.nvcf.nvidia.com", + } +) + + +class NimDiscoveryError(ValueError): + """Raised when NIM discovery cannot safely proceed (bad URL, etc.).""" + + +def validate_nim_models_url(models_url: str) -> str: + """Return a normalized models catalog URL or raise ``NimDiscoveryError``. + + Authenticated requests only go to allowlisted NVIDIA HTTPS hosts with + path ``/v1/models`` (optional trailing slash). No user-controlled host + may receive ``NVIDIA_NIM_API_KEY``. + """ + if not isinstance(models_url, str) or not models_url.strip(): + raise NimDiscoveryError("models_url must be a non-empty string") + parsed = urlparse(models_url.strip()) + if parsed.scheme != "https": + raise NimDiscoveryError("models_url must use https") + if parsed.username or parsed.password: + raise NimDiscoveryError("models_url must not embed credentials") + if parsed.port not in (None, 443): + raise NimDiscoveryError("models_url must use the default HTTPS port") + hostname = (parsed.hostname or "").lower() + if hostname not in ALLOWED_NIM_MODELS_HOSTS: + raise NimDiscoveryError( + f"models_url host {hostname!r} is not an allowlisted NVIDIA catalog host" + ) + path = parsed.path.rstrip("/") or "" + if path != "/v1/models": + raise NimDiscoveryError("models_url path must be /v1/models") + if parsed.query or parsed.fragment: + raise NimDiscoveryError("models_url must not include query or fragment") + return f"https://{hostname}/v1/models" + + +def discover_nim_models( + *, + models_url: str = DEFAULT_NIM_MODELS_URL, + credential_name: str = NIM_CREDENTIAL_NAME, + timeout_seconds: float = 30.0, + transport: Any | None = None, +) -> dict[str, Any]: + """Discover model IDs from a NIM-compatible OpenAI ``/models`` list endpoint. + + Parameters + ---------- + models_url: + Absolute HTTPS URL of the models listing endpoint. Must pass + :func:`validate_nim_models_url` before any credential is attached. + credential_name: + KV credential name. Defaults to ``NVIDIA_NIM_API_KEY``. + timeout_seconds: + Socket timeout for the listing request. + transport: + Optional callable ``(request, timeout) -> bytes`` for tests. When set, + ``measurement_status`` is ``offline_fixture`` (not live catalog). + + Returns + ------- + dict + ``measurement_status`` (``live_nim_catalog`` | ``offline_fixture`` | + ``credential_missing``), ``model_ids`` (sorted unique strings), and + ``source_url``. Never includes the raw API key. + """ + safe_url = validate_nim_models_url(models_url) + api_key = get_credential(credential_name) + if not api_key: + return { + "measurement_status": "credential_missing", + "model_ids": [], + "source_url": safe_url, + "credential_name": credential_name, + } + + request = urllib.request.Request( # nosemgrep -- dynamic-urllib-use: URL validated by validate_nim_models_url allowlist before auth header is attached. + safe_url, + headers={ + "authorization": f"Bearer {api_key}", + "accept": "application/json", + }, + method="GET", + ) + + if transport is not None: + raw = transport(request, timeout_seconds) + status = "offline_fixture" + else: + context = ssl.create_default_context() + with urllib.request.urlopen( # nosec B310 - URL validated by validate_nim_models_url (HTTPS allowlist). # nosemgrep -- dynamic-urllib-use: URL validated by validate_nim_models_url allowlist before auth header is attached. + request, timeout=timeout_seconds, context=context + ) as response: + content_length = response.headers.get("Content-Length") + if content_length is not None: + try: + declared = int(content_length) + except ValueError as exc: + raise NimDiscoveryError("catalog Content-Length is not an integer") from exc + if declared < 0 or declared > NIM_CATALOG_MAX_BYTES: + raise NimDiscoveryError( + f"catalog Content-Length {declared} exceeds bound {NIM_CATALOG_MAX_BYTES}" + ) + raw = response.read(NIM_CATALOG_MAX_BYTES + 1) + status = "live_nim_catalog" + + if not isinstance(raw, (bytes, bytearray)): + raise NimDiscoveryError("catalog transport must return bytes") + if len(raw) > NIM_CATALOG_MAX_BYTES: + raise NimDiscoveryError( + f"catalog response exceeds bound of {NIM_CATALOG_MAX_BYTES} bytes" + ) + + payload = json.loads(raw.decode("utf-8")) + model_ids = _extract_model_ids(payload) + return { + "measurement_status": status, + "model_ids": model_ids, + "source_url": safe_url, + "model_count": len(model_ids), + } + + +def models_to_agent_pool_entries( + model_ids: list[str], + *, + base_url: str = "https://integrate.api.nvidia.com/v1", + credential_key: str = NIM_CREDENTIAL_NAME, + tags: tuple[str, ...] = ("reasoning", "writing"), +) -> list[dict[str, Any]]: + """Map discovered model IDs to agent-pool JSON dicts (multi-word snake_case ids). + + Each model becomes one agent with a unique deterministic ``id``. Colliding + slugs (normalization or 48-char truncation) get a stable numeric suffix. + After ``model_group`` race lands on main (issue #102 / PR #114), operators + may add ``model_group`` keys for replica race. + """ + entries: list[dict[str, Any]] = [] + used_ids: set[str] = set() + for index, model_id in enumerate(model_ids): + agent_id = _unique_agent_id(model_id, used_ids) + used_ids.add(agent_id) + entries.append( + { + "id": agent_id, + "model": model_id, + "base_url": base_url, + "credential_key": credential_key, + "tags": list(tags), + "priority": max(0, 10 - index), + } + ) + return entries + + +def _unique_agent_id(model_id: str, used_ids: set[str]) -> str: + """Build ``nim__agent`` and append ``_N`` when the id already exists.""" + slug = _slug_model_id(model_id) + base = f"nim_{slug}_agent" + if base not in used_ids: + return base + suffix = 2 + while True: + candidate = f"{base}_{suffix}" + if candidate not in used_ids: + return candidate + suffix += 1 + + +def _extract_model_ids(payload: Any) -> list[str]: + """Parse OpenAI-style ``{data: [{id: ...}]}`` or a bare list of ids/objects.""" + ids: list[str] = [] + if isinstance(payload, dict): + data = payload.get("data", payload.get("models", [])) + else: + data = payload + if not isinstance(data, list): + return [] + for item in data: + if isinstance(item, str) and item.strip(): + ids.append(item.strip()) + elif isinstance(item, dict): + mid = item.get("id") or item.get("model") + if isinstance(mid, str) and mid.strip(): + ids.append(mid.strip()) + return sorted(set(ids)) + + +def _slug_model_id(model_id: str) -> str: + """Convert a provider model id into a multi-word-friendly snake_case token.""" + cleaned = [] + for char in model_id.lower(): + if char.isalnum(): + cleaned.append(char) + else: + cleaned.append("_") + slug = "".join(cleaned).strip("_") + while "__" in slug: + slug = slug.replace("__", "_") + if not slug: + slug = "unnamed_model" + # require_object_name needs two semantic words — ensure underscore present + if "_" not in slug: + slug = f"{slug}_model" + return slug[:48] + + +# Capability labels used in offline inventory / dry-run benchmark plans (issue #86). +CAPABILITY_CHAT = "chat" +CAPABILITY_EMBEDDINGS = "embeddings" +CAPABILITY_IMAGE = "image" +CAPABILITY_AUDIO = "audio" +CAPABILITY_VIDEO = "video" +CAPABILITY_UNSUPPORTED = "unsupported" +CAPABILITY_UNKNOWN = "unknown" + +_BENCHMARK_POLICY_NAMES = ( + "direct_worker", + "route_once", + "bounded_conduct", + "hindsight_best_single", +) + + +def classify_model_capability_hint(model_id: str) -> str: + """Return a coarse capability label from the model id string alone (offline). + + This is a catalog hint for dry-run inventory, not a live probe. Live + capability probing (issue #86) must opt in with ``RUN_LIVE_NIM_TESTS=1`` and + never invent success for unsupported modalities. + """ + if not isinstance(model_id, str) or not model_id.strip(): + return CAPABILITY_UNSUPPORTED + token = model_id.lower() + if any(part in token for part in ("embed", "embedding", "e5-", "bge-")): + return CAPABILITY_EMBEDDINGS + if any(part in token for part in ("image", "vision", "sdxl", "flux", "dall-e", "stable-diffusion")): + return CAPABILITY_IMAGE + if any(part in token for part in ("audio", "whisper", "tts", "speech", "asr")): + return CAPABILITY_AUDIO + if any(part in token for part in ("video", "luma", "runway", "sora")): + return CAPABILITY_VIDEO + if any(part in token for part in ("gpt", "llama", "gemma", "mistral", "claude", "qwen", "nemotron", "instruct", "chat")): + return CAPABILITY_CHAT + return CAPABILITY_UNKNOWN + + +def build_capability_inventory(model_ids: list[str]) -> dict[str, Any]: + """Build a secret-free capability inventory from discovered model ids. + + Offline-only: classifies each id via :func:`classify_model_capability_hint`. + ``measurement_status`` is always ``offline_capability_hints`` so operators + never confuse this with a live probe. + """ + rows: list[dict[str, str]] = [] + for model_id in sorted({mid.strip() for mid in model_ids if isinstance(mid, str) and mid.strip()}): + rows.append( + { + "model_id": model_id, + "capability_hint": classify_model_capability_hint(model_id), + } + ) + by_capability: dict[str, int] = {} + for row in rows: + label = row["capability_hint"] + by_capability[label] = by_capability.get(label, 0) + 1 + return { + "measurement_status": "offline_capability_hints", + "model_count": len(rows), + "capability_rows": rows, + "capability_counts": dict(sorted(by_capability.items())), + } + + +def build_benchmark_plan_dry_run( + model_ids: list[str], + *, + task_manifest_id: str = "locked_eval_v1", + max_steps: int = 5, + hard_request_budget: int = 100, +) -> dict[str, Any]: + """Return a fail-closed dry-run benchmark plan for issue #86 fair comparisons. + + Never attaches secrets. Hypothetical cost fields stay ``unknown`` until a + versioned pricing scenario exists (honest cost reporting — never invent zero). + Policies cover direct worker, route_once, bounded conduct, and hindsight + best-single baselines per the product research contract. + """ + if max_steps < 1 or max_steps > 5: + raise NimDiscoveryError("max_steps must be between 1 and 5 for Conductor/TRINITY-bounded dry runs") + if hard_request_budget < 1: + raise NimDiscoveryError("hard_request_budget must be >= 1") + if not isinstance(task_manifest_id, str) or not task_manifest_id.strip(): + raise NimDiscoveryError("task_manifest_id must be a non-empty string") + + chat_eligible = [ + mid + for mid in sorted({m.strip() for m in model_ids if isinstance(m, str) and m.strip()}) + if classify_model_capability_hint(mid) in {CAPABILITY_CHAT, CAPABILITY_UNKNOWN} + ] + inventory = build_capability_inventory(model_ids) + cells: list[dict[str, Any]] = [] + # Worst-case provider calls per chat-eligible model: + # direct_worker=1, route_once=1, bounded_conduct<=max_steps. + # hindsight_best_single reuses direct scores (zero extra egress). + per_model_call_budget = 1 + 1 + max_steps + planned_calls = len(chat_eligible) * per_model_call_budget + for policy_name in _BENCHMARK_POLICY_NAMES: + for model_id in chat_eligible: + cells.append( + { + "policy_name": policy_name, + "model_id": model_id, + "task_manifest_id": task_manifest_id.strip(), + "max_steps": max_steps, + "planned_provider_calls": ( + max_steps if policy_name == "bounded_conduct" else ( + 0 if policy_name == "hindsight_best_single" else 1 + ) + ), + "actual_api_cost": "unknown", + "hypothetical_paid_cost": "unknown", + "pricing_scenario_id": None, + } + ) + + fits_budget = planned_calls <= hard_request_budget + return { + "measurement_status": "dry_run_plan", + "task_manifest_id": task_manifest_id.strip(), + "max_steps": max_steps, + "hard_request_budget": hard_request_budget, + "planned_request_count": planned_calls, + "per_model_call_budget": per_model_call_budget, + "fits_hard_request_budget": fits_budget, + "chat_eligible_model_count": len(chat_eligible), + "capability_inventory": inventory, + "comparison_cells": cells, + "admission_status": "admitted" if fits_budget else "rejected_budget_exceeded", + } + + +# Probe outcome labels for issue #86 capability inventory (offline dry-run + live opt-in). +PROBE_OUTCOME_CHAT = "chat" +PROBE_OUTCOME_EMBEDDINGS = "embeddings" +PROBE_OUTCOME_IMAGE = "image" +PROBE_OUTCOME_AUDIO = "audio" +PROBE_OUTCOME_VIDEO = "video" +PROBE_OUTCOME_UNSUPPORTED = "unsupported" +PROBE_OUTCOME_RATE_LIMITED = "rate_limited" +PROBE_OUTCOME_UNAVAILABLE = "unavailable" +PROBE_OUTCOME_TIMEOUT = "timeout" +PROBE_OUTCOME_MALFORMED = "malformed" +PROBE_OUTCOME_FAILED = "failed" +PROBE_OUTCOME_SKIPPED = "skipped" + +_PROBE_OUTCOMES = frozenset( + { + PROBE_OUTCOME_CHAT, + PROBE_OUTCOME_EMBEDDINGS, + PROBE_OUTCOME_IMAGE, + PROBE_OUTCOME_AUDIO, + PROBE_OUTCOME_VIDEO, + PROBE_OUTCOME_UNSUPPORTED, + PROBE_OUTCOME_RATE_LIMITED, + PROBE_OUTCOME_UNAVAILABLE, + PROBE_OUTCOME_TIMEOUT, + PROBE_OUTCOME_MALFORMED, + PROBE_OUTCOME_FAILED, + PROBE_OUTCOME_SKIPPED, + } +) + + +def classify_probe_http_status(status_code: int) -> str: + """Map an HTTP status from a capability probe to a machine-readable outcome.""" + if not isinstance(status_code, int) or isinstance(status_code, bool): + raise NimDiscoveryError("probe status_code must be an int") + if status_code == 200: + return PROBE_OUTCOME_CHAT # refined by response body shape in classify_probe_result + if status_code == 429: + return PROBE_OUTCOME_RATE_LIMITED + if status_code in (401, 403): + return PROBE_OUTCOME_UNAVAILABLE + if status_code == 404: + return PROBE_OUTCOME_UNSUPPORTED + if status_code == 408 or status_code == 504: + return PROBE_OUTCOME_TIMEOUT + if 400 <= status_code < 500: + return PROBE_OUTCOME_UNSUPPORTED + if 500 <= status_code < 600: + return PROBE_OUTCOME_FAILED + raise NimDiscoveryError(f"unsupported probe status_code: {status_code}") + + +def classify_probe_result( + *, + model_id: str, + probe_kind: str, + status_code: int, + body: Any | None = None, + error_class: str | None = None, +) -> dict[str, Any]: + """Classify one capability probe into a secret-free evidence row. + + Offline dry-run supplies fixture status/body; live probes (opt-in) reuse the + same classifier. Never embeds credentials or raw provider secrets. + """ + if not isinstance(model_id, str) or not model_id.strip(): + raise NimDiscoveryError("model_id must be a non-empty string") + if not isinstance(probe_kind, str) or not probe_kind.strip(): + raise NimDiscoveryError("probe_kind must be a non-empty string") + kind = probe_kind.strip().lower() + if error_class: + err = str(error_class) + if "timeout" in err.lower() or err in {"TimeoutError", "socket.timeout"}: + outcome = PROBE_OUTCOME_TIMEOUT + else: + outcome = PROBE_OUTCOME_FAILED + return { + "model_id": model_id.strip(), + "probe_kind": kind, + "status_code": status_code if isinstance(status_code, int) else None, + "outcome": outcome, + "error_class": err, + "skip_reason": None, + } + + status_outcome = classify_probe_http_status(status_code) + if status_outcome != PROBE_OUTCOME_CHAT: + return { + "model_id": model_id.strip(), + "probe_kind": kind, + "status_code": status_code, + "outcome": status_outcome, + "error_class": None, + "skip_reason": None, + } + + # 200: refine by body shape when present. + if body is None: + outcome = { + "chat": PROBE_OUTCOME_CHAT, + "embeddings": PROBE_OUTCOME_EMBEDDINGS, + "image": PROBE_OUTCOME_IMAGE, + "audio": PROBE_OUTCOME_AUDIO, + "video": PROBE_OUTCOME_VIDEO, + }.get(kind, PROBE_OUTCOME_CHAT) + elif not isinstance(body, dict): + outcome = PROBE_OUTCOME_MALFORMED + elif kind == "embeddings" and isinstance(body.get("data"), list): + outcome = PROBE_OUTCOME_EMBEDDINGS + elif kind == "chat" and isinstance(body.get("choices"), list) and body.get("choices"): + outcome = PROBE_OUTCOME_CHAT + elif kind in {"image", "audio", "video"} and body: + outcome = { + "image": PROBE_OUTCOME_IMAGE, + "audio": PROBE_OUTCOME_AUDIO, + "video": PROBE_OUTCOME_VIDEO, + }[kind] + else: + outcome = PROBE_OUTCOME_MALFORMED + + return { + "model_id": model_id.strip(), + "probe_kind": kind, + "status_code": status_code, + "outcome": outcome, + "error_class": None, + "skip_reason": None, + } + + +def build_capability_probe_plan( + model_ids: list[str], + *, + hard_request_budget: int = 100, + probe_kinds: tuple[str, ...] = ("chat", "embeddings"), +) -> dict[str, Any]: + """Build a fail-closed dry-run probe plan for discovered model ids. + + Does not perform network I/O. Planned probe count = models x probe_kinds. + """ + if hard_request_budget < 1: + raise NimDiscoveryError("hard_request_budget must be >= 1") + if not probe_kinds: + raise NimDiscoveryError("probe_kinds must be non-empty") + for kind in probe_kinds: + if not isinstance(kind, str) or not kind.strip(): + raise NimDiscoveryError("probe_kinds entries must be non-empty strings") + unique_ids = sorted({m.strip() for m in model_ids if isinstance(m, str) and m.strip()}) + kinds = tuple(k.strip().lower() for k in probe_kinds) + planned = len(unique_ids) * len(kinds) + cells = [ + { + "model_id": mid, + "probe_kind": kind, + "planned_provider_calls": 1, + "hint": classify_model_capability_hint(mid), + } + for mid in unique_ids + for kind in kinds + ] + fits = planned <= hard_request_budget + return { + "measurement_status": "offline_probe_plan", + "hard_request_budget": hard_request_budget, + "planned_request_count": planned, + "fits_hard_request_budget": fits, + "admission_status": "admitted" if fits else "rejected_budget_exceeded", + "model_count": len(unique_ids), + "probe_kinds": list(kinds), + "probe_cells": cells, + } + + +def run_capability_probes_dry_run( + fixture_rows: list[dict[str, Any]], + *, + hard_request_budget: int = 100, +) -> dict[str, Any]: + """Execute offline capability probes from fixture rows (no network, no secrets). + + Each fixture row must include ``model_id``, ``probe_kind``, and either + ``status_code`` or ``error_class``. Optional ``body`` refines 200 outcomes. + """ + if hard_request_budget < 1: + raise NimDiscoveryError("hard_request_budget must be >= 1") + if not isinstance(fixture_rows, list): + raise NimDiscoveryError("fixture_rows must be a list") + if len(fixture_rows) > hard_request_budget: + raise NimDiscoveryError( + f"fixture probe count {len(fixture_rows)} exceeds hard_request_budget {hard_request_budget}" + ) + + results: list[dict[str, Any]] = [] + for index, row in enumerate(fixture_rows): + if not isinstance(row, dict): + raise NimDiscoveryError(f"fixture row {index} must be an object") + model_id = row.get("model_id") + probe_kind = row.get("probe_kind") + if row.get("error_class"): + classified = classify_probe_result( + model_id=str(model_id or ""), + probe_kind=str(probe_kind or "chat"), + status_code=0, + error_class=str(row["error_class"]), + ) + else: + status = row.get("status_code") + if not isinstance(status, int) or isinstance(status, bool): + raise NimDiscoveryError(f"fixture row {index} needs int status_code or error_class") + classified = classify_probe_result( + model_id=str(model_id or ""), + probe_kind=str(probe_kind or "chat"), + status_code=status, + body=row.get("body"), + ) + results.append(classified) + + by_outcome: dict[str, int] = {} + for row in results: + by_outcome[row["outcome"]] = by_outcome.get(row["outcome"], 0) + 1 + + chat_eligible = sorted( + { + r["model_id"] + for r in results + if r["outcome"] == PROBE_OUTCOME_CHAT and r["probe_kind"] == "chat" + } + ) + return { + "measurement_status": "offline_probe_results", + "probe_count": len(results), + "outcome_counts": dict(sorted(by_outcome.items())), + "chat_eligible_model_ids": chat_eligible, + "probe_rows": results, + "hard_request_budget": hard_request_budget, + "fits_hard_request_budget": True, + "admission_status": "admitted", + "notes": ( + "Offline fixture classification only — not live NIM probe evidence. " + "Live probes require RUN_LIVE_NIM_TESTS=1 and NVIDIA_NIM_API_KEY via KV." + ), + } diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 0097b722e..29e36c358 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -169,6 +169,9 @@ class OrchestrationPolicy: # report; "model" asks a verifier-selected model to reply ACCEPT/REJECT (fixes the # known term-matching false negative on risk-vocabulary verifier outputs). verifier_judge: str = "terms" + # Role-differentiated sampling temperature (reasoning effort proxy for ablation). + # Thinker/verifier lower for stability; worker slightly higher for exploration. + role_temperature: dict[str, float] | None = None def as_dict(self) -> dict[str, Any]: """Return the API-safe policy snapshot for workflow records.""" @@ -179,10 +182,29 @@ def as_dict(self) -> dict[str, Any]: "workflow_planning": self.workflow_planning, "verifier_judge": self.verifier_judge, "max_workflow_steps": self.max_workflow_steps, + "role_temperature": dict(self.role_temperature or self.default_role_temperature()), "workflow_steps": ["thinker", "worker", "verifier", "synthesizer"], "supported_locales": ["en", "ko"], } + @staticmethod + def default_role_temperature() -> dict[str, float]: + """Default per-role temperatures for route/conduct ablation studies.""" + return { + "thinker": 0.1, + "worker": 0.2, + "verifier": 0.0, + "synthesizer": 0.15, + } + + def temperature_for_role(self, role: str) -> float: + """Return the sampling temperature configured for a paper role.""" + table = self.role_temperature or self.default_role_temperature() + try: + return float(table.get(role, 0.2)) + except (TypeError, ValueError): + return 0.2 + # HTTP statuses worth retrying: request timeout, conflict, too-early, rate limit, # and the standard upstream/gateway failures. Everything else (400/401/403/404 ...) @@ -230,7 +252,7 @@ def __init__( @staticmethod def _build_ssl_context(ca_bundle: str | None, verify_tls: bool) -> ssl.SSLContext: if not verify_tls: - return ssl._create_unverified_context() # nosec B323 - explicit dev-only provider TLS opt-out. + return ssl._create_unverified_context() # nosec B323 - explicit dev-only provider TLS opt-out. # nosemgrep -- unverified-ssl-context: intentional, default-secure (verify_tls defaults True) dev-only opt-out for self-signed endpoints. if ca_bundle: if not os.path.isfile(ca_bundle): raise ValueError(f"provider CA bundle does not exist: {ca_bundle}") @@ -307,7 +329,7 @@ def _send(self, agent: ModelAgent, payload: dict[str, Any]) -> str: def _open_provider(self, request: urllib.request.Request) -> Any: """Open a provider request built from a validated provider URL.""" - return urllib.request.urlopen( # nosec B310 - request URL comes from _provider_url after provider validation. + return urllib.request.urlopen( # nosec B310 - request URL comes from _provider_url after provider validation. # nosemgrep -- dynamic-urllib-use: URL is built by _provider_url after scheme/host validation; egress to loopback/private/reserved is blocked. request, timeout=self.timeout, context=self._ssl_context, @@ -1548,10 +1570,11 @@ def _invoke( usage when available (else None), so spend analytics can prefer it. """ candidates = self._failover_candidates(primary, text, role) + temperature = self.policy.temperature_for_role(role) last_error: Exception | None = None for agent in candidates: try: - output = self.client.chat(agent, messages) + output = self.client.chat(agent, messages, temperature=temperature) except Exception as exc: # noqa: BLE001 - one agent failing routes to the next last_error = exc self._record_failure(agent.id) diff --git a/docs/architecture.md b/docs/architecture.md index c0f63a81e..cf8e859f8 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -53,3 +53,57 @@ The product is not a Fugu clone. It is a control-plane prototype for the same pu - replayable evaluation runs before any learned coordinator replaces the deterministic policy. See [product_planning.md](product_planning.md) for the product reboot. + +## NVIDIA NIM discovery + +`contextual_orchestrator.nim_discovery` lists models from a NIM-compatible +OpenAI `/models` endpoint using the KV credential `NVIDIA_NIM_API_KEY` (bootstrap +may seed the KV from process env; request-time resolution stays on `get_credential`). +Authenticated catalog requests only target allowlisted NVIDIA HTTPS hosts with +path `/v1/models` (`validate_nim_models_url`); the API key is never sent to a +caller-controlled origin. Discovered IDs convert to agent-pool entries for +cost-aware routing and optional `model_group` race once that lands. Default CI +uses offline fixtures (`measurement_status=offline_fixture` when a test transport +is injected). Live catalog checks require explicit opt-in +`RUN_LIVE_NIM_TESTS=1` plus a seeded `NVIDIA_NIM_API_KEY`. + +Offline capability inventory (`build_capability_inventory`) and dry-run +benchmark plans (`build_benchmark_plan_dry_run`) advance issue #86 without +network or secrets: cost fields stay `unknown` until a versioned pricing +scenario exists, and plans that exceed `--hard-request-budget` are rejected +fail-closed (`admission_status=rejected_budget_exceeded`). + +## Offline NIM cost-quality (issue #86 post-discovery) + +`contextual_orchestrator.nim_cost_quality` runs fair offline policy comparisons +on a locked task manifest after discovery admission: + +- policies: `direct_worker`, `route_once`, `bounded_conduct`, and + `hindsight_best_single` (max direct score per task); +- strict scorers (`exact_number_match` / `substring_match`) with no prompt + leakage of expected answers; +- `actual_api_cost` always `unknown` offline; `hypothetical_paid_cost` is + numeric only when a versioned pricing scenario prices every model used in the + cell — partial tables stay `unknown` (never coerced to zero); +- quality-latency and quality-hypothetical-cost Pareto frontiers from policy + summaries. + +CLI: `python -m contextual_orchestrator nim-cost-quality-offline`. Scripted +answers keep CI hermetic; `--use-mock-orchestrator` drives Fugu `route_once` and +Conductor/TRINITY `conduct` through `mock://` agents without provider secrets. +Live NIM quality/cost still requires explicit opt-in and is not claimed by this +offline module. + +## Role temperature (reasoning effort proxy) + +`OrchestrationPolicy.role_temperature` sets per-role sampling temperatures +(default: thinker 0.1, worker 0.2, verifier 0.0, synthesizer 0.15) so ablation +studies can vary reasoning effort by paper role without collapsing multi-agent +depth. Used by `_invoke` for route and conduct steps. + +## Offline NIM capability probes (issue #86) + +`build_capability_probe_plan` and `run_capability_probes_dry_run` classify +fixture probe rows (HTTP status + optional body shape) into chat / embeddings / +rate_limited / unsupported / timeout / failed outcomes without network or +secrets. Live probing remains opt-in via `RUN_LIVE_NIM_TESTS=1`. diff --git a/docs/papers/README.md b/docs/papers/README.md index 65a89d2af..591a57a23 100644 --- a/docs/papers/README.md +++ b/docs/papers/README.md @@ -43,6 +43,14 @@ motivate throughput-oriented **batched** inference and the load-balancing that makes the latency-tolerant batch route economical. Those sources are referenced but not vendored here so this repository remains one deployable control plane. +## Offline cost-quality comparison (issue #86) + +The `nim_cost_quality` offline harness reuses the FrugalGPT / RouteLLM / Hybrid +LLM citations above for quality-vs-cost Pareto reporting after NIM discovery. +It does not claim live NIM quality until an opted-in secret-backed run publishes +exact-head evidence; hypothetical prices remain `unknown` without a versioned +scenario. + > Citations are provided for scholarly attribution. Redistribution here relies > on the arXiv non-exclusive distribution license each author granted; no > GPL/AGPL-licensed material is vendored anywhere in this repository. diff --git a/examples/nim_capability_probe_fixtures.json b/examples/nim_capability_probe_fixtures.json new file mode 100644 index 000000000..aaaf37fb4 --- /dev/null +++ b/examples/nim_capability_probe_fixtures.json @@ -0,0 +1,29 @@ +[ + { + "model_id": "meta/llama-3-8b-instruct", + "probe_kind": "chat", + "status_code": 200, + "body": {"choices": [{"message": {"content": "ok"}}]} + }, + { + "model_id": "nvidia/nv-embedqa-e5-v5", + "probe_kind": "embeddings", + "status_code": 200, + "body": {"data": [{"embedding": [0.1, 0.2]}]} + }, + { + "model_id": "vendor/missing-model", + "probe_kind": "chat", + "status_code": 404 + }, + { + "model_id": "vendor/rate-limited", + "probe_kind": "chat", + "status_code": 429 + }, + { + "model_id": "vendor/timeout-model", + "probe_kind": "chat", + "error_class": "TimeoutError" + } +] diff --git a/examples/nim_pricing_scenario_offline.json b/examples/nim_pricing_scenario_offline.json new file mode 100644 index 000000000..2953cde87 --- /dev/null +++ b/examples/nim_pricing_scenario_offline.json @@ -0,0 +1,9 @@ +{ + "scenario_version": "2026-08-13.1", + "scenario_status": "example_unreviewed", + "scenario_notes": "Schema-demonstration hypothetical USD-per-million-token rates for offline dry runs only. Not authoritative NVIDIA pricing. Hosted NIM may be free to the caller (actual_api_cost stays separate). Models absent from usd_per_million_tokens keep hypothetical_paid_cost as unknown.", + "usd_per_million_tokens": { + "mock-scripted": {"input": 0.2, "output": 0.6}, + "dryrun/chat-basic": {"input": 0.2, "output": 0.6} + } +} diff --git a/examples/nim_task_manifest_offline.json b/examples/nim_task_manifest_offline.json new file mode 100644 index 000000000..c510ba3e0 --- /dev/null +++ b/examples/nim_task_manifest_offline.json @@ -0,0 +1,34 @@ +{ + "manifest_version": "2026-08-13.1", + "manifest_notes": "Minimal locked offline split for issue #86 cost-quality dry path. Expected answers never appear in prompts. Live NIM evaluation may use a larger locked set under RUN_LIVE_NIM_TESTS=1.", + "tasks": [ + { + "task_id": "digit_sum_reasoning", + "split": "locked", + "prompt": "Multiply 12 by 12, then add 56 to the result. Answer with the number only.", + "scorer": {"name": "exact_number_match", "version": "1"}, + "expected": {"number": "200"} + }, + { + "task_id": "linear_equation_solution", + "split": "locked", + "prompt": "Solve 3x + 5 = 26 for x. Answer with the number only.", + "scorer": {"name": "exact_number_match", "version": "1"}, + "expected": {"number": "7"} + }, + { + "task_id": "capital_recall_france", + "split": "locked", + "prompt": "Name the capital city of France. Answer with the city name only.", + "scorer": {"name": "substring_match", "version": "1"}, + "expected": {"substring": "Paris"} + }, + { + "task_id": "sequence_next_fibonacci", + "split": "locked", + "prompt": "What number comes next in this sequence: 1, 1, 2, 3, 5, 8, 13? Answer with the number only.", + "scorer": {"name": "exact_number_match", "version": "1"}, + "expected": {"number": "21"} + } + ] +} diff --git a/tests/test_nim_cost_quality.py b/tests/test_nim_cost_quality.py new file mode 100644 index 000000000..26c0d8d67 --- /dev/null +++ b/tests/test_nim_cost_quality.py @@ -0,0 +1,399 @@ +"""Contracts for offline NIM cost-quality comparisons (issue #86 post-discovery).""" + +from __future__ import annotations + +import json +import os +import tempfile +import unittest +from pathlib import Path + +from contextual_orchestrator.nim_cost_quality import ( + CostQualityContractError, + build_orchestrator_policy_runners, + build_pareto_frontiers, + build_scripted_policy_runners, + chat_eligible_model_ids, + format_task_prompt, + hypothetical_cost_usd, + load_pricing_scenario, + load_task_manifest, + locked_evaluation_tasks, + plan_from_discovery_models, + render_cost_quality_markdown, + run_offline_cost_quality, + run_policy_cell, + score_exact_number_match, + score_substring_match, + score_task_answer, + summarize_policy_cells, +) +from contextual_orchestrator.orchestrator import ModelAgent, TaskOrchestrator + +ROOT = Path(__file__).resolve().parents[1] +MANIFEST_PATH = ROOT / "examples" / "nim_task_manifest_offline.json" +PRICING_PATH = ROOT / "examples" / "nim_pricing_scenario_offline.json" + + +class TestScorersAndManifest(unittest.TestCase): + def test_exact_number_standalone_only(self) -> None: + expected = {"number": "21"} + self.assertEqual(score_exact_number_match(expected, "answer is 21."), 1.0) + self.assertEqual(score_exact_number_match(expected, "210"), 0.0) + self.assertEqual(score_exact_number_match(expected, "121"), 0.0) + + def test_substring_case_insensitive(self) -> None: + self.assertEqual(score_substring_match({"substring": "Paris"}, "paris is capital"), 1.0) + self.assertEqual(score_substring_match({"substring": "Paris"}, "Lyon"), 0.0) + + def test_load_manifest_and_locked_split(self) -> None: + manifest = load_task_manifest(str(MANIFEST_PATH)) + locked = locked_evaluation_tasks(manifest) + self.assertGreaterEqual(len(locked), 3) + self.assertTrue(all(t["split"] == "locked" for t in locked)) + + def test_manifest_rejects_leakage(self) -> None: + payload = { + "manifest_version": "x", + "tasks": [ + { + "task_id": "leaky_task_case", + "split": "locked", + "prompt": "The answer is 7 only.", + "scorer": {"name": "exact_number_match", "version": "1"}, + "expected": {"number": "7"}, + } + ], + } + with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as handle: + json.dump(payload, handle) + path = handle.name + try: + with self.assertRaises(CostQualityContractError): + load_task_manifest(path) + finally: + os.unlink(path) + + def test_manifest_rejects_duplicate_task_id(self) -> None: + payload = { + "manifest_version": "x", + "tasks": [ + { + "task_id": "same_task_one", + "split": "locked", + "prompt": "Name a city.", + "scorer": {"name": "substring_match", "version": "1"}, + "expected": {"substring": "Oslo"}, + }, + { + "task_id": "same_task_one", + "split": "locked", + "prompt": "Name another city.", + "scorer": {"name": "substring_match", "version": "1"}, + "expected": {"substring": "Bergen"}, + }, + ], + } + with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as handle: + json.dump(payload, handle) + path = handle.name + try: + with self.assertRaises(CostQualityContractError): + load_task_manifest(path) + finally: + os.unlink(path) + + +class TestPricingHonesty(unittest.TestCase): + def test_missing_scenario_is_unknown(self) -> None: + self.assertEqual( + hypothetical_cost_usd(None, {"m": {"prompt_tokens": 10, "completion_tokens": 5}}), + "unknown", + ) + + def test_partial_price_table_is_unknown(self) -> None: + scenario = load_pricing_scenario(str(PRICING_PATH)) + self.assertIsNotNone(scenario) + self.assertEqual( + hypothetical_cost_usd( + scenario, + {"unpriced-model": {"prompt_tokens": 10, "completion_tokens": 5}}, + ), + "unknown", + ) + + def test_priced_model_returns_finite_cost(self) -> None: + scenario = load_pricing_scenario(str(PRICING_PATH)) + cost = hypothetical_cost_usd( + scenario, + {"mock-scripted": {"prompt_tokens": 1_000_000, "completion_tokens": 1_000_000}}, + ) + self.assertIsInstance(cost, float) + self.assertAlmostEqual(float(cost), 0.8, places=6) + + def test_rejects_non_finite_rate(self) -> None: + payload = { + "scenario_version": "bad", + "scenario_status": "example_unreviewed", + "usd_per_million_tokens": {"m": {"input": float("nan"), "output": 1.0}}, + } + with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as handle: + json.dump(payload, handle) + path = handle.name + try: + with self.assertRaises(CostQualityContractError): + load_pricing_scenario(path) + finally: + os.unlink(path) + + +class TestOfflineComparison(unittest.TestCase): + def setUp(self) -> None: + self.manifest = load_task_manifest(str(MANIFEST_PATH)) + self.tasks = locked_evaluation_tasks(self.manifest) + self.answers = { + "digit_sum_reasoning": { + "direct_worker": "200", + "route_once": "200", + "bounded_conduct": "200", + }, + "linear_equation_solution": { + "direct_worker": "7", + "route_once": "7", + "bounded_conduct": "wrong", + }, + "capital_recall_france": { + "direct_worker": "Paris", + "route_once": "Paris", + "bounded_conduct": "Paris", + }, + "sequence_next_fibonacci": { + "direct_worker": "21", + "route_once": "20", + "bounded_conduct": "21", + }, + } + self.runners = build_scripted_policy_runners(self.answers, model_id="mock-scripted") + self.pricing = load_pricing_scenario(str(PRICING_PATH)) + + def test_format_task_prompt_preserves_scorable_body(self) -> None: + task = self.tasks[0] + formatted = format_task_prompt(task) + self.assertIn(task["prompt"], formatted) + self.assertIn(task["task_id"], formatted) + # Marker must not create scorer leakage against the bare expected value path. + scored = score_task_answer(task, formatted) + self.assertEqual(scored["score"], 0.0) + + def test_run_offline_report_shape(self) -> None: + report = run_offline_cost_quality( + tasks=self.tasks, + policy_runners=self.runners, + model_id="mock-scripted", + pricing_scenario=self.pricing, + ) + self.assertEqual(report["measurement_status"], "offline_cost_quality") + self.assertEqual(report["task_count"], len(self.tasks)) + self.assertGreaterEqual(report["cell_count"], len(self.tasks) * 3) + self.assertTrue(all(c["actual_api_cost"] == "unknown" for c in report["cells"])) + self.assertTrue( + all( + isinstance(c["hypothetical_paid_cost"], float) + or c["hypothetical_paid_cost"] == "unknown" + for c in report["cells"] + ) + ) + names = {s["policy_name"] for s in report["policy_summaries"]} + self.assertIn("route_once", names) + self.assertIn("bounded_conduct", names) + self.assertIn("hindsight_best_single", names) + md = render_cost_quality_markdown(report) + self.assertIn("offline_cost_quality", md) + self.assertNotIn("NVIDIA_NIM_API_KEY", md) + + def test_unpriced_model_keeps_unknown_hypothetical(self) -> None: + runners = build_scripted_policy_runners(self.answers, model_id="never-priced") + report = run_offline_cost_quality( + tasks=self.tasks[:1], + policy_runners=runners, + model_id="never-priced", + pricing_scenario=self.pricing, + include_hindsight_best_single=False, + ) + self.assertTrue(all(c["hypothetical_paid_cost"] == "unknown" for c in report["cells"])) + for summary in report["policy_summaries"]: + self.assertEqual(summary["hypothetical_paid_cost_mean"], "unknown") + + def test_failed_runner_records_outcome(self) -> None: + def boom(_prompt: str) -> dict: + raise RuntimeError("provider down") + + task = self.tasks[0] + cell = run_policy_cell( + task=task, + policy_name="route_once", + runner=boom, + model_id="mock-scripted", + pricing_scenario=self.pricing, + ) + self.assertEqual(cell["outcome"], "failed") + self.assertEqual(cell["error_class"], "RuntimeError") + self.assertEqual(cell["score"], 0.0) + self.assertEqual(cell["actual_api_cost"], "unknown") + self.assertEqual(cell["hypothetical_paid_cost"], "unknown") + self.assertEqual(cell["prompt_tokens"], 0) + self.assertEqual(cell["completion_tokens"], 0) + self.assertEqual(cell["call_count"], 0) + self.assertEqual(cell["usage_source"], "none") + + def test_scripted_answers_reject_non_mapping_values(self) -> None: + from contextual_orchestrator.nim_cost_quality import validate_scripted_answers + + with self.assertRaises(CostQualityContractError): + validate_scripted_answers({"digit_sum_reasoning": "200"}) + with self.assertRaises(CostQualityContractError): + validate_scripted_answers({"digit_sum_reasoning": {"route_once": 7}}) + with self.assertRaises(CostQualityContractError): + build_scripted_policy_runners({"digit_sum_reasoning": "200"}) + + def test_pareto_excludes_unknown_cost(self) -> None: + summaries = [ + { + "policy_name": "route_once", + "mean_score": 0.9, + "mean_latency_ms": 10.0, + "hypothetical_paid_cost_mean": "unknown", + }, + { + "policy_name": "bounded_conduct", + "mean_score": 0.8, + "mean_latency_ms": 20.0, + "hypothetical_paid_cost_mean": 0.01, + }, + { + "policy_name": "direct_worker", + "mean_score": 0.7, + "mean_latency_ms": 5.0, + "hypothetical_paid_cost_mean": 0.005, + }, + ] + frontiers = build_pareto_frontiers(summaries) + cost_names = {row["policy_name"] for row in frontiers["quality_hypothetical_cost"]} + self.assertNotIn("route_once", cost_names) + self.assertIn("direct_worker", cost_names) + + def test_chat_eligible_filters_embeddings(self) -> None: + ids = chat_eligible_model_ids( + ["meta/llama-3-8b-instruct", "nvidia/nv-embedqa-e5-v5", " "] + ) + self.assertEqual(ids, ["meta/llama-3-8b-instruct"]) + + def test_plan_from_discovery_aligns_with_dry_run(self) -> None: + plan = plan_from_discovery_models( + ["meta/llama-3-8b-instruct", "nvidia/nv-embedqa-e5-v5"], + hard_request_budget=50, + ) + self.assertEqual(plan["measurement_status"], "dry_run_plan") + self.assertEqual(plan["admission_status"], "admitted") + self.assertTrue(all(c["actual_api_cost"] == "unknown" for c in plan["comparison_cells"])) + + def test_summarize_empty_raises_via_offline(self) -> None: + with self.assertRaises(CostQualityContractError): + run_offline_cost_quality(tasks=[], policy_runners=self.runners) + + def test_summarize_policy_cells_mean(self) -> None: + cells = [ + { + "policy_name": "route_once", + "score": 1.0, + "latency_ms": 10.0, + "hypothetical_paid_cost": 0.1, + "outcome": "success", + }, + { + "policy_name": "route_once", + "score": 0.0, + "latency_ms": 20.0, + "hypothetical_paid_cost": 0.3, + "outcome": "success", + }, + ] + summaries = summarize_policy_cells(cells) + self.assertEqual(len(summaries), 1) + self.assertAlmostEqual(summaries[0]["mean_score"], 0.5) + self.assertAlmostEqual(float(summaries[0]["hypothetical_paid_cost_mean"]), 0.2) + + +class TestNoSecretInModuleSurface(unittest.TestCase): + def test_module_never_references_copilot_token(self) -> None: + source = (ROOT / "contextual_orchestrator" / "nim_cost_quality.py").read_text( + encoding="utf-8" + ) + self.assertNotIn("COPILOT_GITHUB_TOKEN", source) + self.assertNotIn("os.getenv", source) + self.assertNotIn("os.environ", source) + + +class TestOrchestratorPolicyRunners(unittest.TestCase): + def test_mock_orchestrator_route_and_conduct(self) -> None: + agents = [ + ModelAgent( + "general_agent", + "mock-generalist", + base_url="mock://generalist", + tags=("reasoning", "writing", "planning"), + priority=1, + ), + ModelAgent( + "builder_agent", + "mock-builder", + base_url="mock://builder", + tags=("coding", "debugging"), + priority=2, + ), + ModelAgent( + "reviewer_agent", + "mock-reviewer", + base_url="mock://reviewer", + tags=("verification", "review"), + priority=3, + ), + ] + orchestrator = TaskOrchestrator(agents) + runners = build_orchestrator_policy_runners(orchestrator) + self.assertEqual( + set(runners), + {"direct_worker", "route_once", "bounded_conduct"}, + ) + route = runners["route_once"]("hello cost quality") + self.assertEqual(route["mode"], "route_once") + self.assertTrue(route["answer"]) + self.assertGreaterEqual(len(route["trace"]), 1) + conduct = runners["bounded_conduct"]("hello cost quality") + self.assertEqual(conduct["mode"], "bounded_conduct") + self.assertTrue(conduct["answer"]) + self.assertGreaterEqual(len(conduct["trace"]), 2) + + manifest = load_task_manifest(str(MANIFEST_PATH)) + tasks = locked_evaluation_tasks(manifest)[:2] + report = run_offline_cost_quality( + tasks=tasks, + policy_runners=runners, + model_id="mock-orchestrator", + pricing_scenario=None, + ) + self.assertEqual(report["measurement_status"], "offline_cost_quality") + self.assertTrue(all(c["actual_api_cost"] == "unknown" for c in report["cells"])) + self.assertTrue(all(c["hypothetical_paid_cost"] == "unknown" for c in report["cells"])) + # Mock answers echo the prompt; scorers must not invent quality from that alone. + self.assertTrue(all(c["outcome"] == "success" for c in report["cells"] if c["policy_name"] != "hindsight_best_single" or True)) + + def test_orchestrator_runners_reject_invalid(self) -> None: + with self.assertRaises(CostQualityContractError): + build_orchestrator_policy_runners(None) + with self.assertRaises(CostQualityContractError): + build_orchestrator_policy_runners(object()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_nim_discovery.py b/tests/test_nim_discovery.py new file mode 100644 index 000000000..8b2512a15 --- /dev/null +++ b/tests/test_nim_discovery.py @@ -0,0 +1,454 @@ +"""NIM model discovery — offline fixtures always; live catalog only with explicit opt-in.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import sys + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from contextual_orchestrator.credentials import InMemoryCredentialBackend, set_backend # noqa: E402 +from contextual_orchestrator.nim_discovery import ( # noqa: E402 + DEFAULT_NIM_MODELS_URL, + NIM_CATALOG_MAX_BYTES, + NIM_CREDENTIAL_NAME, + NimDiscoveryError, + build_benchmark_plan_dry_run, + build_capability_inventory, + classify_model_capability_hint, + discover_nim_models, + models_to_agent_pool_entries, + validate_nim_models_url, + _extract_model_ids, + _slug_model_id, +) +from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402 + + +def test_extract_model_ids_from_openai_style_payload() -> None: + payload = { + "data": [ + {"id": "meta/llama-3.1-70b-instruct"}, + {"id": "google/gemma-2-9b-it"}, + {"id": "meta/llama-3.1-70b-instruct"}, # dedupe + ] + } + ids = _extract_model_ids(payload) + assert ids == ["google/gemma-2-9b-it", "meta/llama-3.1-70b-instruct"] + + +def test_slug_model_id_is_multi_word_snake_case() -> None: + assert "_" in _slug_model_id("gpt4") + assert _slug_model_id("meta/llama-3.1-70b-instruct").startswith("meta") + + +def test_discover_without_credential_is_honest_missing_status() -> None: + set_backend(InMemoryCredentialBackend()) + try: + report = discover_nim_models() + finally: + set_backend(None) + assert report["measurement_status"] == "credential_missing" + assert report["model_ids"] == [] + assert report["credential_name"] == NIM_CREDENTIAL_NAME + assert report["source_url"] == DEFAULT_NIM_MODELS_URL + + +def test_discover_with_fixture_transport_returns_sorted_ids() -> None: + backend = InMemoryCredentialBackend() + backend.set(NIM_CREDENTIAL_NAME, "nvapi-test-not-real") + set_backend(backend) + + def transport(request, timeout): # noqa: ANN001 + headers = {k.lower(): v for k, v in request.headers.items()} + assert "authorization" in headers + assert request.full_url == DEFAULT_NIM_MODELS_URL + body = json.dumps( + {"data": [{"id": "z-model"}, {"id": "a-model"}, {"id": "m-model"}]} + ).encode("utf-8") + return body + + try: + report = discover_nim_models(transport=transport) + finally: + set_backend(None) + + assert report["measurement_status"] == "offline_fixture" + assert report["model_ids"] == ["a-model", "m-model", "z-model"] + assert report["model_count"] == 3 + + +def test_discover_rejects_oversized_catalog_body() -> None: + backend = InMemoryCredentialBackend() + backend.set(NIM_CREDENTIAL_NAME, "nvapi-test-not-real") + set_backend(backend) + + def transport(request, timeout): # noqa: ANN001 + return b"x" * (NIM_CATALOG_MAX_BYTES + 1) + + try: + with pytest.raises(NimDiscoveryError, match="exceeds bound"): + discover_nim_models(transport=transport) + finally: + set_backend(None) + + +def test_models_to_agent_pool_entries_are_loadable_agents() -> None: + entries = models_to_agent_pool_entries( + ["meta/llama-3.1-8b-instruct", "google/gemma-2-2b-it"] + ) + assert len(entries) == 2 + agents = [ModelAgent.from_dict(row) for row in entries] + orch = TaskOrchestrator(agents) + # route path still works offline with mock:// default replaced — use mock urls + mock_agents = [ + ModelAgent( + agent.id, + agent.model, + base_url="mock://local", + tags=agent.tags, + priority=agent.priority, + ) + for agent in agents + ] + result = TaskOrchestrator(mock_agents).route_once( + [{"role": "user", "content": "nim pool route"}] + ) + assert result["mode"] == "route" + assert result["answer"] + assert all(entry["credential_key"] == NIM_CREDENTIAL_NAME for entry in entries) + assert all("_" in entry["id"] for entry in entries) + + +def test_agent_ids_unique_on_slug_and_truncation_collision() -> None: + # normalized-slug collision: different punctuation → same slug + entries = models_to_agent_pool_entries(["a/b", "a-b", "a_b"]) + ids = [row["id"] for row in entries] + assert len(ids) == len(set(ids)) + assert ids[0] == "nim_a_b_agent" + assert ids[1] == "nim_a_b_agent_2" + assert ids[2] == "nim_a_b_agent_3" + + # truncation collision: long ids that share the first 48 slug chars + long_a = "prefix/" + ("x" * 80) + "-one" + long_b = "prefix/" + ("x" * 80) + "-two" + # After slug+48, both collapse; ensure unique agent ids. + slug_a = _slug_model_id(long_a) + slug_b = _slug_model_id(long_b) + assert slug_a == slug_b + long_entries = models_to_agent_pool_entries([long_a, long_b]) + long_ids = [row["id"] for row in long_entries] + assert len(long_ids) == len(set(long_ids)) + assert long_ids[0].endswith("_agent") + assert long_ids[1].endswith("_agent_2") + + +def test_validate_nim_models_url_allowlist() -> None: + assert validate_nim_models_url(DEFAULT_NIM_MODELS_URL) == DEFAULT_NIM_MODELS_URL + assert ( + validate_nim_models_url("https://api.nvcf.nvidia.com/v1/models/") + == "https://api.nvcf.nvidia.com/v1/models" + ) + with pytest.raises(NimDiscoveryError): + validate_nim_models_url("http://integrate.api.nvidia.com/v1/models") + with pytest.raises(NimDiscoveryError): + validate_nim_models_url("https://evil.example/v1/models") + with pytest.raises(NimDiscoveryError): + validate_nim_models_url("https://integrate.api.nvidia.com/v1/chat/completions") + with pytest.raises(NimDiscoveryError): + validate_nim_models_url("https://integrate.api.nvidia.com:8443/v1/models") + with pytest.raises(NimDiscoveryError): + validate_nim_models_url("") + with pytest.raises(NimDiscoveryError): + validate_nim_models_url(" ") + with pytest.raises(NimDiscoveryError): + validate_nim_models_url(None) # type: ignore[arg-type] + with pytest.raises(NimDiscoveryError): + validate_nim_models_url("https://user:pass@integrate.api.nvidia.com/v1/models") + with pytest.raises(NimDiscoveryError): + validate_nim_models_url("https://integrate.api.nvidia.com/v1/models?x=1") + with pytest.raises(NimDiscoveryError): + validate_nim_models_url("https://integrate.api.nvidia.com/v1/models#frag") + + +def test_unique_agent_id_suffix_chain_covers_collisions() -> None: + """Force multi-suffix agent ids when more than two models share a slug.""" + from contextual_orchestrator.nim_discovery import _unique_agent_id + + used: set[str] = set() + first = _unique_agent_id("a/b", used) + used.add(first) + second = _unique_agent_id("a-b", used) + used.add(second) + third = _unique_agent_id("a_b", used) + used.add(third) + assert first == "nim_a_b_agent" + assert second == "nim_a_b_agent_2" + assert third == "nim_a_b_agent_3" + assert len(used) == 3 + + +def test_discover_rejects_non_allowlisted_url_before_credential_use() -> None: + backend = InMemoryCredentialBackend() + backend.set(NIM_CREDENTIAL_NAME, "nvapi-must-not-leak") + set_backend(backend) + called = {"n": 0} + + def transport(request, timeout): # noqa: ANN001 + called["n"] += 1 + return b"{}" + + try: + with pytest.raises(NimDiscoveryError): + discover_nim_models( + models_url="https://attacker.example/v1/models", + transport=transport, + ) + finally: + set_backend(None) + assert called["n"] == 0 + + +def test_extract_model_ids_handles_list_and_non_list_shapes() -> None: + assert _extract_model_ids(["alpha_model", "beta_model"]) == ["alpha_model", "beta_model"] + assert _extract_model_ids({"models": [{"model": "x_y"}]}) == ["x_y"] + assert _extract_model_ids({"data": "not-a-list"}) == [] + assert _extract_model_ids(42) == [] + + +def test_slug_model_id_edge_cases() -> None: + assert _slug_model_id("!!!") == "unnamed_model" + assert _slug_model_id("simple") == "simple_model" + assert "__" not in _slug_model_id("a--b__c") + + +def test_discover_uses_urllib_when_transport_omitted(monkeypatch) -> None: # noqa: ANN001 + backend = InMemoryCredentialBackend() + backend.set(NIM_CREDENTIAL_NAME, "nvapi-test") + set_backend(backend) + + class _Resp: + def __init__(self): + body = json.dumps({"data": [{"id": "live_fixture_model"}]}).encode("utf-8") + self._body = body + self.headers = {"Content-Length": str(len(body))} + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self, n=-1): + if n is None or n < 0: + return self._body + return self._body[:n] + + def fake_urlopen(request, timeout=None, context=None): # noqa: ANN001 + assert request.get_header("Authorization") or request.headers + assert request.full_url == DEFAULT_NIM_MODELS_URL + return _Resp() + + monkeypatch.setattr("contextual_orchestrator.nim_discovery.urllib.request.urlopen", fake_urlopen) + try: + report = discover_nim_models() + finally: + set_backend(None) + assert report["measurement_status"] == "live_nim_catalog" + assert report["model_ids"] == ["live_fixture_model"] + + +def test_live_nim_catalog_when_env_seeded_into_kv() -> None: + """Optional live check: requires RUN_LIVE_NIM_TESTS=1 and NVIDIA_NIM_API_KEY. + + Env is bootstrap transport only; discover_nim_models still reads via get_credential. + Default CI stays hermetic (offline fixtures only). + """ + if os.environ.get("RUN_LIVE_NIM_TESTS", "").strip() != "1": + return + raw = os.environ.get("NVIDIA_NIM_API_KEY", "").strip() + if not raw: + return + backend = InMemoryCredentialBackend() + backend.set(NIM_CREDENTIAL_NAME, raw) + set_backend(backend) + try: + report = discover_nim_models() + finally: + set_backend(None) + assert report["measurement_status"] == "live_nim_catalog" + assert report["model_count"] >= 1 + assert all(isinstance(mid, str) and mid for mid in report["model_ids"]) + + +def test_role_temperature_defaults_differ_by_paper_role() -> None: + from contextual_orchestrator.orchestrator import OrchestrationPolicy + + policy = OrchestrationPolicy() + assert policy.temperature_for_role("verifier") < policy.temperature_for_role("worker") + assert policy.temperature_for_role("thinker") <= policy.temperature_for_role("worker") + snap = policy.as_dict() + assert "role_temperature" in snap + assert set(snap["role_temperature"]) >= {"thinker", "worker", "verifier", "synthesizer"} + + +def test_discover_nim_models_cli_prints_credential_missing(capsys) -> None: + from contextual_orchestrator.__main__ import _discover_nim_models_command + set_backend(InMemoryCredentialBackend()) + try: + _discover_nim_models_command([]) + finally: + set_backend(None) + out = json.loads(capsys.readouterr().out) + assert out["measurement_status"] == "credential_missing" + + +def test_discover_nim_models_cli_rejects_evil_url(capsys) -> None: + from contextual_orchestrator.__main__ import _discover_nim_models_command + + with pytest.raises(SystemExit): + _discover_nim_models_command(["--models-url", "https://evil.example/v1/models"]) + + +if __name__ == "__main__": # pragma: no cover + test_extract_model_ids_from_openai_style_payload() + test_slug_model_id_is_multi_word_snake_case() + test_discover_without_credential_is_honest_missing_status() + test_discover_with_fixture_transport_returns_sorted_ids() + test_models_to_agent_pool_entries_are_loadable_agents() + test_agent_ids_unique_on_slug_and_truncation_collision() + test_validate_nim_models_url_allowlist() + test_live_nim_catalog_when_env_seeded_into_kv() + print("ok") + + +def test_classify_model_capability_hint_modalities() -> None: + assert classify_model_capability_hint("meta/llama-3.1-8b-instruct") == "chat" + assert classify_model_capability_hint("nvidia/nv-embedqa-e5-v5") == "embeddings" + assert classify_model_capability_hint("stabilityai/stable-diffusion-xl") == "image" + assert classify_model_capability_hint("openai/whisper-large-v3") == "audio" + assert classify_model_capability_hint("vendor/video-gen-1") == "video" + assert classify_model_capability_hint("") == "unsupported" + assert classify_model_capability_hint("obscure-weights-v9") == "unknown" + + +def test_build_capability_inventory_is_offline_honest() -> None: + inv = build_capability_inventory( + ["meta/llama-3.1-8b-instruct", "nvidia/nv-embedqa-e5-v5", "meta/llama-3.1-8b-instruct"] + ) + assert inv["measurement_status"] == "offline_capability_hints" + assert inv["model_count"] == 2 + assert inv["capability_counts"]["chat"] == 1 + assert inv["capability_counts"]["embeddings"] == 1 + + +def test_benchmark_plan_dry_run_unknown_cost_and_budget_gate() -> None: + plan = build_benchmark_plan_dry_run( + ["meta/llama-3.1-8b-instruct", "nvidia/nv-embedqa-e5-v5"], + hard_request_budget=100, + max_steps=5, + ) + assert plan["measurement_status"] == "dry_run_plan" + assert plan["admission_status"] == "admitted" + assert plan["chat_eligible_model_count"] == 1 + # direct(1)+route(1)+conduct(max_steps); hindsight reuses direct (0 extra) + assert plan["per_model_call_budget"] == 7 + assert plan["planned_request_count"] == 7 + assert all(cell["hypothetical_paid_cost"] == "unknown" for cell in plan["comparison_cells"]) + assert all(cell["actual_api_cost"] == "unknown" for cell in plan["comparison_cells"]) + # embeddings excluded from chat comparison cells + assert all("embed" not in cell["model_id"] for cell in plan["comparison_cells"]) + conduct_cells = [c for c in plan["comparison_cells"] if c["policy_name"] == "bounded_conduct"] + assert all(c["planned_provider_calls"] == 5 for c in conduct_cells) + + tight = build_benchmark_plan_dry_run( + ["a-chat-model", "b-chat-instruct"], + hard_request_budget=1, + ) + assert tight["admission_status"] == "rejected_budget_exceeded" + assert tight["fits_hard_request_budget"] is False + assert tight["planned_request_count"] == 14 + + +def test_benchmark_plan_dry_run_rejects_invalid_bounds() -> None: + with pytest.raises(NimDiscoveryError): + build_benchmark_plan_dry_run(["x"], max_steps=6) + with pytest.raises(NimDiscoveryError): + build_benchmark_plan_dry_run(["x"], hard_request_budget=0) + with pytest.raises(NimDiscoveryError): + build_benchmark_plan_dry_run(["x"], task_manifest_id="") + + +def test_classify_probe_http_status_and_result() -> None: + from contextual_orchestrator.nim_discovery import ( + classify_probe_http_status, + classify_probe_result, + run_capability_probes_dry_run, + build_capability_probe_plan, + ) + + assert classify_probe_http_status(429) == "rate_limited" + assert classify_probe_http_status(404) == "unsupported" + row = classify_probe_result( + model_id="meta/llama-3-8b-instruct", + probe_kind="chat", + status_code=200, + body={"choices": [{"message": {"content": "hi"}}]}, + ) + assert row["outcome"] == "chat" + bad = classify_probe_result( + model_id="x", + probe_kind="chat", + status_code=200, + body={"unexpected": True}, + ) + assert bad["outcome"] == "malformed" + timed = classify_probe_result( + model_id="x", + probe_kind="chat", + status_code=0, + error_class="TimeoutError", + ) + assert timed["outcome"] == "timeout" + + +def test_capability_probe_plan_and_dry_run_budget() -> None: + from contextual_orchestrator.nim_discovery import ( + NimDiscoveryError, + build_capability_probe_plan, + run_capability_probes_dry_run, + ) + + plan = build_capability_probe_plan( + ["meta/llama-3-8b-instruct", "nvidia/nv-embedqa-e5-v5"], + hard_request_budget=10, + probe_kinds=("chat", "embeddings"), + ) + assert plan["measurement_status"] == "offline_probe_plan" + assert plan["planned_request_count"] == 4 + assert plan["admission_status"] == "admitted" + + tight = build_capability_probe_plan(["a", "b", "c"], hard_request_budget=2) + assert tight["admission_status"] == "rejected_budget_exceeded" + + fixtures = [ + { + "model_id": "meta/llama-3-8b-instruct", + "probe_kind": "chat", + "status_code": 200, + "body": {"choices": [{"message": {"content": "ok"}}]}, + }, + {"model_id": "x", "probe_kind": "chat", "status_code": 429}, + ] + report = run_capability_probes_dry_run(fixtures, hard_request_budget=10) + assert report["measurement_status"] == "offline_probe_results" + assert report["outcome_counts"]["chat"] == 1 + assert report["outcome_counts"]["rate_limited"] == 1 + assert report["chat_eligible_model_ids"] == ["meta/llama-3-8b-instruct"] + + with pytest.raises(NimDiscoveryError): + run_capability_probes_dry_run(fixtures, hard_request_budget=1)