Reference

class AsyncBroker:
    """Façade over the LLM pool: route completions, inspect state, edit the catalog."""

    def __init__(  # noqa: PLR0913
        self,
        registry: RegistryProtocol | str | Path | None = None,
        *,
        secrets: SecretsProtocol | None = None,
        store: StoreProtocol | None = None,
        optimize: bool | Optimizer = True,
        scope: str | None = None,
    ) -> None:
        if scope == "":
            raise ValueError("scope must not be empty string; use None for unscoped")
        if registry is None:
            raise ValueError("AsyncBroker requires a `registry` source")

        source_secrets: SecretsProtocol | None = None
        source_store: StoreProtocol | None = None
        if isinstance(registry, (str, Path)):
            registry, source_secrets, source_store = resolve_source(registry)

        secrets = as_secrets(secrets) if secrets is not None else (source_secrets or Secrets())
        store = store if store is not None else (source_store or _default_store(registry))

        if isinstance(optimize, Optimizer):
            self._optimizer: Optimizer | None = optimize
        elif optimize:
            self._optimizer = Optimizer()
        else:
            self._optimizer = None

        self._registry = registry
        self._secrets = secrets
        self._base_store = store
        self._scope = scope

        pool = LLMPool(optimizer=self._optimizer)
        self._pool = pool
        self._catalog = Catalog(registry, secrets, pool, scope=scope)

        self._learning_hook: _LearningHook | None = None
        effective_store: StoreProtocol
        if self._optimizer is not None:
            self._learning_hook = _LearningHook(
                self._optimizer,
                store,
                pool,
                self._catalog.resync,
            )
            effective_store = self._learning_hook
        else:
            effective_store = store

        self._store = effective_store
        self._router = Router(pool, effective_store, scope=scope, optimizer=self._optimizer)
        self._pool_view = PoolView(pool, effective_store)

        self._provisioned = False
        self._provision_lock = asyncio.Lock()
        self._last_underprov_alert: float = float("-inf")
        self._underprov_alert_interval: float = 60.0

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def ensure_pool(self) -> None:
        """Lazy idempotent initializer — provisions the pool exactly once.

        Raises if the registry is empty — call ``sync(preset)`` first.
        """
        if self._provisioned:
            return
        async with self._provision_lock:
            if self._provisioned:
                return
            await self._catalog.provision()
            if self._learning_hook is not None:
                # warm start — provision() above already resynced the registry
                await self._learning_hook.maybe_rebuild(force=True, resync_registry=False)
            self._provisioned = True

    async def sync(self, preset: RegistryProtocol | str | Path) -> None:
        """Mirror ``preset`` into the registry: add new entries, update existing
        ones, delete entries absent from the preset. Explicit and idempotent —
        call it once to initialize a fresh DB, or again whenever the preset changes.
        If the pool is already provisioned, the change takes effect immediately.
        """
        if isinstance(preset, (str, Path)):
            preset = Registry(preset)
        await self._catalog.sync(preset)
        if isinstance(self._base_store, DisabledMapProtocol):
            configs = await self._registry.load()
            await self._base_store.seed_disabled([c.name for c in configs])
        if self._provisioned:
            await self._catalog.resync()

    async def aclose(self) -> None:
        await self._router.aclose()
        for port in (self._registry, self._secrets, self._store):
            if isinstance(port, AsyncResourceProtocol):
                await port.aclose()

    async def __aenter__(self) -> "AsyncBroker":
        await self.ensure_pool()
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.aclose()

    # ------------------------------------------------------------------
    # Routing
    # ------------------------------------------------------------------

    async def ask(
        self,
        prompt: str,
        *,
        operation: str | None = None,
        trace_id: str | None = None,
        wait: float | None = None,
    ) -> AsyncResult:
        await self.ensure_pool()
        try:
            return await self._router.ask(prompt, operation=operation, trace_id=trace_id, wait=wait)
        except NoLLMAvailableError as exc:
            self._maybe_alert_underprov(exc)
            raise

    async def chat(
        self,
        messages: list[dict],
        *,
        tools: list[dict] | None = None,
        operation: str | None = None,
        trace_id: str | None = None,
        wait: float | None = None,
    ) -> AsyncResult:
        await self.ensure_pool()
        try:
            return await self._router.chat(
                messages,
                tools=tools,
                operation=operation,
                trace_id=trace_id,
                wait=wait,
            )
        except NoLLMAvailableError as exc:
            self._maybe_alert_underprov(exc)
            raise

    # ------------------------------------------------------------------
    # Inspection
    # ------------------------------------------------------------------

    async def get(self, name: str) -> AsyncLLM:
        await self.ensure_pool()
        return self._pool_view.get(name)

    async def count(self) -> int:
        await self.ensure_pool()
        return self._pool_view.count()

    async def snapshot(self) -> Mapping[str, LLMSnapshot]:
        await self.ensure_pool()
        return await self._pool_view.snapshot()

    # ------------------------------------------------------------------
    # Manual disable — the one verdict that actually excludes
    # ------------------------------------------------------------------

    async def disable_llm(self, name: str) -> None:
        """Set the manual latch: withdraws the slot, survives preset rolls, covers
        every operation including future ones. Only ``enable_llm`` clears it."""
        await self.ensure_pool()
        self._pool.set_disabled(name)
        if isinstance(self._base_store, DisabledMapProtocol):
            await self._base_store.set_disabled(name, True)

    async def enable_llm(self, name: str) -> None:
        """Clear the manual latch — a re-enabled model rehabilitates through new
        ratings, no quality reset exists."""
        await self.ensure_pool()
        await self._pool.clear_disabled(name)
        if isinstance(self._base_store, DisabledMapProtocol):
            await self._base_store.set_disabled(name, False)

    # ------------------------------------------------------------------
    # Call journal
    # ------------------------------------------------------------------

    async def calls(self, *, limit: int) -> list[Call]:
        return await self._require_queryable().calls(limit=limit, scope=self._scope)

    def _maybe_alert_underprov(self, exc: NoLLMAvailableError) -> None:
        """Fire when zero keyed configs are routable — the genuine "no usable models" alarm.

        A keyless config is never enqueued/acquired/cooled (see the partial-key framing
        in architecture.md), so it must be excluded here: with even one keyless config
        present, an unfiltered check could never observe "all non-AVAILABLE", masking
        the real alarm even when every *keyed* config is COOLING. Only a ``"timeout"``
        reason means the pool is merely temporarily exhausted; every other reason
        (no keys, all disabled, empty pool) already logs its own actionable line.
        """
        if exc.reason != "timeout":
            return
        if self._optimizer is None:
            return
        if not self._pool.configs:
            return
        now = time.monotonic()
        if now - self._last_underprov_alert < self._underprov_alert_interval:
            return
        keyed_names = [name for name in self._pool.configs if self._pool.has_key(name)]
        all_offline = all(
            self._pool.state(name).phase is not LifecyclePhase.AVAILABLE for name in keyed_names
        )
        if all_offline:
            self._last_underprov_alert = now
            logger.warning(
                "pool under-provisioned: all LLMs are COOLING — add more LLMs to the registry",
            )

    def _require_queryable(self) -> QueryableStoreProtocol:
        if not isinstance(self._base_store, QueryableStoreProtocol):
            raise TypeError(
                "this store backend is not queryable — use a queryable backend"
                " (e.g. llmbroker.sqlite.Store) for calls()",
            )
        return cast(QueryableStoreProtocol, self._store)

