""" Bridge between infospace-bench's ``AssistedGenerationAdapter`` protocol and llm-connect's ``RoutingPolicy`` / ``AdaptiveRoutingPolicy`` primitives (LLM-WP-0004). Lets a generation run delegate each stage to a task-typed route without touching ``workflow.py``. The mapping from infospace-bench workflow stage ids to llm-connect task types is the consumer side of LLM-WP-0004's scope guardrail: llm-connect ships the routing primitives, infospace-bench names the tasks. """ from __future__ import annotations from dataclasses import dataclass, field from typing import Any from llm_connect.adapter import LLMAdapter from llm_connect.grading import BaselineGrader from llm_connect.models import RunConfig from llm_connect.quality import QualityLedger from llm_connect.routing import AdaptiveRoutingPolicy, RoutingPolicy from llm_connect.shadowing import ShadowingAdapter from .workflow import AssistedGenerationRequest, AssistedGenerationResult # Default identity mapping: every generation stage shipped by the # generic-source profile is its own task type. Callers can override # individual stages via the ``stage_to_task_type`` field — for example to # collapse ``extract-entities`` and ``extract-relations`` into a single # ``extraction`` route, or to widen ``evaluate-entity`` to ``judge``. STAGE_TO_TASK_TYPE_DEFAULT: dict[str, str] = { "summarize-source": "summarize-source", "extract-entities": "extract-entities", "extract-relations": "extract-relations", "evaluate-entity": "evaluate-entity", "synthesize-report": "synthesize-report", } @dataclass(frozen=True) class RoutingAssistedGenerationAdapter: """Route assisted-generation requests through an llm-connect policy. On each ``generate(request)`` call: 1. Resolves ``task_type`` from ``request.stage_id`` (overridable via ``stage_to_task_type``; default falls back to the stage id itself). 2. Asks the policy for an adapter. When the policy is an ``AdaptiveRoutingPolicy`` and ``quality_floor`` is set, the adaptive path is used; otherwise the policy resolves statically. 3. Calls the resolved llm-connect ``LLMAdapter.execute_prompt`` with a ``RunConfig`` built from ``default_run_config``. 4. Maps the ``LLMResponse`` back to an ``AssistedGenerationResult`` and preserves model, usage, finish_reason, and the resolved task_type / adapter_id in ``metadata``. """ policy: RoutingPolicy stage_to_task_type: dict[str, str] = field(default_factory=dict) default_run_config: RunConfig = field(default_factory=RunConfig) quality_floor: float | None = None estimated_cost_per_1k: float | None = None def generate( self, request: AssistedGenerationRequest ) -> AssistedGenerationResult: task_type = self._task_type_for(request.stage_id) adapter = self._resolve(task_type) response = adapter.execute_prompt(request.prompt, self.default_run_config) adapter_id = _identify_adapter(adapter) metadata: dict[str, Any] = { "task_type": task_type, "adapter_id": adapter_id, "model": response.model or self.default_run_config.model_name, "usage": dict(response.usage or {}), "finish_reason": response.finish_reason, } if response.metadata: metadata.update(response.metadata) return AssistedGenerationResult( markdown=response.content, provider=_provider_tag(adapter), metadata=metadata, ) def _resolve(self, task_type: str) -> LLMAdapter: if isinstance(self.policy, AdaptiveRoutingPolicy) and self.quality_floor is not None: return self.policy.resolve( task_type, estimated_cost_per_1k=self.estimated_cost_per_1k, quality_floor=self.quality_floor, ) return self.policy.resolve( task_type, estimated_cost_per_1k=self.estimated_cost_per_1k, ) def _task_type_for(self, stage_id: str) -> str: merged = dict(STAGE_TO_TASK_TYPE_DEFAULT) merged.update(self.stage_to_task_type) return merged.get(stage_id, stage_id) def _identify_adapter(adapter: LLMAdapter) -> str: """Best-effort stable id for an llm-connect adapter instance. Prefers an explicit ``adapter_id`` attribute (some adapters set it), falls back to ``{class_name}:{model_attr}`` when a model attribute is present, otherwise just the class name. """ adapter_id = getattr(adapter, "adapter_id", "") if adapter_id: return str(adapter_id) model = ( getattr(adapter, "model", "") or getattr(adapter, "model_name", "") or getattr(adapter, "_model", "") ) name = type(adapter).__name__ if model: return f"{name}:{model}" return name def wrap_with_shadow_sampling( *, candidate: LLMAdapter, baseline: LLMAdapter, grader: BaselineGrader, ledger: QualityLedger, task_type: str, adapter_id: str | None = None, baseline_adapter_id: str | None = None, shadow_rate: float = 0.1, async_shadow: bool = True, on_shadow_error: Any | None = None, ) -> ShadowingAdapter: """Wrap ``candidate`` with llm-connect's ``ShadowingAdapter``. Sampled baseline grading collects QualityLedger observations without changing the response the caller sees. Errors in the shadow path (baseline outage, grader failure, ledger write error) never alter the candidate response — failures land on ``on_shadow_error`` when provided, else are silently swallowed by the underlying adapter. The returned ``ShadowingAdapter`` is still an ``LLMAdapter``, so it can be slotted into a ``RoutingPolicy`` rule and used through ``RoutingAssistedGenerationAdapter`` without further changes. """ return ShadowingAdapter( candidate_adapter=candidate, baseline_adapter=baseline, grader=grader, ledger=ledger, task_type=task_type, adapter_id=adapter_id or _identify_adapter(candidate), baseline_adapter_id=baseline_adapter_id or _identify_adapter(baseline), shadow_rate=shadow_rate, async_shadow=async_shadow, on_shadow_error=on_shadow_error, ) def summarise_quality_ledger( ledger_path: str | Any, ) -> list[dict[str, Any]]: """Roll up a QualityLedger into one row per (task_type, adapter_id). Useful as a CLI helper or a quick budget-style inspection without loading llm-connect's full ledger API at the call site. """ from pathlib import Path ledger = QualityLedger(path=Path(ledger_path)) observations = ledger.read_all() grouped: dict[tuple[str, str], dict[str, Any]] = {} for obs in observations: key = (obs.task_type, obs.adapter_id) bucket = grouped.setdefault( key, { "task_type": obs.task_type, "adapter_id": obs.adapter_id, "observations": 0, "mean_quality": 0.0, "mean_cost_usd": 0.0, "total_tokens_in": 0, "total_tokens_out": 0, }, ) bucket["observations"] += 1 bucket["mean_quality"] += float(obs.quality_score) bucket["mean_cost_usd"] += float(obs.cost_usd) bucket["total_tokens_in"] += int(getattr(obs, "tokens_in", 0) or 0) bucket["total_tokens_out"] += int(getattr(obs, "tokens_out", 0) or 0) rows: list[dict[str, Any]] = [] for bucket in grouped.values(): count = bucket["observations"] if count: bucket["mean_quality"] = round(bucket["mean_quality"] / count, 4) bucket["mean_cost_usd"] = round(bucket["mean_cost_usd"] / count, 6) rows.append(bucket) rows.sort(key=lambda row: (row["task_type"], row["adapter_id"])) return rows def _provider_tag(adapter: LLMAdapter) -> str: """Coarse provider tag matching the strings already used in run records. Returns ``openrouter`` / ``claude_code`` / ``openai`` / ``gemini`` / ``routing`` so existing tooling (budget rollup buckets, archive metadata) keeps its bucket keys stable. """ name = type(adapter).__name__.lower() if "openrouter" in name: return "openrouter" if "claudecode" in name or "claude_code" in name: return "claude_code" if "openai" in name: return "openai" if "gemini" in name: return "gemini" if "mock" in name or "static" in name: return "mock" return "routing"