diff --git a/src/infospace_bench/routing_config.py b/src/infospace_bench/routing_config.py index 514267a..5418c94 100644 --- a/src/infospace_bench/routing_config.py +++ b/src/infospace_bench/routing_config.py @@ -11,9 +11,10 @@ deterministically testable. from __future__ import annotations +import os from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, Callable, Mapping import yaml @@ -27,6 +28,16 @@ SUPPORTED_PROVIDERS: frozenset[str] = frozenset( {"openrouter", "claude_code", "openai", "gemini"} ) +# Default env var per provider when the config does not name one explicitly. +DEFAULT_API_KEY_ENV: dict[str, str] = { + "openrouter": "OPENROUTER_API_KEY", + "openai": "OPENAI_API_KEY", + "gemini": "GOOGLE_API_KEY", + # claude_code shells out to the local CLI and needs no API key +} + +AdapterFactory = Callable[["RoutingCandidateConfig", Mapping[str, str]], Any] + @dataclass(frozen=True) class RoutingCandidateConfig: @@ -263,3 +274,131 @@ def _optional_quality_floor(value: Any, name: str, source: str) -> float | None: {"source": source, "name": name, "value": floor}, ) return floor + + +# --------------------------------------------------------------------------- +# T02 — build live policies and adapters from a parsed config +# --------------------------------------------------------------------------- + + +def build_routing_policy_from_config( + config: RoutingConfig, + *, + workspace: str | Path | None = None, + env: Mapping[str, str] | None = None, + adapter_factory: AdapterFactory | None = None, +) -> Any: + """Materialise a parsed config into a live llm-connect routing policy. + + Returns an ``AdaptiveRoutingPolicy`` when the config sets a + ``default_quality_floor``, any per-task ``quality_floor``, or a + ``ledger_path``; otherwise returns a static ``RoutingPolicy``. + + ``adapter_factory`` is an opt-in override that builds an + ``LLMAdapter`` from a ``RoutingCandidateConfig`` + env mapping. Tests + inject a factory to avoid hitting real provider constructors; the + default factory resolves API keys from ``env`` and instantiates the + matching llm-connect adapter. + + Fails fast (before any network call) when a candidate's required API + key env var is missing from ``env``. + """ + from llm_connect.routing import AdaptiveRoutingPolicy, RoutingPolicy, RoutingRule + + environment: Mapping[str, str] = env if env is not None else os.environ + factory: AdapterFactory = adapter_factory or _default_adapter_factory + + adapters_by_id: dict[str, Any] = {} + rules: list[RoutingRule] = [] + for task in config.task_types: + candidates: list[Any] = [] + for candidate in task.candidates: + if candidate.id not in adapters_by_id: + adapters_by_id[candidate.id] = factory(candidate, environment) + candidates.append(adapters_by_id[candidate.id]) + prefer = candidates[0] + prefer_candidate = task.candidates[0] + fallback = candidates[1] if len(candidates) > 1 else None + rules.append( + RoutingRule( + task_type=task.task_type, + prefer=prefer, + max_cost_per_1k=prefer_candidate.max_cost_per_1k, + fallback=fallback, + ) + ) + + use_adaptive = ( + config.default_quality_floor is not None + or any(task.quality_floor is not None for task in config.task_types) + or config.ledger_path is not None + ) + if not use_adaptive: + return RoutingPolicy(rules=rules) + + from llm_connect.quality import QualityLedger + + ledger: QualityLedger | None = None + if config.ledger_path: + ledger_path = Path(config.ledger_path) + if not ledger_path.is_absolute() and workspace is not None: + ledger_path = Path(workspace) / ledger_path + ledger_path.parent.mkdir(parents=True, exist_ok=True) + ledger = QualityLedger(path=ledger_path) + return AdaptiveRoutingPolicy( + rules=rules, + ledger=ledger, + adapters_by_id=dict(adapters_by_id), + ) + + +def _default_adapter_factory( + candidate: RoutingCandidateConfig, env: Mapping[str, str] +) -> Any: + """Build a real llm-connect adapter for one config candidate. + + API keys are resolved from ``env`` before construction; a missing key + raises ``missing_routing_api_key`` rather than letting the adapter + blow up later mid-run. + """ + provider = candidate.provider + if provider == "claude_code": + from llm_connect.claude_code import ClaudeCodeAdapter + + return ClaudeCodeAdapter(model=candidate.model) + + env_var = candidate.api_key_env or DEFAULT_API_KEY_ENV.get(provider, "") + api_key = env.get(env_var, "") if env_var else "" + if not api_key: + raise InfospaceError( + "missing_routing_api_key", + f"Candidate {candidate.id!r} ({provider}) needs API key from " + f"env var {env_var!r}, but it is unset", + { + "candidate_id": candidate.id, + "provider": provider, + "api_key_env": env_var, + }, + ) + + if provider == "openrouter": + from llm_connect.openrouter import OpenRouterAdapter + + return OpenRouterAdapter(model=candidate.model, api_key=api_key) + if provider == "openai": + from llm_connect.openai import OpenAIAdapter + + return OpenAIAdapter(model=candidate.model, api_key=api_key) + if provider == "gemini": + from llm_connect.gemini import GeminiAdapter + + return GeminiAdapter(model=candidate.model, api_key=api_key) + + # Should have been rejected by the parser; defensive guard for callers + # that build a RoutingConfig programmatically without going through + # parse_routing_config. + raise InfospaceError( + "unsupported_routing_provider", + f"Cannot build adapter for unsupported provider {provider!r}", + {"candidate_id": candidate.id, "provider": provider}, + ) diff --git a/tests/test_routing_config.py b/tests/test_routing_config.py index bda0869..69e2ddc 100644 --- a/tests/test_routing_config.py +++ b/tests/test_routing_config.py @@ -261,6 +261,189 @@ def test_rejects_non_mapping_stage_map() -> None: assert exc_info.value.code == "invalid_routing_stage_map" +# --------------------------------------------------------------------------- +# T02 — loader that materialises a config into a live llm-connect policy +# --------------------------------------------------------------------------- + + +def _fake_adapter_factory_record(record: list): + """Return a factory that records calls and returns a sentinel string.""" + def _factory(candidate, env): + record.append({"id": candidate.id, "provider": candidate.provider, "model": candidate.model}) + return f"adapter:{candidate.id}" + return _factory + + +def test_build_routing_policy_returns_static_when_no_adaptive_signals() -> None: + from llm_connect.routing import RoutingPolicy + from infospace_bench.routing_config import build_routing_policy_from_config + + config = parse_routing_config(MINIMAL) + calls: list[dict] = [] + policy = build_routing_policy_from_config( + config, adapter_factory=_fake_adapter_factory_record(calls) + ) + + assert isinstance(policy, RoutingPolicy) + assert type(policy).__name__ == "RoutingPolicy", "no adaptive signals -> static policy" + assert len(policy.rules) == 1 + assert policy.rules[0].task_type == "summarize-source" + assert policy.rules[0].prefer == "adapter:openrouter:gpt-4o-mini" + assert policy.rules[0].fallback is None + assert calls and calls[0]["provider"] == "openrouter" + + +def test_build_routing_policy_returns_adaptive_when_ledger_path_set(tmp_path: Path) -> None: + from llm_connect.routing import AdaptiveRoutingPolicy + from infospace_bench.routing_config import build_routing_policy_from_config + + data = { + **MINIMAL, + "ledger_path": "output/routing/quality.jsonl", + } + config = parse_routing_config(data) + policy = build_routing_policy_from_config( + config, + workspace=tmp_path, + adapter_factory=_fake_adapter_factory_record([]), + ) + + assert isinstance(policy, AdaptiveRoutingPolicy) + assert policy.ledger is not None + expected_path = tmp_path / "output" / "routing" / "quality.jsonl" + assert Path(policy.ledger.path) == expected_path + + +def test_build_routing_policy_returns_adaptive_when_quality_floor_set() -> None: + from llm_connect.routing import AdaptiveRoutingPolicy + from infospace_bench.routing_config import build_routing_policy_from_config + + data = { + **MINIMAL, + "default_quality_floor": 0.8, + } + config = parse_routing_config(data) + policy = build_routing_policy_from_config( + config, adapter_factory=_fake_adapter_factory_record([]) + ) + + assert isinstance(policy, AdaptiveRoutingPolicy) + assert policy.ledger is None # no ledger_path set + + +def test_build_routing_policy_routes_fallback_for_multi_candidate_rule() -> None: + from infospace_bench.routing_config import build_routing_policy_from_config + + data = { + "schema_version": 1, + "task_types": { + "extract-entities": { + "candidates": [ + { + "id": "openrouter:cheap", + "provider": "openrouter", + "model": "openai/gpt-4o-mini", + "max_cost_per_1k": 0.001, + }, + { + "id": "openrouter:smart", + "provider": "openrouter", + "model": "anthropic/claude-3.5-sonnet", + }, + ], + }, + }, + } + config = parse_routing_config(data) + policy = build_routing_policy_from_config( + config, adapter_factory=_fake_adapter_factory_record([]) + ) + + rule = policy.rules[0] + assert rule.prefer == "adapter:openrouter:cheap" + assert rule.max_cost_per_1k == 0.001 + assert rule.fallback == "adapter:openrouter:smart" + + +def test_build_routing_policy_resolves_api_key_from_env() -> None: + from infospace_bench.routing_config import ( + build_routing_policy_from_config, + _default_adapter_factory, + ) + + config = parse_routing_config(MINIMAL) + # Smoke: real factory with a fake env should construct an OpenRouterAdapter. + env = {"OPENROUTER_API_KEY": "sk-fake-test-key"} + policy = build_routing_policy_from_config(config, env=env) + rule = policy.rules[0] + # The constructed adapter is an OpenRouterAdapter from llm-connect. + from llm_connect.openrouter import OpenRouterAdapter + assert isinstance(rule.prefer, OpenRouterAdapter) + + +def test_build_routing_policy_fails_fast_on_missing_api_key() -> None: + from infospace_bench.routing_config import build_routing_policy_from_config + + config = parse_routing_config(MINIMAL) + # Empty env — the candidate's required env var is unset. + with pytest.raises(InfospaceError) as exc_info: + build_routing_policy_from_config(config, env={}) + assert exc_info.value.code == "missing_routing_api_key" + assert exc_info.value.detail["provider"] == "openrouter" + + +def test_build_routing_policy_claude_code_needs_no_api_key() -> None: + from infospace_bench.routing_config import build_routing_policy_from_config + from llm_connect.claude_code import ClaudeCodeAdapter + + data = { + "schema_version": 1, + "task_types": { + "baseline": { + "candidates": [ + {"id": "claude-code", "provider": "claude_code", "model": "claude-opus-4-7"}, + ], + }, + }, + } + config = parse_routing_config(data) + policy = build_routing_policy_from_config(config, env={}) + + assert isinstance(policy.rules[0].prefer, ClaudeCodeAdapter) + + +def test_build_routing_policy_honours_custom_api_key_env() -> None: + from infospace_bench.routing_config import build_routing_policy_from_config + from llm_connect.openrouter import OpenRouterAdapter + + data = { + "schema_version": 1, + "task_types": { + "summarize-source": { + "candidates": [ + { + "id": "openrouter:gpt-4o-mini", + "provider": "openrouter", + "model": "openai/gpt-4o-mini", + "api_key_env": "ALT_OPENROUTER_KEY", + }, + ], + }, + }, + } + config = parse_routing_config(data) + + with pytest.raises(InfospaceError) as exc_info: + build_routing_policy_from_config(config, env={"OPENROUTER_API_KEY": "wrong-default"}) + assert exc_info.value.code == "missing_routing_api_key" + assert exc_info.value.detail["api_key_env"] == "ALT_OPENROUTER_KEY" + + policy = build_routing_policy_from_config( + config, env={"ALT_OPENROUTER_KEY": "sk-fake"} + ) + assert isinstance(policy.rules[0].prefer, OpenRouterAdapter) + + def test_rejects_non_string_ledger_path() -> None: payload = { "schema_version": 1, diff --git a/workplans/IB-WP-0020-provider-routing-cli.md b/workplans/IB-WP-0020-provider-routing-cli.md index bd56ce5..698f2f7 100644 --- a/workplans/IB-WP-0020-provider-routing-cli.md +++ b/workplans/IB-WP-0020-provider-routing-cli.md @@ -93,7 +93,7 @@ state_hub_task_id: "39597441-22ab-4dcf-b68d-b045823a9374" ```task id: IB-WP-0020-T02 -status: todo +status: done priority: high state_hub_task_id: "5e38514b-ad6a-4d39-8716-f812f241d9fd" ```