class AsyncLLM:
    """Handle returned by ``AsyncBroker.get(name)`` — live view into broker internals."""

    def __init__(
        self,
        name: str,
        config: LLMConfig,
        pool: LLMPool,
        store: StoreProtocol,
    ) -> None:
        self._name = name
        self._config = config
        self._pool = pool
        self._store = store

    @property
    def config(self) -> LLMConfig:
        return self._config

    @property
    def disabled(self) -> bool:
        return self._pool.is_disabled(self._name)

    async def state(self) -> LLMState:
        return self._pool.state(self._name)

    async def metrics(self) -> LLMMetrics:
        all_metrics = await resolve_metrics_map(self._store)
        return all_metrics.get(self._name, LLMMetrics(0, None, None))

class AsyncResult:
    """Returned by AsyncBroker.ask()/chat()."""

    def __init__(  # noqa: PLR0913
        self,
        *,
        text: str,
        tool_calls: list[dict] | None,
        usage: Usage | None,
        call_id: str,
        llm_name: str,
        operation: str | None = None,
        store: StoreProtocol,
    ) -> None:
        self.text = text
        self.tool_calls = tool_calls
        self.usage = usage
        self._call_id = call_id
        self._llm_name = llm_name
        self._operation = operation
        self._store = store

    async def record_quality(self, score: float) -> None:
        await self._store.record_quality(
            self._llm_name,
            self._operation,
            score,
            call_id=self._call_id,
        )

