Coverage for src/lilbee/providers/engine_params.py: 100%

1"""Engine-neutral parameter helpers: model-path resolution, context and GPU-layer 

2sizing, and chat-option translation. 

3 

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""" 

8 

9from __future__ import annotations 

10 

11import logging 

12from pathlib import Path 

13from typing import TYPE_CHECKING, Any 

14 

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 

19from lilbee.providers.base import ( 

20 CONTEXT_WINDOW_MARGIN_TOKENS, 

21 GENERATION_RESERVE_TOKENS, 

22 ProviderError, 

23 ProviderErrorKind, 

24 estimate_budget_tokens, 

25 normalize_generation_options, 

26) 

27from lilbee.providers.gguf_meta import read_gguf_metadata, train_ctx_from_meta 

28from lilbee.providers.model_cache import ( 

29 compute_dynamic_ctx, 

30 get_available_memory, 

31 kv_bytes_per_token, 

32) 

33 

34log = logging.getLogger(__name__) 

35 

36EMBED_FALLBACK_CTX = 2048 

37"""Context used for embed/rerank when a GGUF reports junk (e.g. context_length=0).""" 

38 

39# Sized above chunk_size so BOS re-added on re-tokenization doesn't overflow a full-chunk input. 

40_EMBED_CTX_MARGIN = 8 

41 

42 

43def resolve_embed_ctx(meta: dict[str, str] | None, model_path: Path) -> int: 

44 """Embed/rerank context: worst-case chunk tokenization, capped by trained context. 

45 

46 ``chunk_size`` is token-denominated but the chunker enforces a CHARACTER 

47 budget (``chunk_size * CHARS_PER_TOKEN``). A BPE token is at least one 

48 character, so that char budget is also the provable token ceiling for any 

49 chunk the chunker can emit: size the context to it and embed-time 

50 truncation becomes impossible. Token-dense text (numeric tables, dense 

51 identifiers) otherwise reaches ~2x chunk_size tokens against a 1x cap and 

52 silently loses its tail at embed time.""" 

53 from lilbee.data.extract.chunk import CHARS_PER_TOKEN 

54 

55 train_ctx = train_ctx_from_meta(meta, fallback=EMBED_FALLBACK_CTX, model_path=model_path) 

56 return min(train_ctx, cfg.chunk_size * CHARS_PER_TOKEN + _EMBED_CTX_MARGIN) 

57 

58 

59_LLM_RERANK_HEADROOM = 512 

60"""Tokens reserved above chunk_size for an LLM reranker's query, prompt, and 1-token answer.""" 

61 

62 

63def resolve_llm_rerank_ctx(meta: dict[str, str] | None, model_path: Path) -> int: 

64 """LLM-reranker context: a query+candidate pair, capped by the model's trained context.""" 

65 train_ctx = train_ctx_from_meta(meta, fallback=EMBED_FALLBACK_CTX, model_path=model_path) 

66 return min(train_ctx, cfg.chunk_size + _LLM_RERANK_HEADROOM) 

67 

68 

69_VISION_FALLBACK_N_CTX = 4096 

70"""Context for a vision load when the GGUF reports no usable context_length.""" 

71 

72_VISION_PAGE_CTX_CAP = 32768 

73"""Per-page ceiling on a vision OCR server's context: covers a single high-res page's 

74image tokens plus prompt, while keeping a long-context VLM placeable beside a chat giant.""" 

75 

76_N_GPU_LAYERS_AUTO = -1 

77"""llama.cpp's "fit as many layers as the device holds" value for n_gpu_layers. 

78 

79The engine measures free VRAM at load and picks the count itself, spilling the 

80rest to system memory. That is a better answer than any number lilbee can 

81compute ahead of time, because it is taken on the real device after every other 

82tenant, so the planner passes this rather than a layer count of its own. 

83""" 

84# llama.cpp's "offload nothing"; the user's CPU-only opt-out rather than a budget. 

85_N_GPU_LAYERS_NONE = 0 

