Coverage for src/lilbee/providers/base.py: 100%
111 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-04 17:08 +0000
1"""Base protocol and exceptions for LLM providers."""
3from __future__ import annotations
5from collections.abc import Callable, Iterator
6from dataclasses import dataclass
7from enum import StrEnum
8from pathlib import Path
9from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, overload, runtime_checkable
11from pydantic import BaseModel
13from lilbee.core.vectors import Vector
14from lilbee.providers.roles import WorkerRole
16if TYPE_CHECKING:
17 from lilbee.providers.warm_progress import WarmProgress
19T_co = TypeVar("T_co", covariant=True)
21# The inline reasoning markers lilbee's pipeline speaks. A provider whose server
22# extracts reasoning into a separate field re-inlines it with these tags at the
23# client boundary, so every downstream consumer parses one format.
24# What every provider holds back from the context window for the answer it is
25# about to generate, plus a small margin for the chat template's own framing.
26# One owner for both numbers: the fleet ENFORCES this default budget (rejecting a
27# prompt that exceeds it), while retrieval FITS its context to it. An explicit
28# over-large output reservation (num_predict past the default) is clamped back to
29# this default rather than rejected, so an agent that over-reserves still fits.
30# When the two disagreed, retrieval assembled prompts up to the margin larger than
31# the engine would accept, and a grounded turn failed with a 400 the caller could
32# do nothing about.
33GENERATION_RESERVE_TOKENS = 1024
34CONTEXT_WINDOW_MARGIN_TOKENS = 128
36# Chars-per-token assumed when BUDGETING context, deliberately harsher than the
37# display estimator's 4: dense OCR/legal text tokenizes at ~2.5-3 chars per token.
38BUDGET_CHARS_PER_TOKEN = 3
41def prompt_token_budget(ctx: int, num_predict: int | None = None) -> int:
42 """Tokens a prompt may occupy in a *ctx*-token window, reserve and margin removed."""
43 return ctx - (num_predict or GENERATION_RESERVE_TOKENS) - CONTEXT_WINDOW_MARGIN_TOKENS
46def estimate_budget_tokens(text: str) -> int:
47 """Conservative token cost of *text* for budgeting (see BUDGET_CHARS_PER_TOKEN)."""
48 return max(1, len(text) // BUDGET_CHARS_PER_TOKEN)
51THINK_OPEN_TAG = "<think>"
52THINK_CLOSE_TAG = "</think>"
55@runtime_checkable
56class ClosableIterator(Iterator[T_co], Protocol[T_co]):
57 """An iterator that releases resources when ``close()`` is called.
59 Streaming chat responses use this to guarantee upstream resources (the
60 fleet's in-flight request slot) are released even when callers truncate
61 the stream before exhaustion. Generators satisfy this implicitly.
62 """
64 def close(self) -> None: ...
67class LLMOptions(BaseModel):
68 """Validated options passed to LLM providers.
69 Only these fields are forwarded: everything else is rejected
70 to prevent injection of sensitive parameters like api_base or api_key.
71 """
73 temperature: float | None = None
74 top_p: float | None = None
75 top_k: int | None = None
76 seed: int | None = None
77 num_predict: int | None = None
78 repeat_penalty: float | None = None
79 frequency_penalty: float | None = None
80 presence_penalty: float | None = None
81 num_ctx: int | None = None
82 stop: list[str] | None = None
83 # Thinking-template control for structured internal calls (schema
84 # induction and similar): a small thinking model can burn its whole
85 # token budget inside <think> and emit nothing. llama-server maps this
86 # to chat_template_kwargs and Ollama to its think field; hosted-API translators drop it.
87 think: bool | None = None
88 # Structured-output request for the internal calls that want a bare JSON
89 # value back. Must be listed here or the allowlist above silently drops it.
90 response_format: dict[str, Any] | None = None
92 def to_dict(self) -> dict[str, Any]:
93 """Return only non-None values as a dict."""
94 return {k: v for k, v in self.model_dump().items() if v is not None}
97def filter_options(options: dict[str, Any]) -> dict[str, Any]:
98 """Validate and filter generation options through LLMOptions model."""
99 return LLMOptions(**options).to_dict()
102def aux_options(num_predict: int, **extra: Any) -> dict[str, Any]:
103 """Options for an internal call: a cap sized for its answer, with thinking off.
105 A thinking model spends a small cap inside <think>, which ``strip_reasoning``
106 then deletes whole, so the call returns nothing.
107 """
108 return {"num_predict": num_predict, "think": False, **extra}
111def normalize_generation_options(options: dict[str, Any] | None) -> dict[str, Any]:
112 """Validate options and map them to the per-call set an OpenAI/llama-server body takes.
114 ``filter_options`` validates against :class:`LLMOptions`; ``num_predict`` then
115 becomes ``max_tokens`` and ``num_ctx`` is dropped (a model-load param, not a
116 per-call one). Shared by the fleet and SDK option translators so the mapping
117 lives in one place.
118 """
119 if not options:
120 return {}
121 filtered = filter_options(options)
122 if "num_predict" in filtered:
123 filtered["max_tokens"] = filtered.pop("num_predict")
124 filtered.pop("num_ctx", None)
125 return filtered
128class ProviderErrorKind(StrEnum):
129 """Provider-agnostic category of a failed provider call.
131 Classified by exception type at each backend boundary so callers can
132 branch on the kind instead of matching message strings (which are
133 provider-specific and drift between SDK versions).
134 """
136 AUTH = "auth"
137 RATE_LIMIT = "rate_limit"
138 CONTEXT_OVERFLOW = "context_overflow"
139 NOT_FOUND = "not_found"
140 BAD_REQUEST = "bad_request"
141 CONNECTION = "connection"
142 SERVER = "server"
143 CAPACITY = "capacity"
144 PORT_CONFLICT = "port_conflict"
145 UNKNOWN = "unknown"
148class ProviderError(Exception):
149 """Raised when an LLM provider operation fails.
151 ``kind`` is the provider-agnostic category; backends that can't classify a
152 failure leave it ``UNKNOWN``.
153 """
155 def __init__(
156 self,
157 message: str,
158 *,
159 provider: str = "",
160 kind: ProviderErrorKind = ProviderErrorKind.UNKNOWN,
161 ) -> None:
162 self.provider = provider
163 self.kind = kind
164 super().__init__(message)
167# Human word per role for the not-configured error ("embedding model", not "embed model").
168_ROLE_WORDS: dict[WorkerRole, str] = {
169 WorkerRole.CHAT: "chat",
170 WorkerRole.EMBED: "embedding",
171 WorkerRole.RERANK: "reranker",
172 WorkerRole.VISION: "vision",
173}
176def require_role_ref(ref: str, role: WorkerRole, *, provider: str = "") -> str:
177 """Reject an unconfigured role with a clean error instead of parsing ''."""
178 if not ref:
179 raise ProviderError(
180 f"No {_ROLE_WORDS[role]} model is configured. Pick one from the catalog "
181 f"or run 'lilbee model pull <model>'.",
182 provider=provider,
183 kind=ProviderErrorKind.NOT_FOUND,
184 )
185 return ref
188ChatMessage = dict[str, str]
191@dataclass(frozen=True)
192class ToolCall:
193 """One tool/function call the model requested.
195 ``arguments`` is the raw JSON-encoded argument object (OpenAI's shape), left
196 as a string so the caller decides how to parse and validate it. ``id`` is the
197 server-assigned call id, echoed back in the tool result message.
198 """
200 id: str
201 name: str
202 arguments: str
205@dataclass(frozen=True)
206class ChatToolResult:
207 """A chat turn that may carry tool calls alongside (or instead of) text.
209 ``tool_calls`` is empty for an ordinary text answer; ``content`` is empty when
210 the model returned only tool calls. Both can be populated when a model emits
211 commentary plus a call.
212 """
214 content: str
215 tool_calls: list[ToolCall]
218class FinishReason(StrEnum):
219 """Why a chat completion stopped, mirroring OpenAI's vocabulary."""
221 STOP = "stop"
222 LENGTH = "length"
223 TOOL_CALLS = "tool_calls"
224 CONTENT_FILTER = "content_filter"
226 @classmethod
227 def coerce(cls, raw: object) -> FinishReason:
228 """Map a backend-supplied finish_reason to a member, defaulting to STOP.
230 Both the streaming and non-streaming paths read finish_reason from the
231 backend; an unknown or non-string value (a model that omits it) falls
232 back to STOP so the dispatch reports an ordinary end-of-turn.
233 """
234 if isinstance(raw, str):
235 try:
236 return cls(raw)
237 except ValueError:
238 return cls.STOP
239 return cls.STOP
242@dataclass(frozen=True)
243class TokenUsage:
244 """Prompt / completion token counts for one chat call.
246 Defaults to zero so a backend that reports no usage block still yields a
247 well-formed result; the fleet populates these from llama-server's ``usage``.
248 """
250 prompt_tokens: int = 0
251 completion_tokens: int = 0
254@dataclass(frozen=True)
255class ChatResult:
256 """Structured result from a non-streaming chat call.
258 ``tool_calls`` is empty for an ordinary text answer; ``text`` is empty when
259 the model returned only tool calls. ``usage`` carries the backend's token
260 counts (zero when unreported). The canonical chat dispatch reads these to
261 build its OpenAI/Anthropic-shaped response.
262 """
264 text: str
265 tool_calls: tuple[ToolCall, ...]
266 finish_reason: FinishReason
267 usage: TokenUsage = TokenUsage()
270@dataclass(frozen=True)
271class ToolCallDelta:
272 """Partial tool-call delta in a streaming response, accumulated by ``index``.
274 ``id`` and ``name`` arrive on the opener frame for a call; ``arguments_delta``
275 accumulates across subsequent frames at the same ``index``.
276 """
278 index: int
279 id: str | None
280 name: str | None
281 arguments_delta: str | None
284@dataclass(frozen=True)
285class StreamFinish:
286 """Terminal frame carrying why a streaming chat call stopped.
288 Emitted once, near the end of the stream, so the dispatch can report the
289 same finish_reason the non-streaming path already surfaces, notably
290 ``length`` on a max_tokens truncation. Tool-call streams already infer
291 TOOL_USE from their deltas, so a finish frame never downgrades that.
292 """
294 reason: FinishReason
297ChatStreamItem = str | ToolCallDelta | TokenUsage | StreamFinish
298"""One frame yielded by a streaming chat call: text token, tool-call delta, the
299final token-usage summary, or the finish-reason terminator (each emitted once,
300last, when the backend reports them)."""
303class LLMProvider(Protocol):
304 """Protocol for pluggable LLM backends."""
306 def embed(self, texts: list[str]) -> list[Vector]:
307 """Embed a batch of texts, return list of vectors."""
308 ...
310 def count_tokens(self, text: str) -> int:
311 """Exact token count of *text* under the embedding model's tokenizer.
313 Raise ``NotImplementedError`` when the backend has no local tokenizer (cloud
314 SDK backends); token-budgeted chunk sizing then falls back to a character
315 estimate.
316 """
317 ...
319 @overload
320 def chat(
321 self,
322 messages: list[ChatMessage],
323 *,
324 stream: Literal[False] = False,
325 options: dict[str, Any] | None = None,
326 model: str | None = None,
327 tools: list[dict[str, Any]] | None = None,
328 tool_choice: str | dict[str, Any] | None = None,
329 ) -> ChatResult: ...
331 @overload
332 def chat(
333 self,
334 messages: list[ChatMessage],
335 *,
336 stream: Literal[True],
337 options: dict[str, Any] | None = None,
338 model: str | None = None,
339 tools: list[dict[str, Any]] | None = None,
340 tool_choice: str | dict[str, Any] | None = None,
341 ) -> ClosableIterator[ChatStreamItem]: ...
343 def chat(
344 self,
345 messages: list[ChatMessage],
346 *,
347 stream: bool = False,
348 options: dict[str, Any] | None = None,
349 model: str | None = None,
350 tools: list[dict[str, Any]] | None = None,
351 tool_choice: str | dict[str, Any] | None = None,
352 ) -> ChatResult | ClosableIterator[ChatStreamItem]:
353 """Chat completion.
355 Non-streaming returns a :class:`ChatResult` (assistant text, any
356 tool-call frames, and a finish reason). Streaming returns a
357 :class:`ClosableIterator` of :data:`ChatStreamItem` (text tokens
358 interleaved with :class:`ToolCallDelta` frames). ``tools`` is the
359 OpenAI function-tool list; ``tool_choice`` is ``"auto"`` / ``"none"`` /
360 ``"required"`` or a ``{"type": "function", ...}`` selector. A model
361 that lacks tool support returns an empty ``tool_calls`` / yields no
362 tool deltas rather than erroring.
363 """
364 ...
366 def supports_tools(self, model_ref: str) -> bool:
367 """Return True iff the backend can route tool calls for *model_ref*.
369 Default False so backends without a tool path are never offered tools;
370 tool-capable backends override this with a real probe.
371 """
372 return False
374 def chat_with_tools(
375 self,
376 messages: list[ChatMessage],
377 *,
378 tools: list[dict[str, Any]],
379 tool_choice: str | dict[str, Any] | None = None,
380 options: dict[str, Any] | None = None,
381 model: str | None = None,
382 ) -> ChatToolResult:
383 """Non-streaming chat that may return tool calls.
385 ``tools`` is the OpenAI function-tool list; ``tool_choice`` is ``"auto"``
386 / ``"none"`` / ``"required"`` or a specific ``{"type": "function", ...}``
387 selector. Backends without tool support raise :class:`ProviderError`.
388 """
389 raise ProviderError("This backend does not support tool calling.")
391 def vision_ocr(
392 self,
393 png_bytes: bytes,
394 model: str,
395 prompt: str = "",
396 *,
397 timeout: float | None = None,
398 ) -> str:
399 """OCR one page image; ``timeout`` seconds, ``None``/``0`` = no cap."""
400 ...
402 def vision_slot_capacity(self) -> int | None:
403 """Fitted concurrent-OCR slots if the vision fleet is running, else None.
405 The ingest fan-out uses this to size itself to the servers' real
406 continuous-batching capacity rather than the requested concurrency,
407 which a memory-constrained card cannot always fit. ``None`` means the
408 capacity isn't known yet (no local vision backend, or the fleet hasn't
409 started); the caller falls back to its own estimate.
410 """
411 ...
413 def list_models(self) -> list[str]:
414 """List available model identifiers."""
415 ...
417 def list_chat_models(self, provider: str) -> list[str]:
418 """List frontier chat models the provider is aware of for *provider*.
420 Returns the unfiltered upstream catalog (whatever litellm
421 exposes for API providers; an empty list for local backends
422 like the llama-server fleet that have no notion of external
423 catalogs).
424 """
425 ...
427 def pull_model(self, model: str, *, on_progress: Callable[..., Any] | None = None) -> None:
428 """Download a model. Raises NotImplementedError if not supported."""
429 ...
431 def show_model(self, model: str) -> dict[str, Any] | None:
432 """Return model metadata, or None if backend doesn't expose it."""
433 ...
435 def get_capabilities(self, model: str) -> list[str]:
436 """Return capability tags (e.g. ``["completion", "vision"]``) for *model*.
438 Returns an empty list when the backend does not support capability
439 reporting or the model is not found.
440 """
441 ...
443 def rerank(self, query: str, candidates: list[str]) -> list[float]:
444 """Score *candidates* for their relevance to *query*, one float per candidate.
446 The backend resolves the reranker model from ``cfg.reranker_model``.
447 Callers MUST check ``cfg.reranker_model`` is non-empty before
448 calling; use :meth:`supports_rerank` for UI-render decisions.
450 Returns: list of floats in input order, higher = more relevant.
451 Empty ``candidates`` returns ``[]``.
452 Raises :class:`ProviderError` when the backend does not support
453 reranking, ``cfg.reranker_model`` is empty, or the model scored no
454 candidate. A backend must raise rather than return uniform scores,
455 which would silently preserve the caller's input order.
456 """
457 ...
459 def supports_rerank(self) -> bool:
460 """Capability probe: can this backend rerank *if* a model is configured?
462 Pure capability check, NOT "a reranker is currently active". An
463 empty ``cfg.reranker_model`` returns ``True`` so the settings UI
464 keeps the picker visible; callers that need to know whether
465 reranking is actually configured must check ``bool(cfg.reranker_model)``
466 separately. ``rerank()`` is the gated path that requires a
467 non-empty value.
468 """
469 return False
471 def shutdown(self) -> None:
472 """Release resources (e.g. background threads). No-op if nothing to clean up."""
473 ...
475 def invalidate_load_cache(self, model_path: Path | None = None) -> None:
476 """Drop loaded-model state; ``None`` evicts all, else only that path. No-op default."""
477 return
479 def drop_loaded_models_async(self) -> None:
480 """Drop all loaded-model state off the caller's thread. No-op default.
482 Like :meth:`invalidate_load_cache` with no path, but the teardown (which
483 stops every server and waits on each process) runs on a background thread
484 so a settings change that touches a role-agnostic load key never blocks
485 the UI / request thread. The next call rebuilds with current cfg.
486 """
487 self.invalidate_load_cache()
489 def warm_up_pool(self) -> None:
490 """Eagerly start the configured role servers so the first call lands warm.
492 Default no-op so providers without managed servers (SDK / routing
493 wrappers) can be passed to ``Services`` unchanged. Implemented by
494 :class:`FleetProvider` to spawn the chat / embed / rerank / vision
495 servers whose model is configured.
496 """
497 return
499 def cancel_inference(self) -> None:
500 """Interrupt any in-flight generation. No-op default.
502 The fleet engine severs its live chat streams (llama-server stops
503 generating when the connection drops); the SDK wrapper has nothing to
504 interrupt here.
505 """
506 return
508 def reload_role(self, role: WorkerRole, *, wait: bool = False) -> None:
509 """Drop and respawn just *role*'s model so it picks up changed cfg.
511 Default no-op for providers without per-role model servers. The fleet
512 respawns only that role's server; other roles and their in-flight work
513 are left untouched. ``wait=True`` blocks until the respawn finishes (for a
514 caller already off the event loop); the default returns immediately.
515 """
516 return
518 def reload_placement(self, *, wait: bool = False) -> None:
519 """Re-plan GPU placement with current cfg, restarting only moved roles.
521 Default no-op for providers without GPU-placed servers. The fleet diffs
522 the fresh plan against the running fleet and respawns only the roles
523 whose placement changed, so an untouched role's loaded model stays
524 resident. ``wait=True`` blocks until the restarted proxies are healthy.
525 """
526 return
528 def role_ready(self, role: WorkerRole) -> bool:
529 """Whether *role* has a healthy server now, without starting one.
531 Default ``True``: providers without managed servers (SDK / routing
532 wrappers) are always reachable. The fleet returns ``False`` while a role
533 is still cold-starting so surfaces can show a warming state.
534 """
535 del role
536 return True
538 def max_concurrent_chats(self) -> int:
539 """Upper bound on simultaneous chat generations this provider can serve.
541 Default ``1``: a single in-process model cannot take concurrent generate
542 calls, so chat is serialized. A server-backed provider that batches (the
543 fleet) overrides this with its slot capacity, so the chat admission gate
544 lets that many run at once instead of one at a time.
545 """
546 return 1
548 def served_chat_ctx(self) -> int | None:
549 """Per-slot context the active chat server runs with, or None if unknown.
551 A client trims its conversation to this so a long agentic session fits
552 the model's actual window instead of overflowing. Default ``None``:
553 providers without a managed context (SDK wrappers) advertise nothing.
554 """
555 return None
557 def served_chat_slots(self) -> int | None:
558 """Batching slots the active chat server runs with, or None if unknown.
560 Unlike :meth:`max_concurrent_chats` (an admission bound that must always
561 yield a number), this reports the granted shape and stays ``None`` until
562 a managed engine is up, so status surfaces can tell "one slot" apart
563 from "no engine yet".
564 """
565 return None
567 def chat_prefill_progress(self) -> tuple[int, int] | None:
568 """``(processed, total)`` prompt tokens of a chat prefill in flight, or None.
570 A large model's first agent turn can spend minutes in prompt processing
571 with no tokens streamed; status surfaces poll this to show the work.
572 Default ``None``: providers without a managed engine report nothing.
573 """
574 return None
576 def warm_pending(self) -> bool:
577 """Whether a warm has been requested and has not finished.
579 True from the moment ``warm_up_pool`` accepts a warm until its background
580 work ends, so a surface can hold before the first phase is stamped. Default
581 ``False``: providers without managed servers never warm.
582 """
583 return False
585 def warm_progress(self) -> WarmProgress | None:
586 """Snapshot of the chat model's cold-load progress, or None when idle.
588 A launcher streams this to render a real progress bar while a large chat
589 model loads. Default ``None``: providers without a managed load (SDK /
590 routing wrappers) expose nothing, so a launcher falls back to a plain
591 spinner. The fleet returns live read / engine-load state.
592 """
593 return None
595 def add_spawn_listener(
596 self,
597 *,
598 on_spawning: Callable[[WorkerRole], None] | None = None,
599 on_spawned: Callable[[WorkerRole], None] | None = None,
600 ) -> None:
601 """Subscribe to server (re)spawn lifecycle events. No-op default.
603 The fleet calls ``on_spawning`` before a role's server starts and
604 ``on_spawned`` once it is healthy, so the TUI can surface cold-start and
605 reload progress. Providers without managed servers ignore it.
606 """
607 return