class Broker:
    """Synchronous client over an AsyncBroker on a background loop thread."""

    def __init__(  # noqa: PLR0913
        self,
        registry: RegistryProtocol | str | Path | None = None,
        *,
        secrets: SecretsProtocol | None = None,
        store: StoreProtocol | None = None,
        optimize: bool | Optimizer = True,
        scope: str | None = None,
    ) -> None:
        self._async = AsyncBroker(
            registry,
            secrets=secrets,
            store=store,
            optimize=optimize,
            scope=scope,
        )
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(
            target=_run_loop,
            args=(self._loop,),
            daemon=True,
            name="llmbroker-loop",
        )
        self._thread.start()
        # Backstop: if the caller never closes the Broker, stop the loop and
        # join the thread when the instance is garbage-collected. The callback
        # holds only loop + thread (never self), so it does not pin the Broker.
        self._finalizer = weakref.finalize(self, _shutdown, self._loop, self._thread)

    def _run(self, coro: Coroutine[Any, Any, Any]) -> Any:
        return asyncio.run_coroutine_threadsafe(coro, self._loop).result()

    def _ensure_pool(self) -> None:
        self._run(self._async.ensure_pool())

    # ── Accessors ──
    def get(self, name: str) -> LLM:
        return LLM(self._run, self._run(self._async.get(name)))

    def count(self) -> int:
        return self._run(self._async.count())

    # ── calls ──
    def ask(
        self,
        prompt: str,
        *,
        operation: str | None = None,
        trace_id: str | None = None,
        wait: float | None = None,
    ) -> Result:
        return Result(
            self._run,
            self._run(self._async.ask(prompt, operation=operation, trace_id=trace_id, wait=wait)),
        )

    def chat(
        self,
        messages: list[dict],
        *,
        tools: list[dict] | None = None,
        operation: str | None = None,
        trace_id: str | None = None,
        wait: float | None = None,
    ) -> Result:
        return Result(
            self._run,
            self._run(
                self._async.chat(
                    messages,
                    tools=tools,
                    operation=operation,
                    trace_id=trace_id,
                    wait=wait,
                ),
            ),
        )

    def snapshot(self) -> Mapping[str, LLMSnapshot]:
        return self._run(self._async.snapshot())

    def sync(self, preset: RegistryProtocol | str | Path) -> None:
        self._run(self._async.sync(preset))

    def disable_llm(self, name: str) -> None:
        self._run(self._async.disable_llm(name))

    def enable_llm(self, name: str) -> None:
        self._run(self._async.enable_llm(name))

    def calls(self, *, limit: int) -> list[Call]:
        return self._run(self._async.calls(limit=limit))

    # ── lifecycle ──
    def close(self) -> None:
        if not self._finalizer.alive:
            return
        self._run(self._async.aclose())
        # Run the same teardown the GC backstop would, and mark it done so the
        # finalizer does not repeat it later.
        self._finalizer()

    def __enter__(self) -> "Broker":
        self._ensure_pool()
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()

class DictSecrets:
    """Read-only secrets resolver backed by an in-memory mapping (tests / preloaded keys)."""

    def __init__(self, mapping: dict[str, str]) -> None:
        self._mapping = dict(mapping)

    async def resolve(self, ref: str) -> str:
        if ref not in self._mapping:
            raise KeyError(f"DictSecrets: ref {ref!r} not found")
        return self._mapping[ref]