86 

87 

88def chat_options_to_kwargs(options: dict[str, Any] | None) -> dict[str, Any]: 

89 """Translate user-facing chat options into generation kwargs. 

90 

91 The output keys (``temperature``/``top_p``/``top_k``/``seed``/``max_tokens``/ 

92 ``repeat_penalty``) are accepted by llama-server's OpenAI body. ``top_k`` is 

93 kept (local llama.cpp honors it), unlike the SDK/API translator which drops it. 

94 ``think`` becomes ``chat_template_kwargs.enable_thinking``, which thinking 

95 templates honor and others ignore. 

96 """ 

97 kwargs = normalize_generation_options(options) 

98 think = kwargs.pop("think", None) 

99 if think is not None: 

100 kwargs["chat_template_kwargs"] = {"enable_thinking": think} 

101 return kwargs 

102 

103 

104def resolve_model_path(model: str, registry: ModelRegistry | None = None) -> Path: 

105 """Resolve a model name to a .gguf file path. 

106 

107 Resolution order: (1) registry (canonical source for installed models), 

108 (2) an absolute path to an existing file. Pass *registry* to resolve without 

109 reaching for ``get_services()`` (callers running inside its construction). 

110 """ 

111 if not model: 

112 raise ProviderError( 

113 "No model is configured for this role. Pick one from the catalog " 

114 "or run 'lilbee model pull <model>'.", 

115 provider="llama-server", 

116 kind=ProviderErrorKind.NOT_FOUND, 

117 ) 

118 if registry is None: 

119 # call-time import: keeps the app-layer container off this module's import graph 

120 from lilbee.app.services import get_services 

121 

122 registry = get_services().registry 

123 try: 

124 return registry.resolve(model) 

125 except (KeyError, ValueError): 

126 pass 

127 

128 candidate = Path(model) 

129 if candidate.is_absolute(): 

130 if candidate.exists(): 

131 return candidate 

132 raise ProviderError( 

133 f"Model file not found: {model}", 

134 provider="llama-server", 

135 kind=ProviderErrorKind.NOT_FOUND, 

136 ) 

137 

138 raise ProviderError( 

139 f"Model {model!r} is not installed. Run 'lilbee model pull {model}' to download it.", 

140 provider="llama-server", 

141 kind=ProviderErrorKind.NOT_FOUND, 

142 ) 

143 

144 

145def _kv_elem_bytes_for_cfg() -> int: 

146 """Bytes per KV element implied by the configured cache type.""" 

147 return KV_CACHE_TYPE_BYTES[cfg.kv_cache_type] 

148 

149 

150def chat_ctx_ceiling(meta: dict[str, str] | None, model_path: Path) -> int: 

151 """Hard upper bound on a chat per-slot n_ctx: trained context, capped by ``cfg.num_ctx_max``.""" 

152 training_ctx = train_ctx_from_meta(meta, fallback=DEFAULT_NUM_CTX, model_path=model_path) 

153 if cfg.num_ctx_max is not None: 

154 return min(training_ctx, cfg.num_ctx_max) 

155 return training_ctx 

156 

157 

158def resolve_chat_ctx( 

159 model_path: Path, meta: dict[str, str] | None, *, available_bytes: int | None = None 

160) -> int: 

161 """Pick a single-GPU n_ctx aiming for ``cfg.chat_n_ctx_target``, clamped to model + host. 

162 

163 When ``cfg.num_ctx_max`` is ``None`` the model's training_ctx is the only 

164 ceiling, so a long-context model can grow past the target if the host has 

165 the RAM to back it. A multi-GPU tensor-split chat is sized separately by the 

166 fleet against its per-device headroom (see :func:`lilbee.providers.fleet.ctx.fit_split_ctx`). 

167 ``available_bytes`` overrides the live host-memory read, and every caller 

168 that is sizing a real launch passes it: the fleet and the surfaces that 

169 mirror it hand over 

170 :func:`lilbee.providers.fleet.planning.plan_sizing_budget`, which reports the 

171 memory of the GPU that will run the model and holds a clean-box snapshot so 

172 a reload sizes ctx like the boot did. 

173 """ 

