Coverage for src/lilbee/modelhub/model_manager/validation.py: 100%

1"""Validate persisted chat/embedding refs against current installation state. 

2 

3Persisted refs in ``~/.lilbee/config.toml`` (or any other config source) 

4become stale when the user removes a GGUF, swaps providers, or moves 

5between machines. The TUI / server / CLI all read these refs at startup 

6and should not get a "model not found" error from the very first prompt. 

7 

8The helpers here decide nothing: callers act on the result (swap the 

9in-memory ``cfg`` field, surface a banner, log a warning, etc.). They 

10read the model registry, which repairs a manifest it can recover from 

11the HuggingFace cache, but they never rewrite the config file, so the 

12user's declared intent is preserved across reinstalls. 

13""" 

14 

15from __future__ import annotations 

16 

17import logging 

18from dataclasses import dataclass 

19 

20from lilbee.catalog.query import reclassify_by_name 

21from lilbee.catalog.types import KeyStatus, ModelTask 

22from lilbee.core.config import cfg 

23from lilbee.modelhub.install_state import InstallState, install_state 

24from lilbee.modelhub.model_manager.discovery import ( 

25 classify_remote_models, 

26 discover_api_models, 

27) 

28from lilbee.modelhub.model_manager.types import ValidationResult 

29from lilbee.modelhub.registry import ModelRegistry 

30from lilbee.providers import litellm_sdk 

31from lilbee.providers.key_check import provider_key_status 

32from lilbee.providers.local_servers import LocalServerSpec 

33from lilbee.providers.local_servers.config_urls import base_url_for 

34from lilbee.providers.local_servers.registry import LOCAL_SERVER_KEYS, local_server_for_key 

35from lilbee.providers.model_ref import ProviderModelRef, format_remote_ref, parse_model_ref 

36from lilbee.providers.sdk_backend import PROVIDER_API_KEY_FIELD 

37 

38log = logging.getLogger(__name__) 

39 

40# User-facing reasons a persisted ref is unusable, shared across surfaces. 

41REASON_LITELLM_MISSING = "the litellm extra isn't installed" 

42REASON_SERVER_UNREACHABLE = "the model server at {base_url} isn't reachable" 

43REASON_NO_API_KEY = "no API key is configured for {provider}" 

44REASON_INVALID_API_KEY = "{provider} rejected the configured API key" 

45REASON_NOT_INSTALLED = "it isn't installed" 

46REASON_UNAVAILABLE = "it isn't available" 

47 

48# Reachability-probe timeout for ollama/lm_studio refs. 

49_PROBE_TIMEOUT_S = 1.0 

50 

51_KEY_STATUS_VALIDATION: dict[KeyStatus, tuple[ValidationResult, str | None]] = { 

52 KeyStatus.READY: (ValidationResult.OK, None), 

53 KeyStatus.MISSING_KEY: (ValidationResult.NO_KEY, REASON_NO_API_KEY), 

54 KeyStatus.INVALID_KEY: (ValidationResult.INVALID_KEY, REASON_INVALID_API_KEY), 

55} 

56 

57 

58@dataclass(frozen=True) 

59class CanonicalRef: 

60 """Result of canonicalizing a persisted ref. 

61 

62 ``effective`` is what callers should use this session. ``original`` 

63 is what the user persisted; if it differs from ``effective`` the 

64 caller should surface the swap. ``reason`` is a human-readable 

65 explanation of why ``original`` was unusable, set whenever 

66 ``status`` is not ``OK``. 

67 """ 

68 

69 original: str 

70 effective: str 

71 status: ValidationResult 

72 reason: str | None = None 

73 

74 

75def _local_server_reachable(spec: LocalServerSpec, base_url: str) -> bool: 

76 """True if the local model server lists at least one model within the probe budget.""" 

77 try: 

78 return bool(classify_remote_models(base_url, spec, timeout=_PROBE_TIMEOUT_S)) 

79 except Exception: # a listing the strategy cannot walk raises past the request guard 

80 log.debug("Local model server probe failed for %r", base_url, exc_info=True) 

81 return False 

82 

83 

84def _classify_local_server_ref(spec: LocalServerSpec) -> tuple[ValidationResult, str | None]: 

85 """Classify an ollama/lm_studio ref: needs the litellm extra and a live server.""" 

86 if not litellm_sdk.litellm_available(): 

87 return ValidationResult.UNKNOWN, REASON_LITELLM_MISSING 

88 base_url = base_url_for(spec.key) 

89 if not _local_server_reachable(spec, base_url): 

90 return ValidationResult.UNKNOWN, REASON_SERVER_UNREACHABLE.format(base_url=base_url) 

91 return ValidationResult.OK, None 

92 

93 