class FileStore:
    """Day-split JSONL call journal plus a YAML disabled-verdict map, under one directory."""

    def __init__(self, directory: str | Path, *, retention: timedelta = _DEFAULT_RETENTION) -> None:
        self._dir = Path(directory)
        self._calls_dir = self._dir / "calls"
        self._disabled_path = self._dir / "disabled.yml"
        self._retention = retention
        self._last_purge = float("-inf")

    def _day_path(self, ts: datetime) -> Path:
        return self._calls_dir / f"{ts.date().isoformat()}.jsonl"

    def _append(self, call: Call) -> None:
        ts = call.ts or datetime.now(UTC)
        path = self._day_path(ts)
        path.parent.mkdir(parents=True, exist_ok=True)
        line = json.dumps(_call_to_jsonable(call))
        with path.open("a", encoding="utf-8") as fh:
            fh.write(line + "\n")

    async def record(self, call: Call) -> None:
        await asyncio.to_thread(self._append, call)
        await self._maybe_purge()

    async def record_quality(
        self,
        llm_name: str,
        operation: str | None,
        score: float,
        *,
        call_id: str | None = None,
    ) -> None:
        await self.record(_new_quality_call(llm_name, operation, score, call_id))

    def _day_files_newest_first(self) -> list[Path]:
        if not self._calls_dir.exists():
            return []
        return sorted(self._calls_dir.glob("*.jsonl"), reverse=True)

    def _read_tail(self, limit: int, scope: str | None) -> list[Call]:
        result: list[Call] = []
        for path in self._day_files_newest_first():
            lines = path.read_text(encoding="utf-8").splitlines()
            for raw_line in reversed(lines):
                stripped = raw_line.strip()
                if not stripped:
                    continue
                call = _call_from_jsonable(json.loads(stripped))
                if scope is not None and call.scope != scope:
                    continue
                result.append(call)
                if len(result) >= limit:
                    return result
        return result

    async def calls(self, *, limit: int, scope: str | None = None) -> list[Call]:
        """Newest-first tail of the journal, both kinds interleaved — unfiltered by scope
        (learning is global); ``scope`` is accepted for the host-facing filter only."""
        return await asyncio.to_thread(self._read_tail, limit, scope)

    def _purge_old_day_files(self) -> None:
        cutoff = (datetime.now(UTC) - self._retention).date()
        for path in self._day_files_newest_first():
            try:
                file_date = date.fromisoformat(path.stem)
            except ValueError:
                continue
            if file_date < cutoff:
                path.unlink(missing_ok=True)

    async def _maybe_purge(self) -> None:
        now = time.monotonic()
        if now - self._last_purge < _PURGE_INTERVAL_SECONDS:
            return
        self._last_purge = now
        await asyncio.to_thread(self._purge_old_day_files)

    def _read_disabled(self) -> dict[str, bool]:
        if not self._disabled_path.exists():
            return {}
        data = yaml.safe_load(self._disabled_path.read_text(encoding="utf-8"))
        return dict(data) if data else {}

    def _write_disabled(self, data: dict[str, bool]) -> None:
        self._disabled_path.parent.mkdir(parents=True, exist_ok=True)
        body = yaml.safe_dump(data, sort_keys=True)
        self._disabled_path.write_text(_DISABLED_HEADER + body, encoding="utf-8")

    async def get_disabled(self, name: str) -> bool:
        data = await asyncio.to_thread(self._read_disabled)
        return bool(data.get(name, False))

    async def set_disabled(self, name: str, flag: bool) -> None:  # noqa: FBT001
        data = await asyncio.to_thread(self._read_disabled)
        data[name] = flag
        await asyncio.to_thread(self._write_disabled, data)

    async def seed_disabled(self, names: list[str]) -> None:
        data = await asyncio.to_thread(self._read_disabled)
        changed = False
        for name in names:
            if name not in data:
                data[name] = False
                changed = True
        if changed:
            await asyncio.to_thread(self._write_disabled, data)

    async def disabled_map(self) -> dict[str, bool]:
        return await asyncio.to_thread(self._read_disabled)