174 training_ctx = train_ctx_from_meta(meta, fallback=DEFAULT_NUM_CTX, model_path=model_path) 

175 ceiling = cfg.num_ctx_max if cfg.num_ctx_max is not None else training_ctx 

176 

177 try: 

178 model_bytes = model_path.stat().st_size 

179 kv_per_tok = kv_bytes_per_token(meta, _kv_elem_bytes_for_cfg()) 

180 if available_bytes is None: 

181 available_bytes = get_available_memory(cfg.gpu_memory_fraction) 

182 return compute_dynamic_ctx( 

183 model_bytes=model_bytes, 

184 available_bytes=available_bytes, 

185 training_ctx=training_ctx, 

186 kv_bytes_per_tok=kv_per_tok, 

187 ceiling=ceiling, 

188 target=cfg.chat_n_ctx_target, 

189 ) 

190 except (OSError, ValueError): 

191 log.debug("dynamic ctx sizing failed for %s, using static cap", model_path, exc_info=True) 

192 return min(training_ctx, cfg.chat_n_ctx_target) 

193 

194 

195# Tokens the minimum grounded prompt allows for the question plus the context 

196# template's framing, beyond the system prompt and one retrieved source. 

197_GROUNDED_QUESTION_TOKENS = 128 

198 

199 

200def min_usable_chat_ctx() -> int: 

201 """Smallest chat window that serves one grounded answer: the system prompt, 

202 one retrieved source, the question, and the generation reserve plus margin.""" 

203 return ( 

204 estimate_budget_tokens(cfg.rag_system_prompt) 

205 + cfg.chunk_size 

206 + _GROUNDED_QUESTION_TOKENS 

207 + GENERATION_RESERVE_TOKENS 

208 + CONTEXT_WINDOW_MARGIN_TOKENS 

209 ) 

210 

211 

212def resolve_n_gpu_layers(*, embedding: bool) -> int: 

213 """Resolve ``cfg.n_gpu_layers`` (None=all) to llama.cpp's offload integer. 

214 

215 Zero is honoured for every role. It is not a layer budget but the way a user 

216 says "run this on the CPU", and the search roles used to take the 

217 full-offload sentinel before the setting was read, so embed, rerank and 

218 vision kept loading onto the GPU that had just been excluded. 

219 

220 Any other value is a chat-shaped budget and says nothing useful about a small 

221 embedding model, which still offloads fully. 

222 """ 

223 if cfg.n_gpu_layers == _N_GPU_LAYERS_NONE: 

224 return _N_GPU_LAYERS_NONE 

225 if embedding or cfg.n_gpu_layers is None: 

226 return _N_GPU_LAYERS_AUTO 

227 return cfg.n_gpu_layers 

228 

229 

230def resolve_vision_ctx(model_path: Path) -> int: 

231 """Pick n_ctx for a vision OCR load: the model's training context, capped per page. 

232 

233 Uses the model's ``<arch>.context_length`` (not the chat-tuned ``cfg.num_ctx``: a 

234 vision pass packs image-token embeddings plus the prompt, and a small chat ctx 

235 truncates OCR output) but caps it at ``_VISION_PAGE_CTX_CAP``. OCR processes one page 

236 per request, so a single page never exceeds the cap, yet a long-context VLM's full 

237 context would otherwise estimate too large to place alongside a chat giant. 

238 """ 

239 try: 

240 meta = read_gguf_metadata(model_path) 

241 except Exception: 

242 log.debug("read_gguf_metadata failed for vision %s", model_path, exc_info=True) 

243 meta = None 

244 train_ctx = train_ctx_from_meta(meta, fallback=_VISION_FALLBACK_N_CTX, model_path=model_path) 

245 return min(train_ctx, _VISION_PAGE_CTX_CAP)