Coverage for src/lilbee/modelhub/model_manager/discovery.py: 100%
144 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Remote model discovery and task classification."""
3import logging
4import os
5import time
6from collections.abc import Callable
7from functools import lru_cache
8from threading import Lock
10import httpx
12from lilbee.app.services import get_services
13from lilbee.catalog.formatting import agent_model_id
14from lilbee.catalog.query import (
15 EMBEDDING_NAME_PATTERNS,
16 RERANKER_NAME_PATTERNS,
17 VISION_NAME_PATTERNS,
18)
19from lilbee.catalog.types import ModelTask
20from lilbee.core.config.model import cfg
21from lilbee.modelhub.model_manager.types import RemoteModel
22from lilbee.providers.backend_names import BackendName
23from lilbee.providers.local_servers import (
24 LM_STUDIO,
25 OLLAMA,
26 LocalServerSpec,
27 openai_models_url,
28)
29from lilbee.providers.local_servers.config_urls import configured_local_servers
30from lilbee.providers.model_ref import format_remote_ref
31from lilbee.providers.sdk_backend import PROVIDER_KEYS
33log = logging.getLogger(__name__)
35_EMBEDDING_FAMILIES = frozenset({"bert", "nomic-bert", "e5", "bge"})
37_CLASSIFY_DEFAULT_TIMEOUT_S = 5.0
40@lru_cache(maxsize=1)
41def _discovery_client() -> httpx.Client:
42 """One shared client for the local-server discovery probes.
44 ``httpx.get`` builds a fresh ``Client`` per call, and every ``Client``
45 construction creates an SSL context that loads the system CA bundle. The
46 catalog re-runs discovery on tab activations and refreshes, which made
47 ``ssl.create_default_context`` 16% of a real-terminal TUI py-spy session.
48 One client builds it once and reuses connections between probes (same
49 pattern as the engine probes in ``fleet.swap_manager``). Unlike those
50 loopback-only probes, ``trust_env`` stays on: these base URLs are
51 user-configurable and a LAN server may sit behind a proxy.
52 """
53 return httpx.Client()
56def _http_get(url: str, *, timeout: float) -> httpx.Response:
57 """GET via the shared discovery client (module seam; tests stub here)."""
58 return _discovery_client().get(url, timeout=timeout)
61def _classify_remote_task(name: str, family: str) -> ModelTask:
62 """Classify a remote model as rerank, embedding, vision, or chat (in that order).
64 Embedding matches by family tag or name pattern; the name path covers
65 servers like LM Studio that report no family.
66 """
67 name_lower = name.lower()
68 if any(rp in name_lower for rp in RERANKER_NAME_PATTERNS):
69 return ModelTask.RERANK
70 family_lower = family.lower()
71 if any(ef in family_lower for ef in _EMBEDDING_FAMILIES) or any(
72 ep in name_lower for ep in EMBEDDING_NAME_PATTERNS
73 ):
74 return ModelTask.EMBEDDING
75 if any(vp in name_lower for vp in VISION_NAME_PATTERNS):
76 return ModelTask.VISION
77 return ModelTask.CHAT
80def classify_remote_models(
81 base_url: str,
82 spec: LocalServerSpec,
83 *,
84 timeout: float = _CLASSIFY_DEFAULT_TIMEOUT_S,
85) -> list[RemoteModel]:
86 """Discover and classify all models from one local server by task.
88 The strategy and provider label come from *spec* (Ollama ``/api/tags`` vs
89 LM Studio ``/v1/models``), so a server reached at a non-default host is
90 classified correctly. Returns ``[]`` on any error so read-only callers stay
91 responsive when the backend is down.
92 """
93 discover = _DISCOVERY_BY_KEY[spec.key]
94 return discover(base_url, spec.display_name, timeout)
97def classify_all_remote_models(
98 *,
99 timeout: float = _CLASSIFY_DEFAULT_TIMEOUT_S,
100) -> list[RemoteModel]:
101 """Classify models across every configured local server, source-labeled."""
102 result: list[RemoteModel] = []
103 for spec, base_url in configured_local_servers():
104 result.extend(classify_remote_models(base_url, spec, timeout=timeout))
105 return result
108def _discover_via_ollama_tags(
109 base_url: str, provider: BackendName, timeout: float
110) -> list[RemoteModel]:
111 """Classify models from Ollama's ``/api/tags`` using family metadata."""
112 try:
113 resp = _http_get(f"{base_url}/api/tags", timeout=timeout)
114 resp.raise_for_status()
115 raw_models = resp.json().get("models", [])
116 except Exception:
117 log.debug("Failed to classify remote models", exc_info=True)
118 return []
120 result: list[RemoteModel] = []
121 for model in raw_models:
122 name = model.get("name", "")
123 details = model.get("details", {})
124 family = details.get("family", "")
125 param_size = details.get("parameter_size", "")
126 task = _classify_remote_task(name, family)
127 result.append(
128 RemoteModel(
129 name=name,
130 task=task,
131 family=family,
132 parameter_size=param_size,
133 provider=provider,
134 )
135 )
136 return result
139def _discover_via_openai_models(
140 base_url: str, provider: BackendName, timeout: float
141) -> list[RemoteModel]:
142 """Classify models from an OpenAI-compatible ``/v1/models`` endpoint.
144 These servers report only ids (no family), so task detection runs off the
145 name patterns, which LM Studio ids usually carry. Every id is surfaced: LM
146 Studio presents LM Link remote/cloud models here as if local, so the list
147 is intentionally not filtered to locally-downloaded models.
148 """
149 try:
150 resp = _http_get(openai_models_url(base_url), timeout=timeout)
151 resp.raise_for_status()
152 raw_models = resp.json().get("data", [])
153 except Exception:
154 log.debug("Failed to classify remote models", exc_info=True)
155 return []
157 result: list[RemoteModel] = []
158 for model in raw_models:
159 name = model.get("id", "")
160 if not name:
161 continue
162 task = _classify_remote_task(name, "")
163 result.append(
164 RemoteModel(
165 name=name,
166 task=task,
167 family="",
168 parameter_size="",
169 provider=provider,
170 )
171 )
172 return result
175# Listing strategy per local-server routing key. Module-level so it stays a
176# single source of truth as servers are added to the registry.
177_DISCOVERY_BY_KEY: dict[str, Callable[[str, BackendName, float], list[RemoteModel]]] = {
178 OLLAMA.key: _discover_via_ollama_tags,
179 LM_STUDIO.key: _discover_via_openai_models,
180}
183def _has_provider_key(cfg_field: str, env_var: str) -> bool:
184 """Return True if a usable API key exists via env var or lilbee config."""
185 if os.environ.get(env_var):
186 return True
187 return bool(getattr(cfg, cfg_field, ""))
190def discover_api_models() -> dict[str, list[RemoteModel]]:
191 """Return frontier chat models grouped by provider.
193 Returns whatever the active provider's backend exposes for each
194 configured API key, no curation. Short-circuits before touching
195 the SDK when no keys are present.
196 """
197 active = [
198 (prov, cfg_f, env, label)
199 for prov, cfg_f, env, label in PROVIDER_KEYS
200 if _has_provider_key(cfg_f, env)
201 ]
202 if not active:
203 return {}
205 provider = get_services().provider
207 result: dict[str, list[RemoteModel]] = {}
208 for prov, _cfg_field, _env_var, display_name in active:
209 chat_models = [
210 RemoteModel(
211 name=model_name,
212 task=ModelTask.CHAT,
213 family="",
214 parameter_size="",
215 provider=display_name,
216 )
217 for model_name in provider.list_chat_models(prov)
218 ]
219 if chat_models:
220 result[display_name] = chat_models
221 return result
224def detect_remote_embedding_models() -> list[str]:
225 """Return embedding-model names across every configured local server."""
226 return [m.name for m in classify_all_remote_models() if m.task == ModelTask.EMBEDDING]
229def _installed_native_refs() -> set[str]:
230 """Canonical refs from the native registry; empty set if the walk fails."""
231 try:
232 return {m.ref for m in get_services().registry.list_installed()}
233 except Exception:
234 log.warning("Native registry walk failed; contributing no installed refs", exc_info=True)
235 return set()
238def gather_known_model_refs() -> set[str]:
239 """Canonical refs from the native registry, every configured local server, and APIs.
241 Each primitive swallows its own failures, so a backend being down contributes an
242 empty subset rather than raising.
243 """
244 refs = _installed_native_refs()
245 for rm in classify_all_remote_models():
246 refs.add(format_remote_ref(rm.name, rm.provider))
247 for models in discover_api_models().values():
248 for rm in models:
249 refs.add(format_remote_ref(rm.name, rm.provider))
250 return refs
253class KnownModelCache:
254 """TTL-cached union of native + remote + frontier model refs.
256 Not a ``cachetools.TTLCache``: the generation counter is the point. A
257 fan-out already in flight when :meth:`invalidate` runs would otherwise
258 install its pre-pull answer with a full TTL, hiding a freshly pulled model
259 for 30s; bumping the generation makes that late writer publish without
260 renewing the expiry. The fan-out runs off the lock because it hits network.
261 """
263 DEFAULT_TTL_S = 30.0
265 def __init__(self, ttl_s: float = DEFAULT_TTL_S) -> None:
266 self._ttl_s = ttl_s
267 self._refs: frozenset[str] = frozenset()
268 self._expires_at: float = 0.0
269 self._generation: int = 0
270 self._lock = Lock()
272 def refs(self) -> frozenset[str]:
273 """Cached canonical-ref set, refreshing past the TTL (fan-out runs off the lock)."""
274 with self._lock:
275 if time.monotonic() < self._expires_at:
276 return self._refs
277 captured_generation = self._generation
278 fresh = frozenset(gather_known_model_refs())
279 with self._lock:
280 self._refs = fresh
281 if self._generation == captured_generation:
282 self._expires_at = time.monotonic() + self._ttl_s
283 return self._refs
285 def resolve(self, model: str) -> str | None:
286 """Resolve *model* to its canonical ref, or None if unknown.
288 Accepts the canonical ref, an Ollama ``name:tag`` shorthand, and the clean
289 agent-facing id (:func:`agent_model_id`) an agent config pins in place of
290 the full GGUF path. 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 "/" not in model and ":" in model:
298 prefixed = OLLAMA.qualify(model)
299 if prefixed in refs:
300 return prefixed
301 aliased = [ref for ref in refs if agent_model_id(ref) == model]
302 if len(aliased) == 1:
303 return aliased[0]
304 return None
306 def invalidate(self) -> None:
307 """Force the next ``refs()`` call to re-probe."""
308 with self._lock:
309 self._expires_at = 0.0
310 self._generation += 1