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 are pure and side-effect-free: callers decide what to 

9do with the result (swap in-memory ``cfg`` field, surface a banner, log 

10a warning, etc.). The persisted file is never rewritten, so the user's 

11declared intent is preserved across reinstalls. 

12""" 

13 

14from __future__ import annotations 

15 

16import logging 

17from dataclasses import dataclass 

18 

19from lilbee.catalog.query import reclassify_by_name 

20from lilbee.catalog.types import ModelTask 

21from lilbee.core.config import cfg 

22from lilbee.modelhub.model_manager.discovery import ( 

23 classify_remote_models, 

24 discover_api_models, 

25) 

26from lilbee.modelhub.model_manager.types import ValidationResult 

27from lilbee.modelhub.registry import ModelRegistry 

28from lilbee.providers.litellm_sdk import litellm_available 

29from lilbee.providers.local_servers import LocalServerSpec 

30from lilbee.providers.local_servers.config_urls import base_url_for 

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

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

33from lilbee.providers.sdk_backend import PROVIDER_API_KEY_FIELD, provider_has_key 

34 

35log = logging.getLogger(__name__) 

36 

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

38REASON_LITELLM_MISSING = "the litellm extra isn't installed; run pip install 'lilbee[litellm]'" 

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

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

41REASON_NOT_INSTALLED = "it isn't installed" 

42REASON_UNAVAILABLE = "it isn't available" 

43 

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

45_PROBE_TIMEOUT_S = 1.0 

46 

47 

48@dataclass(frozen=True) 

49class CanonicalRef: 

50 """Result of canonicalizing a persisted ref. 

51 

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

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

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

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

56 ``status`` is not ``OK``. 

57 """ 

58 

59 original: str 

60 effective: str 

61 status: ValidationResult 

62 reason: str | None = None 

63 

64 

65def _is_local_installed(ref: str) -> bool: 

66 """True iff ``ref`` resolves to an installed GGUF in the local registry.""" 

67 try: 

68 registry = ModelRegistry(cfg.models_dir) 

69 installed_models = registry.list_installed() 

70 installed = {m.ref for m in installed_models} | {m.hf_repo for m in installed_models} 

71 return ref in installed 

72 except Exception: # pragma: no cover - defensive for fresh installs 

73 log.debug("Local registry probe failed for %r", ref, exc_info=True) 

74 return False 

75 

76 

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

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

79 try: 

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

81 except Exception: # pragma: no cover - defensive; classify swallows its own errors 

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

83 return False 

84 

85 

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

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

88 if not litellm_available(): 

89 return ValidationResult.UNKNOWN, REASON_LITELLM_MISSING 

90 base_url = base_url_for(spec.key) 

91 if not _local_server_reachable(spec, base_url): 

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

93 return ValidationResult.OK, None 

94 

95 

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

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

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

99 if provider in LOCAL_SERVER_KEYS: 

100 spec = local_server_for_key(provider) 

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

102 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

103 return _classify_local_server_ref(spec) 

104 if provider in PROVIDER_API_KEY_FIELD: 

105 if provider_has_key(provider): 

106 return ValidationResult.OK, None 

107 return ValidationResult.NO_KEY, REASON_NO_API_KEY.format(provider=provider) 

108 if not parsed.is_remote: 

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

110 return ValidationResult.NOT_INSTALLED, REASON_NOT_INSTALLED 

111 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

112 

113 

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

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

116 

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

118 the configured model server. Never mutates persisted state. 

119 """ 

120 if not ref: 

121 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

122 if _is_local_installed(ref): 

123 return ValidationResult.OK, None 

124 try: 

125 parsed = parse_model_ref(ref) 

126 except Exception: 

127 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

128 return _classify_uninstalled_ref(parsed) 

129 

130 

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

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

133 status, _reason = _classify_ref(ref) 

134 return status 

135 

136 

137def _first_available_api_chat_ref() -> str | None: 

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

139 try: 

140 groups = discover_api_models() 

141 except Exception: 

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

143 return None 

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

145 if models: 

146 first = models[0] 

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

148 return None 

149 

150 

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

152 """Return the first installed local ref whose task matches *want*. 

153 

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

155 """ 

156 try: 

157 registry = ModelRegistry(cfg.models_dir) 

158 installed = list(registry.list_installed()) 

159 except Exception: 

160 log.debug("Local registry probe failed during canonicalization", exc_info=True) 

161 return None 

162 for manifest in installed: 

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

164 return manifest.ref 

165 return None 

166 

167 

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

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

170 

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

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

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

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

175 """ 

176 status, reason = _classify_ref(original) 

177 if status == ValidationResult.OK: 

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

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

180 if allow_api: 

181 candidates.append(_first_available_api_chat_ref()) 

182 candidates.append(_first_installed_local_ref(want_task)) 

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

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

185 

186 

187def canonicalize_chat_model() -> CanonicalRef: 

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

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

190 

191 

192def canonicalize_embedding_model() -> CanonicalRef: 

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

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