Coverage for src/lilbee/providers/routing_provider.py: 100%
190 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Routing provider: prefix-based dispatch between the SDK backend and the local engine."""
3from __future__ import annotations
5import contextlib
6import threading
7from collections.abc import Callable
8from typing import TYPE_CHECKING, Any, Literal, overload
10from lilbee.app.services import get_services
11from lilbee.catalog.refs import is_bare_hf_repo
12from lilbee.core.config import cfg
13from lilbee.core.health_warnings import HealthWarning, WarningCode
14from lilbee.core.vectors import Vector
15from lilbee.providers.base import (
16 ChatMessage,
17 ChatResult,
18 ChatStreamItem,
19 ChatToolResult,
20 ClosableIterator,
21 LLMProvider,
22 ProviderError,
23 require_role_ref,
24)
25from lilbee.providers.factory import create_sdk_provider
26from lilbee.providers.model_ref import ProviderModelRef, parse_model_ref, routes_to_native_gguf
27from lilbee.providers.roles import ROLE_REGISTRY, WorkerRole
28from lilbee.runtime.cancellation import CancelSignal
30if TYPE_CHECKING:
31 from lilbee.providers.sdk_llm_provider import SdkLLMProvider
32 from lilbee.providers.warm_progress import WarmProgress
35class RoutingProvider(LLMProvider):
36 """Dispatches calls based on the model ref prefix.
38 ``ollama/``, ``openai/``, ``anthropic/``, ``gemini/`` go to the SDK
39 provider. Other refs (the HuggingFace ``<org>/<repo>/<file>.gguf``
40 shape) go to the local llama-server engine, which resolves them against the native
41 registry. A registry miss surfaces the native ProviderError
42 unchanged, rather than silently falling through to a remote backend.
43 """
45 def __init__(
46 self, *, hold_warm: bool = False, config_refresh: Callable[[], None] | None = None
47 ) -> None:
48 self._local: LLMProvider | None = None
49 self._sdk_provider: SdkLLMProvider | None = None
50 # Carried until the local fleet is lazily built, so a provider created for
51 # an interactive session hands that down to the FleetProvider it composes.
52 self._hold_warm = hold_warm
53 self._config_refresh = config_refresh
54 # Guards both lazy inits so two concurrent first-callers on the shared
55 # daemon don't each build a backend and leak the loser. (Construction is
56 # cheap field init; FleetProvider defers the role-server spawn to first
57 # use and single-flights it internally, so this lock is the singleton
58 # guard, not the spawn guard.)
59 self._init_lock = threading.Lock()
61 def _get_local(self) -> LLMProvider:
62 if self._local is None:
63 # FleetProvider composes the llama-server stack; its role servers
64 # spawn lazily on first use, not here. Double-checked under
65 # _init_lock so only the first concurrent caller builds it.
66 with self._init_lock:
67 if self._local is None:
68 from lilbee.providers.fleet.provider import FleetProvider
70 self._local = FleetProvider(
71 hold_warm=self._hold_warm, config_refresh=self._config_refresh
72 )
73 return self._local
75 def _get_sdk_provider(self) -> SdkLLMProvider:
76 if self._sdk_provider is None:
77 with self._init_lock:
78 if self._sdk_provider is None:
79 self._sdk_provider = create_sdk_provider()
80 return self._sdk_provider
82 def _pick_backend(self, ref: ProviderModelRef) -> LLMProvider:
83 """Pick the backend for *ref* purely by prefix."""
84 if ref.is_remote:
85 return self._get_sdk_provider()
86 return self._get_local()
88 def embed(self, texts: list[str]) -> list[Vector]:
89 ref = parse_model_ref(require_role_ref(cfg.embedding_model, WorkerRole.EMBED))
90 return self._pick_backend(ref).embed(texts)
92 def count_tokens(self, text: str) -> int:
93 ref = parse_model_ref(require_role_ref(cfg.embedding_model, WorkerRole.EMBED))
94 return self._pick_backend(ref).count_tokens(text)
96 def count_chat_prompt_tokens(
97 self,
98 messages: list[ChatMessage],
99 *,
100 options: dict[str, Any] | None = None,
101 model: str | None = None,
102 tools: list[dict[str, Any]] | None = None,
103 tool_choice: str | dict[str, Any] | None = None,
104 ) -> int:
105 """Count on the backend the chat ref routes to, same rules as :meth:`chat`."""
106 ref = parse_model_ref(require_role_ref(model or cfg.chat_model, WorkerRole.CHAT))
107 return self._pick_backend(ref).count_chat_prompt_tokens(
108 messages, options=options, model=model, tools=tools, tool_choice=tool_choice
109 )
111 @overload
112 def chat(
113 self,
114 messages: list[dict[str, str]],
115 *,
116 stream: Literal[False] = False,
117 options: dict[str, Any] | None = None,
118 model: str | None = None,
119 tools: list[dict[str, Any]] | None = None,
120 tool_choice: str | dict[str, Any] | None = None,
121 ) -> ChatResult: ...
123 @overload
124 def chat(
125 self,
126 messages: list[dict[str, str]],
127 *,
128 stream: Literal[True],
129 options: dict[str, Any] | None = None,
130 model: str | None = None,
131 tools: list[dict[str, Any]] | None = None,
132 tool_choice: str | dict[str, Any] | None = None,
133 ) -> ClosableIterator[ChatStreamItem]: ...
135 def chat(
136 self,
137 messages: list[dict[str, str]],
138 *,
139 stream: bool = False,
140 options: dict[str, Any] | None = None,
141 model: str | None = None,
142 tools: list[dict[str, Any]] | None = None,
143 tool_choice: str | dict[str, Any] | None = None,
144 ) -> ChatResult | ClosableIterator[ChatStreamItem]:
145 ref = parse_model_ref(require_role_ref(model or cfg.chat_model, WorkerRole.CHAT))
146 backend = self._pick_backend(ref)
147 # Split on stream so each call resolves to a specific overload; the
148 # base impl signature accepts bool but the @overloads on the LLMProvider
149 # Protocol require Literal narrowing at the boundary.
150 if stream:
151 return backend.chat(
152 messages,
153 stream=True,
154 options=options,
155 model=model,
156 tools=tools,
157 tool_choice=tool_choice,
158 )
159 return backend.chat(
160 messages,
161 stream=False,
162 options=options,
163 model=model,
164 tools=tools,
165 tool_choice=tool_choice,
166 )
168 def supports_tools(self, model_ref: str) -> bool:
169 """Delegate the tool-capability probe to the backend the ref routes to."""
170 resolved = model_ref or cfg.chat_model
171 if not resolved:
172 return False # no model configured: nothing to advertise tools
173 return self._pick_backend(parse_model_ref(resolved)).supports_tools(resolved)
175 def chat_with_tools(
176 self,
177 messages: list[dict[str, str]],
178 *,
179 tools: list[dict[str, Any]],
180 tool_choice: str | dict[str, Any] | None = None,
181 options: dict[str, Any] | None = None,
182 model: str | None = None,
183 ) -> ChatToolResult:
184 """Dispatch a tool-enabled chat turn to the backend the ref routes to."""
185 ref = parse_model_ref(require_role_ref(model or cfg.chat_model, WorkerRole.CHAT))
186 backend = self._pick_backend(ref)
187 return backend.chat_with_tools(
188 messages, tools=tools, tool_choice=tool_choice, options=options, model=model
189 )
191 def vision_ocr(
192 self,
193 png_bytes: bytes,
194 model: str,
195 prompt: str = "",
196 *,
197 timeout: float | None = None,
198 cancel: CancelSignal | None = None,
199 ) -> str:
200 """Dispatch by ``model``'s ref prefix, same rules as :meth:`chat`."""
201 ref = parse_model_ref(model)
202 backend = self._pick_backend(ref)
203 return backend.vision_ocr(png_bytes, model, prompt, timeout=timeout, cancel=cancel)
205 def vision_slot_capacity(self) -> int | None:
206 """Delegate to the local fleet, but never build it just to size the fan-out."""
207 return self._local.vision_slot_capacity() if self._local is not None else None
209 def list_models(self) -> list[str]:
210 """Return the union of native and SDK-visible models.
212 Both halves are wrapped so an unreachable remote backend or a
213 missing native registry does not mask the other.
214 """
215 native: set[str] = set()
216 with contextlib.suppress(Exception):
217 native = set(self._get_local().list_models())
218 sdk = self._get_sdk_provider()
219 if not sdk.available():
220 return sorted(native)
221 try:
222 remote = set(sdk.list_models())
223 except Exception:
224 return sorted(native)
225 return sorted(native | remote)
227 def list_chat_models(self, provider: str) -> list[str]:
228 """Delegate to the SDK backend; the native engine has no catalog."""
229 sdk = self._get_sdk_provider()
230 if not sdk.available():
231 return []
232 return sdk.list_chat_models(provider)
234 def pull_model(self, model: str, *, on_progress: Callable[..., Any] | None = None) -> None:
235 """Pull via the SDK backend if installed, otherwise raise."""
236 sdk = self._get_sdk_provider()
237 if not sdk.available():
238 raise ProviderError(f"Cannot pull model {model!r}: no pull-capable backend available")
239 sdk.pull_model(model, on_progress=on_progress)
241 def show_model(self, model: str) -> dict[str, Any] | None:
242 """Show model info from the backend selected by the ref prefix."""
243 ref = parse_model_ref(model)
244 return self._pick_backend(ref).show_model(model)
246 def get_capabilities(self, model: str) -> list[str]:
247 """Return capability tags from the backend selected by the ref prefix."""
248 ref = parse_model_ref(model)
249 return self._pick_backend(ref).get_capabilities(model)
251 def rerank(self, query: str, candidates: list[str]) -> list[float]:
252 """Dispatch rerank to the backend that owns ``cfg.reranker_model``.
254 Native GGUF refs go to the local engine; hosted refs go through the SDK
255 provider. Raises ``ProviderError`` when ``cfg.reranker_model`` is
256 empty or the selected backend does not support reranking.
257 """
258 if not cfg.reranker_model:
259 raise ProviderError("No reranker configured. Set cfg.reranker_model first.")
260 if _is_native_rerank_ref(cfg.reranker_model):
261 return self._get_local().rerank(query, candidates)
262 sdk = self._get_sdk_provider()
263 if not sdk.supports_rerank():
264 raise ProviderError(
265 f"Cannot rerank with {cfg.reranker_model!r}: "
266 "hosted rerank backend not available. "
267 "Install the 'providers' extra to enable hosted reranking."
268 )
269 return sdk.rerank(query, candidates)
271 def supports_rerank(self) -> bool:
272 """Capability probe: can the routed backend rerank if configured?
274 Pure capability check, NOT "a reranker is currently active". An
275 empty ``cfg.reranker_model`` returns ``True`` so the settings UI
276 keeps the picker visible; callers that need to know whether
277 reranking is actually configured must check ``bool(cfg.reranker_model)``
278 separately. Delegates to the backend that would handle the
279 configured model when one is set.
280 """
281 model = cfg.reranker_model
282 if not model:
283 return True
284 if _is_native_rerank_ref(model):
285 return self._get_local().supports_rerank()
286 return self._get_sdk_provider().supports_rerank()
288 def shutdown(self) -> None:
289 """Shut down sub-providers to release resources."""
290 if self._local is not None:
291 self._local.shutdown()
292 if self._sdk_provider is not None:
293 self._sdk_provider.shutdown()
295 def warm_up_pool(self) -> None:
296 """Forward to the native side; the SDK side has no servers to warm.
298 Lazily constructs the local engine if it isn't already up so
299 eager-start during ``Services`` boot still warms the configured
300 native roles, even when the user hasn't issued a chat call yet.
301 """
302 self._get_local().warm_up_pool()
304 def cancel_inference(self) -> None:
305 """Forward to the native engine; the SDK side has nothing to interrupt."""
306 if self._local is not None:
307 self._local.cancel_inference()
309 def reload_role(self, role: WorkerRole, *, wait: bool = False) -> None:
310 """Forward to the native engine; the SDK side has no per-role servers."""
311 if self._local is not None:
312 self._local.reload_role(role, wait=wait)
314 def replan_fleet(self, *, wait: bool = False) -> None:
315 """Forward to the native engine; the SDK side has no GPU placement."""
316 if self._local is not None:
317 self._local.replan_fleet(wait=wait)
319 def role_ready(self, role: WorkerRole) -> bool:
320 """Whether *role* can serve a request right now.
322 A role whose configured ref routes to the SDK backend needs no local
323 server, so it is always ready; the local fleet's readiness is
324 irrelevant to it. A native ref with no local engine built yet cannot
325 serve a token, so it reports not-ready (without building the engine);
326 health's ``chat_ready`` and the cold-start waits all rely on this
327 being positive readiness, not reachability.
328 """
329 if self._role_routes_remote(role):
330 return True
331 if self._local is None:
332 return False
333 return self._local.role_ready(role)
335 @staticmethod
336 def _role_routes_remote(role: WorkerRole) -> bool:
337 """Whether *role*'s configured model ref dispatches to the SDK backend."""
338 ref = str(getattr(cfg, ROLE_REGISTRY[role].config_field))
339 return bool(ref) and parse_model_ref(ref).is_remote
341 def max_concurrent_chats(self) -> int:
342 """Chat concurrency of the local engine; 1 until one exists."""
343 if self._local is None:
344 return 1
345 return self._local.max_concurrent_chats()
347 def served_chat_ctx(self) -> int | None:
348 """Per-slot chat context of the local engine, or None when none exists."""
349 if self._local is None:
350 return None
351 return self._local.served_chat_ctx()
353 def served_chat_slots(self) -> int | None:
354 """Chat batching slots of the local engine, or None when none exists."""
355 if self._local is None:
356 return None
357 return self._local.served_chat_slots()
359 def _local_embedder(self) -> LLMProvider | None:
360 """The local engine when it serves the configured embedder, else None."""
361 ref = cfg.embedding_model
362 if not ref or parse_model_ref(ref).is_remote:
363 return None
364 return self._get_local()
366 def embed_token_cap(self) -> int | None:
367 """Embed token cap of the local engine; None for a remote or unset embedder."""
368 local = self._local_embedder()
369 return None if local is None else local.embed_token_cap()
371 def health_warnings(self) -> list[HealthWarning]:
372 """Serving degradations of the local engine; the embed warning needs a local embedder."""
373 local = self._local_embedder()
374 if local is not None:
375 return local.health_warnings()
376 if self._local is None:
377 return []
378 return [
379 w
380 for w in self._local.health_warnings()
381 if w.code != WarningCode.EMBED_WINDOW_BELOW_CHUNK
382 ]
384 def chat_prefill_progress(self) -> tuple[int, int] | None:
385 """Chat prefill progress of the local engine, or None when none exists."""
386 if self._local is None:
387 return None
388 return self._local.chat_prefill_progress()
390 def warm_pending(self) -> bool:
391 """Forward to the native side; the SDK side never warms."""
392 return self._get_local().warm_pending()
394 def warm_progress(self) -> WarmProgress | None:
395 """Cold-load progress of the local engine, or None when none exists yet."""
396 if self._local is None:
397 return None
398 return self._local.warm_progress()
400 def add_spawn_listener(
401 self,
402 *,
403 on_spawning: Callable[[WorkerRole], None] | None = None,
404 on_spawned: Callable[[WorkerRole], None] | None = None,
405 ) -> None:
406 """Register on the native engine so its server spawns reach the TUI.
408 Builds the local engine if it isn't up yet so the listener is attached
409 before the first spawn, matching ``warm_up_pool``'s eager construction.
410 """
411 self._get_local().add_spawn_listener(on_spawning=on_spawning, on_spawned=on_spawned)
414def _is_native_rerank_ref(model: str) -> bool:
415 """Return True iff *model* should route to the native llama-server rerank path.
417 Two acceptance paths:
419 1. :func:`routes_to_native_gguf` accepts the ref: a native GGUF shape that
420 no local-server prefix (``ollama/``, ``lm_studio/``) claims, matching
421 :func:`parse_model_ref`'s exemption.
422 2. The bare ``<org>/<repo>`` names a repo with an installed quant.
424 The model's name is deliberately not consulted. Hosted rerankers are
425 usually called rerankers too (``cohere/rerank-english-v3.0``), so matching
426 on the name captures them and starves the SDK backend. The registry answers
427 "is this one of ours" without guessing. Non-GGUF refs without a known SDK
428 prefix still raise downstream through :func:`parse_model_ref`.
430 An empty registry reports nothing installed; a registry that cannot be read
431 raises, and the fault surfaces to the caller.
432 """
433 if not model:
434 return False
435 if routes_to_native_gguf(model):
436 return True
437 if not is_bare_hf_repo(model):
438 return False
439 return get_services().registry.installed_ref_for_repo(model) is not None