Coverage for src/lilbee/providers/engine_params.py: 100%
111 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Engine-neutral parameter helpers: model-path resolution, context and GPU-layer
2sizing, and chat-option translation.
4These derive launch/generation parameters from cfg + a model's GGUF metadata
5without loading the model, so both the local llama-server engine and any other
6provider compute them identically. No native binding.
7"""
9from __future__ import annotations
11import logging
12from pathlib import Path
13from typing import TYPE_CHECKING, Any
15if TYPE_CHECKING:
16 from lilbee.modelhub.registry import ModelRegistry
17from lilbee.core.config import DEFAULT_NUM_CTX, cfg
18from lilbee.core.config.enums import KV_CACHE_TYPE_BYTES, KvCacheType
19from lilbee.core.health_warnings import HealthWarning, WarningCode
20from lilbee.providers.base import (
21 CONTEXT_WINDOW_MARGIN_TOKENS,
22 GENERATION_RESERVE_TOKENS,
23 ProviderError,
24 ProviderErrorKind,
25 estimate_budget_tokens,
26 normalize_generation_options,
27)
28from lilbee.providers.gguf_meta import read_gguf_metadata, train_ctx_from_meta
29from lilbee.providers.model_cache import (
30 compute_dynamic_ctx,
31 get_available_memory,
32 kv_bytes_per_token,
33)
34from lilbee.providers.model_ref import is_loose_model_file
36log = logging.getLogger(__name__)
38EMBED_FALLBACK_CTX = 2048
39"""Context used for embed/rerank when a GGUF reports junk (e.g. context_length=0)."""
41# Sized above chunk_size so BOS re-added on re-tokenization doesn't overflow a full-chunk input.
42_EMBED_CTX_MARGIN = 8
45def resolve_embed_ctx(meta: dict[str, str] | None, model_path: Path) -> int:
46 """Embed/rerank context: worst-case chunk tokenization, capped by trained context.
48 ``chunk_size`` is token-denominated but the chunker enforces a CHARACTER
49 budget (``chunk_size * CHARS_PER_TOKEN``). A BPE token is at least one
50 character, so that char budget is also the provable token ceiling for any
51 chunk the chunker can emit: size the context to it and embed-time
52 truncation becomes impossible. Token-dense text (numeric tables, dense
53 identifiers) otherwise reaches ~2x chunk_size tokens against a 1x cap and
54 silently loses its tail at embed time. A trained context below that budget
55 caps the window instead, and the chunker then sizes chunks in the embedder's
56 own tokens at the served cap (see :func:`embed_token_cap`)."""
57 from lilbee.data.extract.chunk import CHARS_PER_TOKEN
59 train_ctx = train_ctx_from_meta(meta, fallback=EMBED_FALLBACK_CTX, model_path=model_path)
60 return min(train_ctx, cfg.chunk_size * CHARS_PER_TOKEN + _EMBED_CTX_MARGIN)
63def embed_token_cap(ctx: int) -> int:
64 """Tokens an embed server truncates one input to, given its per-slot context."""
65 return max(1, ctx - _EMBED_CTX_MARGIN)
68def embed_window_warning(cap: int) -> HealthWarning | None:
69 """The embed-window warning for a real, actionable gap, else None.
71 Plain chunking self-corrects to token sizing when *cap* binds, so it loses
72 no text and needs no warning. Semantic chunking ignores token sizing and
73 stays character-sized, so a bound *cap* loses text there regardless of the
74 ``token_sizing`` setting.
75 """
76 from lilbee.data.extract.chunk import CHARS_PER_TOKEN
78 if cfg.semantic_chunking:
79 chars = cfg.chunk_size * CHARS_PER_TOKEN
80 if cap >= chars:
81 return None
82 return HealthWarning(
83 code=WarningCode.EMBED_WINDOW_BELOW_CHUNK,
84 message=(
85 f"The embedding model accepts {cap} tokens per input, below the {chars} "
86 f"characters that chunk_size {cfg.chunk_size} allows. Semantic chunking "
87 f"sizes chunks by characters and ignores token sizing, so a chunk over "
88 f"{cap} tokens loses text at embedding time."
89 ),
90 remedy="Turn off semantic_chunking or lower chunk_size.",
91 )
92 if cfg.token_sizing and cap < cfg.chunk_size:
93 return HealthWarning(
94 code=WarningCode.EMBED_WINDOW_BELOW_CHUNK,
95 message=(
96 f"The embedding model accepts {cap} tokens per input, below the configured "
97 f"chunk_size {cfg.chunk_size} in tokens. Chunks are capped to {cap} tokens "
98 f"instead."
99 ),
100 remedy=f"Set chunk_size to {cap} to match the window.",
101 )
102 return None
105_LLM_RERANK_HEADROOM = 512
106"""Tokens reserved above chunk_size for an LLM reranker's query, prompt, and 1-token answer."""
109def resolve_llm_rerank_ctx(meta: dict[str, str] | None, model_path: Path) -> int:
110 """LLM-reranker context: a query+candidate pair, capped by the model's trained context."""
111 train_ctx = train_ctx_from_meta(meta, fallback=EMBED_FALLBACK_CTX, model_path=model_path)
112 return min(train_ctx, cfg.chunk_size + _LLM_RERANK_HEADROOM)
115_VISION_FALLBACK_N_CTX = 4096
116"""Context for a vision load when the GGUF reports no usable context_length."""
118_VISION_PAGE_CTX_CAP = 32768
119"""Per-page ceiling on a vision OCR server's context: covers a single high-res page's
120image tokens plus prompt, while keeping a long-context VLM placeable beside a chat giant."""
122N_GPU_LAYERS_AUTO = -1
123"""llama.cpp's "fit as many layers as the device holds" value for n_gpu_layers.
125The engine measures free VRAM at load and picks the count itself, spilling the
126rest to system memory. That is a better answer than any number lilbee can
127compute ahead of time, because it is taken on the real device after every other
128tenant, so the planner passes this rather than a layer count of its own.
129"""
130# llama.cpp's "offload nothing"; the user's CPU-only opt-out rather than a budget.
131_N_GPU_LAYERS_NONE = 0
134def chat_options_to_kwargs(options: dict[str, Any] | None) -> dict[str, Any]:
135 """Translate user-facing chat options into generation kwargs.
137 The output keys (``temperature``/``top_p``/``top_k``/``seed``/``max_tokens``/
138 ``repeat_penalty``) are accepted by llama-server's OpenAI body. ``top_k`` is
139 kept (local llama.cpp honors it), unlike the SDK/API translator which drops it.
140 ``think`` becomes ``chat_template_kwargs.enable_thinking``, which thinking
141 templates honor and others ignore.
142 """
143 kwargs = normalize_generation_options(options)
144 think = kwargs.pop("think", None)
145 if think is not None:
146 kwargs["chat_template_kwargs"] = {"enable_thinking": think}
147 return kwargs
150def resolve_model_path(model: str, registry: ModelRegistry | None = None) -> Path:
151 """Resolve a model name to a .gguf file path.
153 Resolution order: (1) registry (canonical source for installed models),
154 (2) an absolute path to an existing file. Pass *registry* to resolve without
155 reaching for ``get_services()`` (callers running inside its construction).
156 """
157 if not model:
158 raise ProviderError(
159 "No model is configured for this role. Pick one from the catalog "
160 "or run 'lilbee model pull <model>'.",
161 provider="llama-server",
162 kind=ProviderErrorKind.NOT_FOUND,
163 )
164 if registry is None:
165 # call-time import: keeps the app-layer container off this module's import graph
166 from lilbee.app.services import get_services
168 registry = get_services().registry
169 try:
170 return registry.resolve(model)
171 except (KeyError, ValueError):
172 pass
174 if is_loose_model_file(model):
175 return Path(model)
176 if Path(model).is_absolute():
177 raise ProviderError(
178 f"Model file not found: {model}",
179 provider="llama-server",
180 kind=ProviderErrorKind.NOT_FOUND,
181 )
183 raise ProviderError(
184 f"Model {model!r} is not installed. Run 'lilbee model pull {model}' to download it.",
185 provider="llama-server",
186 kind=ProviderErrorKind.NOT_FOUND,
187 )
190def chat_kv_elem_bytes() -> float:
191 """Per-element byte cost of the K and V cache a chat launch allocates.
193 Reads the flag the launch passes for both K and V
194 (:func:`lilbee.providers.fleet.planning.chat_cache_type_flag`), so the
195 granted window stays in step with the cache the engine allocates.
196 """
197 # call-time import: planning imports this module at load
198 from lilbee.providers.fleet.planning import chat_cache_type_flag
200 flag = chat_cache_type_flag()
201 return KV_CACHE_TYPE_BYTES[KvCacheType(flag) if flag else KvCacheType.F16]
204def chat_ctx_ceiling(meta: dict[str, str] | None, model_path: Path) -> int:
205 """Hard upper bound on a chat per-slot n_ctx: trained context, capped by ``cfg.num_ctx_max``."""
206 training_ctx = train_ctx_from_meta(meta, fallback=DEFAULT_NUM_CTX, model_path=model_path)
207 if cfg.num_ctx_max is not None:
208 return min(training_ctx, cfg.num_ctx_max)
209 return training_ctx
212def resolve_chat_ctx(
213 model_path: Path, meta: dict[str, str] | None, *, available_bytes: int | None = None
214) -> int:
215 """Pick a single-GPU chat n_ctx aiming for ``cfg.chat_n_ctx_target``, clamped to model + host.
217 A gguf-parser fit answers first
218 (:func:`lilbee.providers.fleet.planning.fit_chat_ctx`), because it prices the
219 cache each layer of this architecture holds. Header math takes over when the
220 estimator cannot answer; it charges every layer as dense attention over the
221 whole window, which under-grants linear-attention, sliding-window and MLA
222 models. Either way the window stops at the smallest of the trained context,
223 ``cfg.num_ctx_max`` and the target.
225 A multi-GPU tensor-split chat is sized separately by the fleet against its
226 per-device headroom (see :func:`lilbee.providers.fleet.ctx.fit_split_ctx`).
227 ``available_bytes`` overrides the live host-memory read, and every caller
228 that is sizing a real launch passes it: the fleet and the surfaces that
229 mirror it hand over
230 :func:`lilbee.providers.fleet.planning.plan_sizing_budget`, which reports the
231 memory of the GPU that will run the model and holds a clean-box snapshot so
232 a reload sizes ctx like the boot did.
233 """
234 # call-time import: planning imports this module at load
235 from lilbee.providers.fleet.planning import fit_chat_ctx
237 training_ctx = train_ctx_from_meta(meta, fallback=DEFAULT_NUM_CTX, model_path=model_path)
238 ceiling = cfg.num_ctx_max if cfg.num_ctx_max is not None else training_ctx
239 if available_bytes is None:
240 available_bytes = get_available_memory(cfg.gpu_memory_fraction)
241 upper = min(training_ctx, ceiling, cfg.chat_n_ctx_target)
243 try:
244 return fit_chat_ctx(model_path, meta, available_bytes=available_bytes, ctx_ceiling=upper)
245 except (ProviderError, OSError, ValueError):
246 log.debug("gguf-parser ctx fit failed for %s, using header math", model_path, exc_info=True)
247 return _header_math_chat_ctx(
248 model_path,
249 meta,
250 available_bytes=available_bytes,
251 training_ctx=training_ctx,
252 ceiling=ceiling,
253 )
256def _header_math_chat_ctx(
257 model_path: Path,
258 meta: dict[str, str] | None,
259 *,
260 available_bytes: int,
261 training_ctx: int,
262 ceiling: int,
263) -> int:
264 """Window from GGUF header arithmetic: weights plus a dense-attention cache."""
265 try:
266 model_bytes = model_path.stat().st_size
267 elem_bytes = chat_kv_elem_bytes()
268 kv_per_tok = kv_bytes_per_token(meta, elem_bytes, elem_bytes)
269 return compute_dynamic_ctx(
270 model_bytes=model_bytes,
271 available_bytes=available_bytes,
272 training_ctx=training_ctx,
273 kv_bytes_per_tok=kv_per_tok,
274 ceiling=ceiling,
275 target=cfg.chat_n_ctx_target,
276 )
277 except (OSError, ValueError):
278 log.debug("dynamic ctx sizing failed for %s, using static cap", model_path, exc_info=True)
279 return min(training_ctx, cfg.chat_n_ctx_target)
282# Tokens the minimum grounded prompt allows for the question plus the context
283# template's framing, beyond the system prompt and one retrieved source.
284_GROUNDED_QUESTION_TOKENS = 128
287def min_usable_chat_ctx() -> int:
288 """Smallest chat window that serves one grounded answer: the system prompt,
289 one retrieved source, the question, and the generation reserve plus margin."""
290 return (
291 estimate_budget_tokens(cfg.rag_system_prompt)
292 + cfg.chunk_size
293 + _GROUNDED_QUESTION_TOKENS
294 + GENERATION_RESERVE_TOKENS
295 + CONTEXT_WINDOW_MARGIN_TOKENS
296 )
299def resolve_n_gpu_layers(*, embedding: bool) -> int:
300 """Resolve ``cfg.n_gpu_layers`` (None=all) to llama.cpp's offload integer.
302 Zero is honoured for every role. It is not a layer budget but the way a user
303 says "run this on the CPU", and the search roles used to take the
304 full-offload sentinel before the setting was read, so embed, rerank and
305 vision kept loading onto the GPU that had just been excluded.
307 Any other value is a chat-shaped budget and says nothing useful about a small
308 embedding model, which still offloads fully.
309 """
310 if cfg.n_gpu_layers == _N_GPU_LAYERS_NONE:
311 return _N_GPU_LAYERS_NONE
312 if embedding or cfg.n_gpu_layers is None:
313 return N_GPU_LAYERS_AUTO
314 return cfg.n_gpu_layers
317def resolve_vision_ctx(model_path: Path) -> int:
318 """Pick n_ctx for a vision OCR load: the model's training context, capped per page.
320 Uses the model's ``<arch>.context_length`` (not the chat-tuned ``cfg.num_ctx``: a
321 vision pass packs image-token embeddings plus the prompt, and a small chat ctx
322 truncates OCR output) but caps it at ``_VISION_PAGE_CTX_CAP``. OCR processes one page
323 per request, so a single page never exceeds the cap, yet a long-context VLM's full
324 context would otherwise estimate too large to place alongside a chat giant.
325 """
326 try:
327 meta = read_gguf_metadata(model_path)
328 except Exception:
329 log.debug("read_gguf_metadata failed for vision %s", model_path, exc_info=True)
330 meta = None
331 train_ctx = train_ctx_from_meta(meta, fallback=_VISION_FALLBACK_N_CTX, model_path=model_path)
332 return min(train_ctx, _VISION_PAGE_CTX_CAP)