94def _classify_uninstalled_ref(parsed: ProviderModelRef) -> tuple[ValidationResult, str | None]: 

95 """Classify a parsed ref that is not installed locally, by provider kind.""" 

96 provider = (parsed.provider or "").lower() 

97 if provider in LOCAL_SERVER_KEYS: 

98 spec = local_server_for_key(provider) 

99 if spec is None: # pragma: no cover - LOCAL_SERVER_KEYS guarantees a match 

100 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

101 return _classify_local_server_ref(spec) 

102 if provider in PROVIDER_API_KEY_FIELD: 

103 status, reason = _KEY_STATUS_VALIDATION[provider_key_status(provider)] 

104 return status, reason and reason.format(provider=provider) 

105 if not parsed.is_remote: 

106 # A native GGUF ref that no longer resolves to a file on disk. 

107 return ValidationResult.NOT_INSTALLED, REASON_NOT_INSTALLED 

108 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

109 

110 

111def _classify_ref(ref: str) -> tuple[ValidationResult, str | None]: 

112 """Classify a persisted ref, returning its status and a human-readable reason. 

113 

114 Reads cfg, the local registry, and (for ollama/lm_studio refs) probes 

115 the configured model server. Never mutates persisted state. Raises 

116 ``OSError`` when the registry cannot be read: that is unknown, not absent. 

117 """ 

118 if not ref: 

119 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

120 if install_state(ref, ModelRegistry(cfg.models_dir)) is not InstallState.MISSING: 

121 return ValidationResult.OK, None 

122 try: 

123 parsed = parse_model_ref(ref) 

124 except Exception: 

125 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

126 return _classify_uninstalled_ref(parsed) 

127 

128 

129def validate_persisted_model(ref: str) -> ValidationResult: 

130 """Classify a persisted chat/embedding ref against current state.""" 

131 status, _reason = _classify_ref(ref) 

132 return status 

133 

134 

135def _first_available_api_chat_ref() -> str | None: 

136 """Return the first cloud chat ref backed by a configured API key, or ``None``.""" 

137 try: 

138 groups = discover_api_models() 

139 except Exception: 

140 log.debug("discover_api_models failed during canonicalization", exc_info=True) 

141 return None 

142 for _provider, models in groups.items(): 

143 if models: 

144 first = models[0] 

145 return format_remote_ref(first.name, first.provider) 

146 return None 

147 

148 

149def _first_installed_local_ref(want: ModelTask) -> str | None: 

150 """Return the first registered local ref whose task matches *want*. 

151 

152 Registered only, a narrower bar than the persisted ref clears: a 

153 substitute must be a ref every surface can name and the engine can load, 

154 so a loose GGUF file (in no listing) and a manifest whose split set is 

155 missing a shard (the listing gates on the first shard alone) are both out. 

156 Tasks are name-reclassified so the pick matches the role validator. 

157 

158 Raises ``OSError`` when the registry cannot be read, so an unreadable 

159 tree never reads as "nothing to substitute". 

160 """ 

161 registry = ModelRegistry(cfg.models_dir) 

162 for manifest in registry.list_installed(): 

163 if reclassify_by_name(manifest.ref, manifest.task) != want: 

164 continue 

165 if install_state(manifest.ref, registry) is InstallState.REGISTERED: 

166 return manifest.ref 

167 return None 

168 

169 

170def _canonicalize(original: str, *, allow_api: bool, want_task: ModelTask) -> CanonicalRef: 

171 """Resolve a persisted ref to its effective session value. 

172 

173 ``allow_api`` controls the fallback chain: chat allows an API 

174 fallback first; embedding is local-only because most providers 

175 have no embedding equivalent. The local fallback is restricted to 

176 installed models whose task is ``want_task``. 

177 """ 

178 status, reason = _classify_ref(original) 

179 if status == ValidationResult.OK: 

180 return CanonicalRef(original=original, effective=original, status=status) 

181 candidates: list[str | None] = [] 

182 if allow_api: 

183 candidates.append(_first_available_api_chat_ref()) 

184 candidates.append(_first_installed_local_ref(want_task)) 

185 effective = next((c for c in candidates if c), original) 

186 return CanonicalRef(original=original, effective=effective, status=status, reason=reason) 

187 

188 

189def canonicalize_chat_model() -> CanonicalRef: 

190 """Effective chat ref for this session, falling back API -> local -> original.""" 

191 return _canonicalize(cfg.chat_model, allow_api=True, want_task=ModelTask.CHAT) 

192 

193 

194def canonicalize_embedding_model() -> CanonicalRef: 

195 """Effective embedding ref for this session, falling back local -> original.""" 

196 return _canonicalize(cfg.embedding_model, allow_api=False, want_task=ModelTask.EMBEDDING)