llm-connect/llm_connect/shadowing.py

"""Shadow-mode observation adapter for adaptive routing."""

from __future__ import annotations

import asyncio
import random
import threading
from collections.abc import Callable, Mapping
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass, field, replace
from typing import Any

from llm_connect.adapter import LLMAdapter
from llm_connect.grading import BaselineGrader
from llm_connect.models import LLMResponse, RunConfig
from llm_connect.quality import QualityLedger, QualityObservation


def _default_cost_estimator(response: LLMResponse) -> float:
    for key in ("cost_usd", "estimated_cost_usd", "cost"):
        value = response.metadata.get(key)
        if isinstance(value, (int, float)) and value >= 0:
            return float(value)
    return 0.0


class _StaticResponseAdapter(LLMAdapter):
    """Adapter shim that lets a BaselineGrader reuse an existing response."""

    def __init__(self, response: LLMResponse):
        self._response = response

    def execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse:
        return self._response

    def validate_config(self, config: RunConfig) -> bool:
        return True


@dataclass
class ShadowingAdapter(LLMAdapter):
    """Return candidate responses while recording sampled baseline grades.

    Shadow work is best-effort: baseline, grading, or ledger failures are
    reported to ``on_shadow_error`` when provided, but never alter the candidate
    response returned to the caller.
    """

    candidate_adapter: LLMAdapter
    baseline_adapter: LLMAdapter
    grader: BaselineGrader
    ledger: QualityLedger
    task_type: str
    adapter_id: str
    model_id: str | None = None
    baseline_adapter_id: str | None = None
    shadow_rate: float = 1.0
    async_shadow: bool = False
    random_source: random.Random = field(default_factory=random.Random, repr=False)
    cost_estimator: Callable[[LLMResponse], float] = _default_cost_estimator
    tags: Mapping[str, Any] = field(default_factory=dict)
    on_shadow_error: Callable[[Exception], None] | None = None
    _executor: ThreadPoolExecutor | None = field(default=None, init=False, repr=False)
    _futures: list[Future[None]] = field(default_factory=list, init=False, repr=False)
    _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)

    def __post_init__(self) -> None:
        if not str(self.task_type).strip():
            raise ValueError("task_type must be a non-empty string")
        if not str(self.adapter_id).strip():
            raise ValueError("adapter_id must be a non-empty string")
        if not 0 <= self.shadow_rate <= 1:
            raise ValueError("shadow_rate must be between 0 and 1")
        if self.async_shadow:
            self._executor = ThreadPoolExecutor(max_workers=1)

    def execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse:
        response = self.candidate_adapter.execute_prompt(prompt, config)
        if self._should_shadow():
            self._handle_shadow(prompt, config, response)
        return response

    async def async_execute_prompt(self, prompt: str, config: RunConfig) -> LLMResponse:
        response = await self.candidate_adapter.async_execute_prompt(prompt, config)
        if self._should_shadow():
            if self.async_shadow:
                self._schedule_shadow(prompt, config, response)
            else:
                await asyncio.to_thread(self._run_shadow, prompt, config, response)
        return response

    def validate_config(self, config: RunConfig) -> bool:
        return self.candidate_adapter.validate_config(config)

    def flush(self, timeout: float | None = None) -> None:
        """Wait for currently queued async shadow work to finish."""
        with self._lock:
            futures = list(self._futures)
            self._futures.clear()
        for future in futures:
            future.result(timeout=timeout)

    def shutdown(self, wait: bool = True) -> None:
        """Shut down the background shadow executor if one was created."""
        if self._executor is not None:
            self._executor.shutdown(wait=wait)
            self._executor = None

    def _should_shadow(self) -> bool:
        if self.shadow_rate <= 0:
            return False
        if self.shadow_rate >= 1:
            return True
        with self._lock:
            return self.random_source.random() < self.shadow_rate

    def _handle_shadow(
        self,
        prompt: str,
        config: RunConfig,
        candidate_response: LLMResponse,
    ) -> None:
        if self.async_shadow:
            self._schedule_shadow(prompt, config, candidate_response)
        else:
            self._run_shadow(prompt, config, candidate_response)

    def _schedule_shadow(
        self,
        prompt: str,
        config: RunConfig,
        candidate_response: LLMResponse,
    ) -> None:
        if self._executor is None:
            self._executor = ThreadPoolExecutor(max_workers=1)
        future = self._executor.submit(self._run_shadow, prompt, config, candidate_response)
        with self._lock:
            self._futures = [item for item in self._futures if not item.done()]
            self._futures.append(future)

    def _run_shadow(
        self,
        prompt: str,
        config: RunConfig,
        candidate_response: LLMResponse,
    ) -> None:
        try:
            shadow_config = replace(config, budget_tracker=None)
            result = self.grader.grade(
                self.baseline_adapter,
                _StaticResponseAdapter(candidate_response),
                prompt,
                shadow_config,
            )
            self.ledger.append(
                QualityObservation(
                    task_type=self.task_type,
                    adapter_id=self.adapter_id,
                    model_id=self.model_id or candidate_response.model or config.model_name,
                    cost_usd=self.cost_estimator(candidate_response),
                    quality_score=result.quality_score,
                    latency_ms=float(candidate_response.metadata.get("latency_ms", 0.0)),
                    tokens_in=int(candidate_response.usage.get("prompt_tokens", 0)),
                    tokens_out=int(candidate_response.usage.get("completion_tokens", 0)),
                    baseline_adapter_id=self.baseline_adapter_id,
                    tags=dict(self.tags),
                )
            )
        except Exception as exc:
            self._report_shadow_error(exc)

    def _report_shadow_error(self, exc: Exception) -> None:
        if self.on_shadow_error is None:
            return
        try:
            self.on_shadow_error(exc)
        except Exception:
            pass