infospace-bench/src/infospace_bench/routing.py

"""
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"