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