Coverage for src/lilbee/providers/base.py: 100%
115 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +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 pathlib import Path
10from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, overload, runtime_checkable
12from pydantic import BaseModel
14from lilbee.core.health_warnings import HealthWarning
15from lilbee.core.vectors import Vector
16from lilbee.providers.roles import WorkerRole
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
98 def to_dict(self) -> dict[str, Any]:
99 """Return only non-None values as a dict."""
100 return {k: v for k, v in self.model_dump().items() if v is not None}
103def filter_options(options: dict[str, Any]) -> dict[str, Any]:
104 """Validate and filter generation options through LLMOptions model."""
105 return LLMOptions(**options).to_dict()
108def aux_options(num_predict: int, **extra: Any) -> dict[str, Any]:
109 """Options for an internal call: a cap sized for its answer, with thinking off.
111 A thinking model spends a small cap inside <think>, which ``strip_reasoning``
112 then deletes whole, so the call returns nothing.
113 """
114 return {"num_predict": num_predict, "think": False, **extra}
117def normalize_generation_options(options: dict[str, Any] | None) -> dict[str, Any]:
118 """Validate options and map them to the per-call set an OpenAI/llama-server body takes.
120 ``filter_options`` validates against :class:`LLMOptions`; ``num_predict`` then
121 becomes ``max_tokens`` and ``num_ctx`` is dropped (a model-load param, not a
122 per-call one). Shared by the fleet and SDK option translators so the mapping
123 lives in one place.
124 """
125 if not options:
126 return {}
127 filtered = filter_options(options)
128 if "num_predict" in filtered:
129 filtered["max_tokens"] = filtered.pop("num_predict")
130 filtered.pop("num_ctx", None)
131 return filtered
134class ProviderErrorKind(StrEnum):
135 """Provider-agnostic category of a failed provider call.
137 Classified by exception type at each backend boundary so callers can
138 branch on the kind instead of matching message strings (which are
139 provider-specific and drift between SDK versions).
140 """
142 AUTH = "auth"
143 RATE_LIMIT = "rate_limit"
144 CONTEXT_OVERFLOW = "context_overflow"
145 NOT_FOUND = "not_found"
146 BAD_REQUEST = "bad_request"
147 CONNECTION = "connection"
148 SERVER = "server"
149 CAPACITY = "capacity"
150 PORT_CONFLICT = "port_conflict"
151 UNKNOWN = "unknown"
154class ProviderError(Exception):
155 """Raised when an LLM provider operation fails.
157 ``kind`` is the provider-agnostic category; backends that can't classify a
158 failure leave it ``UNKNOWN``.
159 """
161 def __init__(
162 self,
163 message: str,
164 *,
165 provider: str = "",
166 kind: ProviderErrorKind = ProviderErrorKind.UNKNOWN,
167 ) -> None:
168 self.provider = provider
169 self.kind = kind
170 super().__init__(message)
173# Human word per role for the not-configured error ("embedding model", not "embed model").
174_ROLE_WORDS: dict[WorkerRole, str] = {
175 WorkerRole.CHAT: "chat",
176 WorkerRole.EMBED: "embedding",
177 WorkerRole.RERANK: "reranker",
178 WorkerRole.VISION: "vision",
179}
182def require_role_ref(ref: str, role: WorkerRole, *, provider: str = "") -> str:
183 """Reject an unconfigured role with a clean error instead of parsing ''."""
184 if not ref:
185 raise ProviderError(
186 f"No {_ROLE_WORDS[role]} model is configured. Pick one from the catalog "
187 f"or run 'lilbee model pull <model>'.",
188 provider=provider,
189 kind=ProviderErrorKind.NOT_FOUND,
190 )
191 return ref
194ChatMessage = dict[str, str]
197@dataclass(frozen=True)
198class ToolCall:
199 """One tool/function call the model requested.
201 ``arguments`` is the raw JSON-encoded argument object (OpenAI's shape), left
202 as a string so the caller decides how to parse and validate it. ``id`` is the
203 server-assigned call id, echoed back in the tool result message.
204 """
206 id: str
207 name: str
208 arguments: str
211@dataclass(frozen=True)
212class ChatToolResult:
213 """A chat turn that may carry tool calls alongside (or instead of) text.
215 ``tool_calls`` is empty for an ordinary text answer; ``content`` is empty when
216 the model returned only tool calls. Both can be populated when a model emits
217 commentary plus a call.
218 """
220 content: str
221 tool_calls: list[ToolCall]
224class FinishReason(StrEnum):
225 """Why a chat completion stopped, mirroring OpenAI's vocabulary."""
227 STOP = "stop"
228 LENGTH = "length"
229 TOOL_CALLS = "tool_calls"
230 CONTENT_FILTER = "content_filter"
232 @classmethod
233 def coerce(cls, raw: object) -> FinishReason:
234 """Map a backend-supplied finish_reason to a member, defaulting to STOP.
236 Both the streaming and non-streaming paths read finish_reason from the
237 backend; an unknown or non-string value (a model that omits it) falls
238 back to STOP so the dispatch reports an ordinary end-of-turn.
239 """
240 if isinstance(raw, str):
241 try:
242 return cls(raw)
243 except ValueError:
244 return cls.STOP
245 return cls.STOP
248@dataclass(frozen=True)
249class TokenUsage:
250 """Prompt / completion token counts for one chat call.
252 Defaults to zero so a backend that reports no usage block still yields a
253 well-formed result; the fleet populates these from llama-server's ``usage``.
254 """
256 prompt_tokens: int = 0
257 completion_tokens: int = 0
258 # Prompt tokens the engine served from cache, a subset of prompt_tokens.
259 cached_prompt_tokens: int = 0
262@dataclass(frozen=True)
263class ChatResult:
264 """Structured result from a non-streaming chat call.
266 ``tool_calls`` is empty for an ordinary text answer; ``text`` is empty when
267 the model returned only tool calls. ``usage`` carries the backend's token
268 counts (zero when unreported). The canonical chat dispatch reads these to
269 build its OpenAI/Anthropic-shaped response.
270 """
272 text: str
273 tool_calls: tuple[ToolCall, ...]
274 finish_reason: FinishReason
275 usage: TokenUsage = TokenUsage()
278@dataclass(frozen=True)
279class ToolCallDelta:
280 """Partial tool-call delta in a streaming response, accumulated by ``index``.
282 ``id`` and ``name`` arrive on the opener frame for a call; ``arguments_delta``
283 accumulates across subsequent frames at the same ``index``.
284 """
286 index: int
287 id: str | None
288 name: str | None
289 arguments_delta: str | None
292@dataclass(frozen=True)
293class StreamFinish:
294 """Terminal frame carrying why a streaming chat call stopped.
296 Emitted once, near the end of the stream, so the dispatch can report the
297 same finish_reason the non-streaming path already surfaces, notably
298 ``length`` on a max_tokens truncation. Tool-call streams already infer
299 TOOL_USE from their deltas, so a finish frame never downgrades that.
300 """
302 reason: FinishReason
305ChatStreamItem = str | ToolCallDelta | TokenUsage | StreamFinish
306"""One frame yielded by a streaming chat call: text token, tool-call delta, the
307final token-usage summary, or the finish-reason terminator (each emitted once,
308last, when the backend reports them)."""
311class LLMProvider(Protocol):
312 """Protocol for pluggable LLM backends."""
314 def embed(self, texts: list[str]) -> list[Vector]:
315 """Embed a batch of texts, return list of vectors."""
316 ...
318 def count_tokens(self, text: str) -> int:
319 """Exact token count of *text* under the embedding model's tokenizer.
321 Raise ``NotImplementedError`` when the backend has no local tokenizer (cloud
322 SDK backends); token-budgeted chunk sizing then falls back to a character
323 estimate.
324 """
325 ...
327 def count_chat_prompt_tokens(
328 self,
329 messages: list[ChatMessage],
330 *,
331 options: dict[str, Any] | None = None,
332 model: str | None = None,
333 tools: list[dict[str, Any]] | None = None,
334 tool_choice: str | dict[str, Any] | None = None,
335 ) -> int:
336 """Tokens the chat model prefills for the prompt :meth:`chat` would send.
338 Takes the same arguments as :meth:`chat`. The model's chat template is
339 applied first, so role markers and the tool-call preamble are counted.
340 Assumes no warm state: a backend whose chat model is cold loads it, and
341 the first call against a fresh replica may issue the same one-token
342 template probe :meth:`chat` issues, so this can block for a cold start.
343 Raise ``NotImplementedError`` when the backend cannot render or
344 tokenize, and the caller estimates instead.
345 """
346 ...
348 @overload
349 def chat(
350 self,
351 messages: list[ChatMessage],
352 *,
353 stream: Literal[False] = False,
354 options: dict[str, Any] | None = None,
355 model: str | None = None,
356 tools: list[dict[str, Any]] | None = None,
357 tool_choice: str | dict[str, Any] | None = None,
358 ) -> ChatResult: ...
360 @overload
361 def chat(
362 self,
363 messages: list[ChatMessage],
364 *,
365 stream: Literal[True],
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 ) -> ClosableIterator[ChatStreamItem]: ...
372 def chat(
373 self,
374 messages: list[ChatMessage],
375 *,
376 stream: bool = False,
377 options: dict[str, Any] | None = None,
378 model: str | None = None,
379 tools: list[dict[str, Any]] | None = None,
380 tool_choice: str | dict[str, Any] | None = None,
381 ) -> ChatResult | ClosableIterator[ChatStreamItem]:
382 """Chat completion.
384 Non-streaming returns a :class:`ChatResult` (assistant text, any
385 tool-call frames, and a finish reason). Streaming returns a
386 :class:`ClosableIterator` of :data:`ChatStreamItem` (text tokens
387 interleaved with :class:`ToolCallDelta` frames). ``tools`` is the
388 OpenAI function-tool list; ``tool_choice`` is ``"auto"`` / ``"none"`` /
389 ``"required"`` or a ``{"type": "function", ...}`` selector. A model
390 that lacks tool support returns an empty ``tool_calls`` / yields no
391 tool deltas rather than erroring.
392 """
393 ...
395 def supports_tools(self, model_ref: str) -> bool:
396 """Return True iff the backend can route tool calls for *model_ref*.
398 Default False so backends without a tool path are never offered tools;
399 tool-capable backends override this with a real probe.
400 """
401 return False
403 def chat_with_tools(
404 self,
405 messages: list[ChatMessage],
406 *,
407 tools: list[dict[str, Any]],
408 tool_choice: str | dict[str, Any] | None = None,
409 options: dict[str, Any] | None = None,
410 model: str | None = None,
411 ) -> ChatToolResult:
412 """Non-streaming chat that may return tool calls.
414 ``tools`` is the OpenAI function-tool list; ``tool_choice`` is ``"auto"``
415 / ``"none"`` / ``"required"`` or a specific ``{"type": "function", ...}``
416 selector. Backends without tool support raise :class:`ProviderError`.
417 """
418 raise ProviderError("This backend does not support tool calling.")
420 def vision_ocr(
421 self,
422 png_bytes: bytes,
423 model: str,
424 prompt: str = "",
425 *,
426 timeout: float | None = None,
427 ) -> str:
428 """OCR one page image; ``timeout`` seconds, ``None``/``0`` = no cap."""
429 ...
431 def vision_slot_capacity(self) -> int | None:
432 """Fitted concurrent-OCR slots if the vision fleet is running, else None.
434 The ingest fan-out uses this to size itself to the servers' real
435 continuous-batching capacity rather than the requested concurrency,
436 which a memory-constrained card cannot always fit. ``None`` means the
437 capacity isn't known yet (no local vision backend, or the fleet hasn't
438 started); the caller falls back to its own estimate.
439 """
440 ...
442 def list_models(self) -> list[str]:
443 """List available model identifiers."""
444 ...
446 def list_chat_models(self, provider: str) -> list[str]:
447 """List frontier chat models the provider is aware of for *provider*.
449 Returns the unfiltered upstream catalog (whatever litellm
450 exposes for API providers; an empty list for local backends
451 like the llama-server fleet that have no notion of external
452 catalogs).
453 """
454 ...
456 def pull_model(self, model: str, *, on_progress: Callable[..., Any] | None = None) -> None:
457 """Download a model. Raises NotImplementedError if not supported."""
458 ...
460 def show_model(self, model: str) -> dict[str, Any] | None:
461 """Return model metadata, or None if backend doesn't expose it."""
462 ...
464 def get_capabilities(self, model: str) -> list[str]:
465 """Return capability tags (e.g. ``["completion", "vision"]``) for *model*.
467 Returns an empty list when the backend does not support capability
468 reporting or the model is not found.
469 """
470 ...
472 def rerank(self, query: str, candidates: list[str]) -> list[float]:
473 """Score *candidates* for their relevance to *query*, one float per candidate.
475 The backend resolves the reranker model from ``cfg.reranker_model``.
476 Callers MUST check ``cfg.reranker_model`` is non-empty before
477 calling; use :meth:`supports_rerank` for UI-render decisions.
479 Returns: list of floats in input order, higher = more relevant.
480 Empty ``candidates`` returns ``[]``.
481 Raises :class:`ProviderError` when the backend does not support
482 reranking, ``cfg.reranker_model`` is empty, or the model scored no
483 candidate. A backend must raise rather than return uniform scores,
484 which would silently preserve the caller's input order.
485 """
486 ...
488 def supports_rerank(self) -> bool:
489 """Capability probe: can this backend rerank *if* a model is configured?
491 Pure capability check, NOT "a reranker is currently active". An
492 empty ``cfg.reranker_model`` returns ``True`` so the settings UI
493 keeps the picker visible; callers that need to know whether
494 reranking is actually configured must check ``bool(cfg.reranker_model)``
495 separately. ``rerank()`` is the gated path that requires a
496 non-empty value.
497 """
498 return False
500 def shutdown(self) -> None:
501 """Release resources (e.g. background threads). No-op if nothing to clean up."""
502 ...
504 def invalidate_load_cache(self, model_path: Path | None = None) -> None:
505 """Drop loaded-model state; ``None`` evicts all, else only that path. No-op default."""
506 return
508 def drop_loaded_models_async(self) -> None:
509 """Drop all loaded-model state off the caller's thread. No-op default.
511 Like :meth:`invalidate_load_cache` with no path, but the teardown (which
512 stops every server and waits on each process) runs on a background thread
513 so a settings change that touches a role-agnostic load key never blocks
514 the UI / request thread. The next call rebuilds with current cfg.
515 """
516 self.invalidate_load_cache()
518 def warm_up_pool(self) -> None:
519 """Eagerly start the configured role servers so the first call lands warm.
521 Default no-op so providers without managed servers (SDK / routing
522 wrappers) can be passed to ``Services`` unchanged. Implemented by
523 :class:`FleetProvider` to spawn the chat / embed / rerank / vision
524 servers whose model is configured.
525 """
526 return
528 def cancel_inference(self) -> None:
529 """Interrupt any in-flight generation. No-op default.
531 The fleet engine severs its live chat streams (llama-server stops
532 generating when the connection drops); the SDK wrapper has nothing to
533 interrupt here.
534 """
535 return
537 def reload_role(self, role: WorkerRole, *, wait: bool = False) -> None:
538 """Drop and respawn just *role*'s model so it picks up changed cfg.
540 Default no-op for providers without per-role model servers. The fleet
541 respawns only that role's server; other roles and their in-flight work
542 are left untouched. ``wait=True`` blocks until the respawn finishes (for a
543 caller already off the event loop); the default returns immediately.
544 """
545 return
547 def reload_placement(self, *, wait: bool = False) -> None:
548 """Re-plan GPU placement with current cfg, restarting only moved roles.
550 Default no-op for providers without GPU-placed servers. The fleet diffs
551 the fresh plan against the running fleet and respawns only the roles
552 whose placement changed, so an untouched role's loaded model stays
553 resident. ``wait=True`` blocks until the restarted proxies are healthy.
554 """
555 return
557 def role_ready(self, role: WorkerRole) -> bool:
558 """Whether *role* has a healthy server now, without starting one.
560 Default ``True``: providers without managed servers (SDK / routing
561 wrappers) are always reachable. The fleet returns ``False`` while a role
562 is still cold-starting so surfaces can show a warming state.
563 """
564 del role
565 return True
567 def max_concurrent_chats(self) -> int:
568 """Upper bound on simultaneous chat generations this provider can serve.
570 Default ``1``: a single in-process model cannot take concurrent generate
571 calls, so chat is serialized. A server-backed provider that batches (the
572 fleet) overrides this with its slot capacity, so the chat admission gate
573 lets that many run at once instead of one at a time.
574 """
575 return 1
577 def served_chat_ctx(self) -> int | None:
578 """Per-slot context the active chat server runs with, or None if unknown.
580 A client trims its conversation to this so a long agentic session fits
581 the model's actual window instead of overflowing. Default ``None``:
582 providers without a managed context (SDK wrappers) advertise nothing.
583 """
584 return None
586 def served_chat_slots(self) -> int | None:
587 """Batching slots the active chat server runs with, or None if unknown.
589 Unlike :meth:`max_concurrent_chats` (an admission bound that must always
590 yield a number), this reports the granted shape and stays ``None`` until
591 a managed engine is up, so status surfaces can tell "one slot" apart
592 from "no engine yet".
593 """
594 return None
596 def embed_token_cap(self) -> int | None:
597 """Tokens the embedder truncates one input to, or None when it has no fixed cap.
599 The chunker bounds its budget to this so no chunk loses text at embedding
600 time. Default ``None``: a remote embedder advertises no window.
601 """
602 return None
604 def health_warnings(self) -> list[HealthWarning]:
605 """Serving degradations a client should know about. Default: none."""
606 return []
608 def chat_prefill_progress(self) -> tuple[int, int] | None:
609 """``(processed, total)`` prompt tokens of a chat prefill in flight, or None.
611 A large model's first agent turn can spend minutes in prompt processing
612 with no tokens streamed; status surfaces poll this to show the work.
613 Default ``None``: providers without a managed engine report nothing.
614 """
615 return None
617 def warm_pending(self) -> bool:
618 """Whether a warm has been requested and has not finished.
620 True from the moment ``warm_up_pool`` accepts a warm until its background
621 work ends, so a surface can hold before the first phase is stamped. Default
622 ``False``: providers without managed servers never warm.
623 """
624 return False
626 def warm_progress(self) -> WarmProgress | None:
627 """Snapshot of the chat model's cold-load progress, or None when idle.
629 A launcher streams this to render a real progress bar while a large chat
630 model loads. Default ``None``: providers without a managed load (SDK /
631 routing wrappers) expose nothing, so a launcher falls back to a plain
632 spinner. The fleet returns live read / engine-load state.
633 """
634 return None
636 def add_spawn_listener(
637 self,
638 *,
639 on_spawning: Callable[[WorkerRole], None] | None = None,
640 on_spawned: Callable[[WorkerRole], None] | None = None,
641 ) -> None:
642 """Subscribe to server (re)spawn lifecycle events. No-op default.
644 The fleet calls ``on_spawning`` before a role's server starts and
645 ``on_spawned`` once it is healthy, so the TUI can surface cold-start and
646 reload progress. Providers without managed servers ignore it.
647 """
648 return