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.core.config import cfg 

20from lilbee.modelhub.model_manager.discovery import discover_api_models 

21from lilbee.modelhub.model_manager.types import ValidationResult 

22from lilbee.modelhub.registry import ModelRegistry 

23from lilbee.providers.model_ref import format_remote_ref, parse_model_ref 

24from lilbee.providers.sdk_backend import PROVIDER_API_KEY_FIELD, get_provider_api_key 

25 

26log = logging.getLogger(__name__) 

27 

28 

29@dataclass(frozen=True) 

30class CanonicalRef: 

31 """Result of canonicalizing a persisted ref. 

32 

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

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

35 caller should surface the swap. 

36 """ 

37 

38 original: str 

39 effective: str 

40 status: ValidationResult 

41 

42 

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

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

45 try: 

46 registry = ModelRegistry(cfg.models_dir) 

47 installed = {m.ref for m in registry.list_installed()} | { 

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

49 } 

50 return ref in installed 

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

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

53 return False 

54 

55 

56def _is_api_ref_with_key(ref: str) -> ValidationResult: 

57 """Classify a non-local ref. Returns OK / NO_KEY / UNKNOWN. 

58 

59 OK when the parsed provider matches a configured API key. NO_KEY 

60 when the provider is recognized but the user hasn't set the key. 

61 UNKNOWN for malformed strings or unfamiliar providers. 

62 """ 

63 try: 

64 parsed = parse_model_ref(ref) 

65 except Exception: 

66 return ValidationResult.UNKNOWN 

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

68 if provider not in PROVIDER_API_KEY_FIELD: 

69 return ValidationResult.UNKNOWN 

70 return ValidationResult.OK if get_provider_api_key(provider) else ValidationResult.NO_KEY 

71 

72 

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

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

75 

76 Pure function. Reads cfg and the local registry; never mutates. 

77 """ 

78 if not ref: 

79 return ValidationResult.UNKNOWN 

80 if _is_local_installed(ref): 

81 return ValidationResult.OK 

82 return _is_api_ref_with_key(ref) 

83 

84 

85def _first_available_api_chat_ref() -> str | None: 

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

87 try: 

88 groups = discover_api_models() 

89 except Exception: 

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

91 return None 

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

93 if models: 

94 first = models[0] 

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

96 return None 

97 

98 

99def _first_installed_local_ref() -> str | None: 

100 """Return the first installed local ref, or ``None`` if none.""" 

101 try: 

102 registry = ModelRegistry(cfg.models_dir) 

103 installed = list(registry.list_installed()) 

104 except Exception: 

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

106 return None 

107 return installed[0].ref if installed else None 

108 

109 

110def _canonicalize(original: str, *, allow_api: bool) -> CanonicalRef: 

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

112 

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

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

115 have no embedding equivalent. 

116 """ 

117 status = validate_persisted_model(original) 

118 if status == ValidationResult.OK: 

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

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

121 if allow_api: 

122 candidates.append(_first_available_api_chat_ref()) 

123 candidates.append(_first_installed_local_ref()) 

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

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

126 

127 

128def canonicalize_chat_model() -> CanonicalRef: 

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

130 return _canonicalize(cfg.chat_model, allow_api=True) 

131 

132 

133def canonicalize_embedding_model() -> CanonicalRef: 

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

135 return _canonicalize(cfg.embedding_model, allow_api=False)