class InMemoryStore:
    """Explicit in-memory opt-out — no persistence, session-scoped learning;
    disabled verdicts live only in process memory."""

    def __init__(self) -> None:
        self._disabled: dict[str, bool] = {}

    async def record(self, _call: Call) -> None:
        return

    async def record_quality(
        self,
        _llm_name: str,
        _operation: str | None,
        _score: float,
        *,
        call_id: str | None = None,  # noqa: ARG002
    ) -> None:
        return

    async def get_disabled(self, name: str) -> bool:
        return self._disabled.get(name, False)

    async def set_disabled(self, name: str, flag: bool) -> None:  # noqa: FBT001
        self._disabled[name] = flag

    async def seed_disabled(self, names: list[str]) -> None:
        for name in names:
            self._disabled.setdefault(name, False)

    async def disabled_map(self) -> dict[str, bool]:
        return dict(self._disabled)

class LLM:
    """Synchronous analogue of AsyncLLM."""

    def __init__(self, run_fn: "Callable[[Any], Any]", async_llm: AsyncLLM) -> None:
        self._run = run_fn
        self._async = async_llm

    @property
    def config(self) -> LLMConfig:
        return self._async.config

    @property
    def disabled(self) -> bool:
        return self._async.disabled

    def state(self) -> LLMState:
        return self._run(self._async.state())

    def metrics(self) -> LLMMetrics:
        return self._run(self._async.metrics())

class LLMRequestError(Exception):
    """Base: this request could not be completed."""

class LifecyclePhase(Enum):
    """The FSM label for one LLM's lifecycle, always derived from cooldown_until vs now."""

    AVAILABLE = "available"
    COOLING = "cooling"

class NoLLMAvailableError(LLMRequestError):
    """No LLM slot was available for this request.

    ``reason`` is one of:

    - ``"empty_pool"`` — the pool has zero slots.
    - ``"no_keys"`` — slots exist but none has a resolved key.
    - ``"all_disabled"`` — keyed slots exist but every one is admin-disabled.
    - ``"excluded"`` — every candidate was excluded for this request (internal).
    - ``"timeout"`` — the deadline expired (or nothing is free right now);
      ``retry_at`` carries the earliest known return time, when known.
    """

    def __init__(
        self,
        message: str,
        *,
        reason: str,
        retry_at: datetime | None = None,
    ) -> None:
        super().__init__(message)
        self.reason = reason
        self.retry_at = retry_at

