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.types import ModelTask 

20from lilbee.core.config import cfg 

21from lilbee.modelhub.model_manager.discovery import ( 

22 classify_remote_models, 

23 discover_api_models, 

24 reclassify_by_name, 

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 = {m.ref for m in registry.list_installed()} | { 

70 m.hf_repo for m in registry.list_installed() 

71 } 

72 return ref in installed 

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

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

75 return False 

76 

77 

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

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

80 try: 

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

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

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

84 return False 

85 

86 

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

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

89 if not litellm_available(): 

90 return ValidationResult.UNKNOWN, REASON_LITELLM_MISSING 

91 base_url = base_url_for(spec.key) 

92 if not _local_server_reachable(spec, base_url): 

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

94 return ValidationResult.OK, None 

95 

96 

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

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

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

100 if provider in LOCAL_SERVER_KEYS: 

101 spec = local_server_for_key(provider) 

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

103 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

104 return _classify_local_server_ref(spec) 

105 if provider in PROVIDER_API_KEY_FIELD: 

106 if provider_has_key(provider): 

107 return ValidationResult.OK, None 

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

109 if not parsed.is_remote: 

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

111 return ValidationResult.NOT_INSTALLED, REASON_NOT_INSTALLED 

112 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

113 

114 

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

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

117 

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

119 the configured model server. Never mutates persisted state. 

120 """ 

121 if not ref: 

122 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

123 if _is_local_installed(ref): 

124 return ValidationResult.OK, None 

125 try: 

126 parsed = parse_model_ref(ref) 

127 except Exception: 

128 return ValidationResult.UNKNOWN, REASON_UNAVAILABLE 

129 return _classify_uninstalled_ref(parsed) 

130 

131 

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

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

134 status, _reason = _classify_ref(ref) 

135 return status 

136 

137 

138def _first_available_api_chat_ref() -> str | None: 

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

140 try: 

141 groups = discover_api_models() 

142 except Exception: 

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

144 return None 

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

146 if models: 

147 first = models[0] 

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

149 return None 

150 

151 

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

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

154 

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

156 """ 

157 try: 

158 registry = ModelRegistry(cfg.models_dir) 

159 installed = list(registry.list_installed()) 

160 except Exception: 

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

162 return None 

163 for manifest in installed: 

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

165 return manifest.ref 

166 return None 

167 

168 

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

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

171 

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

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

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

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

176 """ 

177 status, reason = _classify_ref(original) 

178 if status == ValidationResult.OK: 

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

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

181 if allow_api: 

182 candidates.append(_first_available_api_chat_ref()) 

183 candidates.append(_first_installed_local_ref(want_task)) 

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

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

186 

187 

188def canonicalize_chat_model() -> CanonicalRef: 

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

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

191 

192 

193def canonicalize_embedding_model() -> CanonicalRef: 

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

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