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