@dataclass
class Optimizer:
    """Consecutive-failure counter for backoff, plus per-(model, operation) sliding
    windows of raw quality ratings backing the demoted-last selection order."""

    max_delay: float = 3600.0
    backoff_factor: float = 2.0
    quality_floor: float = 0.3
    quality_confidence: float = 0.95  # z for the Wilson upper bound
    quality_window: int = 30  # ratings kept per (model, operation)
    quality_min_count: int = 10  # verdicts need at least this many

    _rl_fail_count: dict[str, int] = field(default_factory=dict, init=False, repr=False)
    _scores: dict[tuple[str, str | None], deque] = field(
        default_factory=dict,
        init=False,
        repr=False,
    )

    def rl_fail_count(self, llm_name: str) -> int:
        return self._rl_fail_count.get(llm_name, 0)

    def on_rate_limited(self, llm_name: str) -> None:
        """Increment the consecutive-failure count the router reads for its backoff exponent."""
        self._rl_fail_count[llm_name] = self._rl_fail_count.get(llm_name, 0) + 1

    def on_success(self, llm_name: str) -> None:
        self._rl_fail_count[llm_name] = 0

    # ------------------------------------------------------------------
    # Quality windows + derived per-operation demotion verdicts
    # ------------------------------------------------------------------

    def wilson_bound(self, llm_name: str, operation: str | None) -> float | None:
        """The Wilson-score upper bound backing ``is_demoted``, for diagnostics."""
        window = self._scores.get((llm_name, operation))
        if not window:
            return None
        return wilson_upper(list(window), _z_score(self.quality_confidence))

    def is_demoted(self, llm_name: str, operation: str | None) -> bool:
        """True iff the window holds at least ``quality_min_count`` ratings and their
        Wilson-score upper bound sits below ``quality_floor``."""
        window = self._scores.get((llm_name, operation))
        if window is None or len(window) < self.quality_min_count:
            return False
        bound = wilson_upper(list(window), _z_score(self.quality_confidence))
        return bound < self.quality_floor

    def demoted_operations(self, llm_name: str) -> frozenset[str | None]:
        return frozenset(
            operation
            for name, operation in self._scores
            if name == llm_name and self.is_demoted(name, operation)
        )

    def record_quality(self, llm_name: str, operation: str | None, score: float) -> None:
        """Fold one rating into the (model, operation) window, oldest evicted first."""
        before = self.is_demoted(llm_name, operation)
        window = self._scores.setdefault(
            (llm_name, operation),
            deque(maxlen=self.quality_window),
        )
        window.append(score)
        self._log_flip(llm_name, operation, before, self.is_demoted(llm_name, operation))

    def load_scores(self, scores: dict[tuple[str, str | None], list[float]]) -> None:
        """Replace every window wholesale — used by the journal rebuild."""
        keys = set(self._scores) | set(scores)
        before = {key: self.is_demoted(*key) for key in keys}
        self._scores = {
            key: deque(values[-self.quality_window :], maxlen=self.quality_window)
            for key, values in scores.items()
        }
        for key in keys:
            self._log_flip(key[0], key[1], before[key], self.is_demoted(*key))

    def _log_flip(
        self,
        llm_name: str,
        operation: str | None,
        before: bool,  # noqa: FBT001
        after: bool,  # noqa: FBT001
    ) -> None:
        if before == after:
            return
        bound = self.wilson_bound(llm_name, operation)
        if after:
            if bound is not None:
                logger.warning(
                    "%s: quality-demoted for operation=%r (wilson upper %.3f < floor %.2f)",
                    llm_name,
                    operation,
                    bound,
                    self.quality_floor,
                )
            else:
                logger.warning("%s: quality-demoted for operation=%r", llm_name, operation)
        elif bound is not None:
            logger.info(
                "%s: quality demotion cleared for operation=%r (wilson upper %.3f)",
                llm_name,
                operation,
                bound,
            )
        else:
            logger.info("%s: quality demotion cleared for operation=%r", llm_name, operation)

class Registry:
    """File-backed read-only registry — ``.toml`` / ``.json`` by extension."""

    def __init__(self, path: str | Path) -> None:
        self._path = Path(path)

    @property
    def path(self) -> Path:
        return self._path

    async def load(self) -> list[LLMConfig]:
        data = _read_data(self._path)
        result: list[LLMConfig] = []
        for entry in data.get("llms", []):
            cfg = _config_from_entry(entry)
            if cfg is not None:
                result.append(cfg)
        return result

    async def key_info(self) -> dict[str, KeyInfo]:
        """Per-provider onboarding metadata from the ``[keys]`` table, keyed by ``api_key_ref``."""
        raw = _read_data(self._path).get("keys", {})
        if not isinstance(raw, dict):
            return {}
        return {str(ref): key_info_from_entry(str(ref), val) for ref, val in raw.items()}

class Result:
    """Synchronous analogue of AsyncResult."""

    def __init__(self, run_fn: "Callable[[Any], Any]", async_result: AsyncResult) -> None:
        self._run = run_fn
        self._async = async_result
        self.text = async_result.text
        self.tool_calls = async_result.tool_calls
        self.usage = async_result.usage

    def record_quality(self, score: float) -> None:
        self._run(self._async.record_quality(score))

class Secrets:
    """Read-only env-backed secrets resolver (the default battery)."""

    async def resolve(self, ref: str) -> str:
        value = os.environ.get(ref)
        if value is None:
            raise KeyError(f"Secrets: env var {ref!r} is not set")
        return value