Coverage for src/lilbee/modelhub/model_manager/discovery.py: 100%
138 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Remote model discovery and task classification."""
3import logging
4import time
5from collections.abc import Callable
6from functools import lru_cache
7from threading import Lock
9import httpx
11from lilbee.app.services import get_services
12from lilbee.catalog.formatting import agent_model_id
13from lilbee.catalog.query import (
14 EMBEDDING_NAME_PATTERNS,
15 RERANKER_NAME_PATTERNS,
16 VISION_NAME_PATTERNS,
17)
18from lilbee.catalog.types import KeyStatus, ModelTask
19from lilbee.modelhub.install_state import InstallState, install_state
20from lilbee.modelhub.model_manager.types import ApiModelGroup, RemoteModel
21from lilbee.providers.backend_names import BackendName
22from lilbee.providers.base import LLMProvider
23from lilbee.providers.key_check import provider_key_set, provider_key_statuses
24from lilbee.providers.local_servers import (
25 LM_STUDIO,
26 OLLAMA,
27 LocalServerSpec,
28 openai_models_url,
29)
30from lilbee.providers.local_servers.config_urls import configured_local_servers
31from lilbee.providers.model_ref import format_remote_ref
32from lilbee.providers.sdk_backend import PROVIDER_KEYS
34log = logging.getLogger(__name__)
36_EMBEDDING_FAMILIES = frozenset({"bert", "nomic-bert", "e5", "bge"})
38_CLASSIFY_DEFAULT_TIMEOUT_S = 5.0
41@lru_cache(maxsize=1)
42def _discovery_client() -> httpx.Client:
43 """One shared client for the local-server discovery probes.
45 ``httpx.get`` builds a fresh ``Client`` per call, and every ``Client``
46 construction creates an SSL context that loads the system CA bundle. The
47 catalog re-runs discovery on tab activations and refreshes, which made
48 ``ssl.create_default_context`` 16% of a real-terminal TUI py-spy session.
49 One client builds it once and reuses connections between probes (same
50 pattern as the engine probes in ``fleet.swap_manager``). Unlike those
51 loopback-only probes, ``trust_env`` stays on: these base URLs are
52 user-configurable and a LAN server may sit behind a proxy.
53 """
54 return httpx.Client()
57def _http_get(url: str, *, timeout: float) -> httpx.Response:
58 """GET via the shared discovery client (module seam; tests stub here)."""
59 return _discovery_client().get(url, timeout=timeout)
62def _classify_remote_task(name: str, family: str) -> ModelTask:
63 """Classify a remote model as rerank, embedding, vision, or chat (in that order).
65 Embedding matches by family tag or name pattern; the name path covers
66 servers like LM Studio that report no family.
67 """
68 name_lower = name.lower()
69 if any(rp in name_lower for rp in RERANKER_NAME_PATTERNS):
70 return ModelTask.RERANK
71 family_lower = family.lower()
72 if any(ef in family_lower for ef in _EMBEDDING_FAMILIES) or any(
73 ep in name_lower for ep in EMBEDDING_NAME_PATTERNS
74 ):
75 return ModelTask.EMBEDDING
76 if any(vp in name_lower for vp in VISION_NAME_PATTERNS):
77 return ModelTask.VISION
78 return ModelTask.CHAT
81def classify_remote_models(
82 base_url: str,
83 spec: LocalServerSpec,
84 *,
85 timeout: float = _CLASSIFY_DEFAULT_TIMEOUT_S,
86) -> list[RemoteModel]:
87 """Discover and classify all models from one local server by task.
89 The strategy and provider label come from *spec* (Ollama ``/api/tags`` vs
90 LM Studio ``/v1/models``), so a server reached at a non-default host is
91 classified correctly. A transport failure or a non-JSON body yields ``[]``
92 so read-only callers stay responsive when the backend is down. A listing
93 the strategy cannot walk raises: the parse runs outside the request guard.
94 """
95 discover = _DISCOVERY_BY_KEY[spec.key]
96 return discover(base_url, spec.display_name, timeout)
99def classify_all_remote_models(
100 *,
101 timeout: float = _CLASSIFY_DEFAULT_TIMEOUT_S,
102) -> list[RemoteModel]:
103 """Classify models across every configured local server, source-labeled."""
104 result: list[RemoteModel] = []
105 for spec, base_url in configured_local_servers():
106 result.extend(classify_remote_models(base_url, spec, timeout=timeout))
107 return result
110def _discover_via_ollama_tags(
111 base_url: str, provider: BackendName, timeout: float
112) -> list[RemoteModel]:
113 """Classify models from Ollama's ``/api/tags`` using family metadata."""
114 try:
115 resp = _http_get(f"{base_url}/api/tags", timeout=timeout)
116 resp.raise_for_status()
117 raw_models = resp.json().get("models", [])
118 except Exception:
119 log.debug("Failed to classify remote models", exc_info=True)
120 return []
122 result: list[RemoteModel] = []
123 for model in raw_models:
124 name = model.get("name", "")
125 details = model.get("details", {})
126 family = details.get("family", "")
127 param_size = details.get("parameter_size", "")
128 task = _classify_remote_task(name, family)
129 result.append(
130 RemoteModel(
131 name=name,
132 task=task,
133 family=family,
134 parameter_size=param_size,
135 provider=provider,
136 )
137 )
138 return result
141def _discover_via_openai_models(
142 base_url: str, provider: BackendName, timeout: float
143) -> list[RemoteModel]:
144 """Classify models from an OpenAI-compatible ``/v1/models`` endpoint.
146 These servers report only ids (no family), so task detection runs off the
147 name patterns, which LM Studio ids usually carry. Every id is surfaced: LM
148 Studio presents LM Link remote/cloud models here as if local, so the list
149 is intentionally not filtered to locally-downloaded models.
150 """
151 try:
152 resp = _http_get(openai_models_url(base_url), timeout=timeout)
153 resp.raise_for_status()
154 raw_models = resp.json().get("data", [])
155 except Exception:
156 log.debug("Failed to classify remote models", exc_info=True)
157 return []
159 result: list[RemoteModel] = []
160 for model in raw_models:
161 name = model.get("id", "")
162 if not name:
163 continue
164 task = _classify_remote_task(name, "")
165 result.append(
166 RemoteModel(
167 name=name,
168 task=task,
169 family="",
170 parameter_size="",
171 provider=provider,
172 )
173 )
174 return result
177# Listing strategy per local-server routing key. Module-level so it stays a
178# single source of truth as servers are added to the registry.
179_DISCOVERY_BY_KEY: dict[str, Callable[[str, BackendName, float], list[RemoteModel]]] = {
180 OLLAMA.key: _discover_via_ollama_tags,
181 LM_STUDIO.key: _discover_via_openai_models,
182}
185def _chat_models_for(provider: LLMProvider, prov: str, display_name: str) -> list[RemoteModel]:
186 """The backend's chat models for hosted provider *prov*, labeled *display_name*."""
187 return [
188 RemoteModel(
189 name=model_name,
190 task=ModelTask.CHAT,
191 family="",
192 parameter_size="",
193 provider=display_name,
194 )
195 for model_name in provider.list_chat_models(prov)
196 ]
199def discover_api_model_groups() -> list[ApiModelGroup]:
200 """Hosted chat models for every provider with a key set, with the key's status.
202 Short-circuits before touching the SDK when no keys are present. Only
203 providers that list models have their key checked.
204 """
205 keyed = [(prov, label) for prov, _cfg, _env, label in PROVIDER_KEYS if provider_key_set(prov)]
206 if not keyed:
207 return []
208 provider = get_services().provider
209 listed = [
210 (prov, label, models)
211 for prov, label in keyed
212 if (models := _chat_models_for(provider, prov, label))
213 ]
214 statuses = provider_key_statuses([prov for prov, _label, _models in listed])
215 return [
216 ApiModelGroup(provider=prov, display_name=label, key_status=statuses[prov], models=models)
217 for prov, label, models in listed
218 ]
221def discover_api_models() -> dict[str, list[RemoteModel]]:
222 """Hosted chat models grouped by provider display name, for accepted keys only."""
223 return {
224 group.display_name: group.models
225 for group in discover_api_model_groups()
226 if group.key_status is KeyStatus.READY
227 }
230def detect_remote_embedding_models() -> list[str]:
231 """Return embedding-model names across every configured local server."""
232 return [m.name for m in classify_all_remote_models() if m.task == ModelTask.EMBEDDING]
235def gather_known_model_refs() -> set[str]:
236 """Canonical refs from the native registry, every configured local server, and APIs.
238 A local server or API that is down contributes an empty subset. An
239 unreadable native registry raises ``OSError`` instead: answering from the
240 remote sources alone would route a request for an installed local model
241 to a hosted provider.
242 """
243 refs = {m.ref for m in get_services().registry.list_installed()}
244 for rm in classify_all_remote_models():
245 refs.add(format_remote_ref(rm.name, rm.provider))
246 for models in discover_api_models().values():
247 for rm in models:
248 refs.add(format_remote_ref(rm.name, rm.provider))
249 return refs
252class KnownModelCache:
253 """TTL-cached union of native + remote + frontier model refs.
255 Not a ``cachetools.TTLCache``: the generation counter is the point. A
256 fan-out already in flight when :meth:`invalidate` runs would otherwise
257 install its pre-pull answer with a full TTL, hiding a freshly pulled model
258 for 30s; bumping the generation makes that late writer publish without
259 renewing the expiry. The fan-out runs off the lock because it hits network.
260 """
262 DEFAULT_TTL_S = 30.0
264 def __init__(self, ttl_s: float = DEFAULT_TTL_S) -> None:
265 self._ttl_s = ttl_s
266 self._refs: frozenset[str] = frozenset()
267 self._expires_at: float = 0.0
268 self._generation: int = 0
269 self._lock = Lock()
271 def refs(self) -> frozenset[str]:
272 """Cached canonical-ref set, refreshing past the TTL (fan-out runs off the lock)."""
273 with self._lock:
274 if time.monotonic() < self._expires_at:
275 return self._refs
276 captured_generation = self._generation
277 fresh = frozenset(gather_known_model_refs())
278 with self._lock:
279 self._refs = fresh
280 if self._generation == captured_generation:
281 self._expires_at = time.monotonic() + self._ttl_s
282 return self._refs
284 def resolve(self, model: str) -> str | None:
285 """Resolve *model* to its canonical ref, or None if unknown.
287 Accepts the canonical ref, an Ollama ``name:tag`` shorthand, and the clean
288 agent-facing id (:func:`agent_model_id`) an agent config pins in place of
289 the full GGUF path, and an absolute path to a GGUF file on disk, which the
290 fleet loads without a manifest. The clean id resolves only when exactly one known ref
291 produces it, so two same-labelled quants stay unresolved rather than
292 routing to the wrong file.
293 """
294 refs = self.refs()
295 if model in refs:
296 return model
297 if install_state(model, get_services().registry) is InstallState.LOOSE_FILE:
298 return model
299 if "/" not in model and ":" in model:
300 prefixed = OLLAMA.qualify(model)
301 if prefixed in refs:
302 return prefixed
303 aliased = [ref for ref in refs if agent_model_id(ref) == model]
304 if len(aliased) == 1:
305 return aliased[0]
306 return None
308 def invalidate(self) -> None:
309 """Force the next ``refs()`` call to re-probe."""
310 with self._lock:
311 self._expires_at = 0.0
312 self._generation += 1