Coverage for src/lilbee/modelhub/role_validator.py: 100%

1"""Role-slot assignment validation for the four model config fields.""" 

2 

3import os 

4import sys 

5 

6from lilbee.catalog import CatalogModel, find_pick 

7from lilbee.catalog.query import reclassify_by_name 

8from lilbee.catalog.refs import is_bare_hf_repo 

9from lilbee.catalog.types import ModelTask 

10from lilbee.core.config import cfg 

11from lilbee.modelhub.registry import ModelRegistry 

12from lilbee.providers.model_ref import PROVIDER_PREFIXES, is_native_gguf_ref 

13 

14# Test-only bypass. Both the env var and pytest must be present so a 

15# leaked env var cannot disable validation in production. 

16_SKIP_MODEL_TASK_VALIDATION_ENV = "LILBEE_SKIP_MODEL_TASK_VALIDATION" 

17 

18MODEL_FIELD_TO_TASK: dict[str, str] = { 

19 "chat_model": "chat", 

20 "embedding_model": "embedding", 

21 "vision_model": "vision", 

22 "reranker_model": "rerank", 

23} 

24 

25 

26class TaskMismatchError(ValueError): 

27 """A role slot was assigned a model whose catalog task does not match. 

28 

29 Carries the structured fields so each surface (HTTP, CLI, TUI, MCP) 

30 can format its own user-facing message. The default ``str()`` form is 

31 surface-neutral so it is safe to surface unmodified. 

32 """ 

33 

34 def __init__(self, ref: str, entry_task: ModelTask, expected_task: ModelTask) -> None: 

35 self.ref = ref 

36 self.entry_task = entry_task 

37 self.expected_task = expected_task 

38 super().__init__(f"Model '{ref}' is a {entry_task} model, not {expected_task}.") 

39 

40 

41def _model_task_validation_bypassed() -> bool: 

42 if not os.environ.get(_SKIP_MODEL_TASK_VALIDATION_ENV): 

43 return False 

44 return sys.modules.get("pytest") is not None 

45 

46 

47def _resolve_installed_task(registry: ModelRegistry, ref: str) -> ModelTask | None: 

48 """Return the manifest's ``ModelTask`` for *ref*, name-reclassified, or ``None``.""" 

49 manifest = registry.get_manifest(ref) 

50 if manifest is None: 

51 return None 

52 return ModelTask(reclassify_by_name(ref, manifest.task)) 

53 

54 

55def _skips_catalog_check(ref: str, *, allow_bypass: bool) -> bool: 

56 """Whether *ref* skips the catalog check.""" 

57 if not ref or not ref.strip(): 

58 return True 

59 if allow_bypass and _model_task_validation_bypassed(): 

60 return True 

61 return ref.split("/", 1)[0] in PROVIDER_PREFIXES 

62 

63 

64def _canonical_pick_ref(ref: str, entry: CatalogModel, want: ModelTask) -> str: 

65 """Role-check a current pick and choose the canonical ref to persist.""" 

66 if entry.task != want: 

67 raise TaskMismatchError(ref, ModelTask(entry.task), want) 

68 # Keep a full ``<repo>/<file>.gguf`` so resolve_model_path lands on the 

69 # exact installed quant; fall back to the pick's own ref otherwise. 

70 if is_native_gguf_ref(ref): 

71 return ref 

72 canonical: str = entry.ref 

73 return canonical 

74 

75 

76def _installed_ref_and_task(ref: str) -> tuple[str, str | None]: 

77 """Canonical installed ref for *ref* and its manifest task, task None if absent. 

78 

79 A bare ``<org>/<repo>`` ref canonicalizes to its installed quant's full ref 

80 so the persisted value always names the exact GGUF file. 

81 """ 

82 registry = ModelRegistry(cfg.models_dir) 

83 if is_bare_hf_repo(ref): 

84 ref = registry.installed_ref_for_repo(ref) or ref 

85 return ref, _resolve_installed_task(registry, ref) 

86 

87 

88def _not_installed(ref: str) -> ValueError: 

89 """The error for a ref that is neither a current pick nor installed.""" 

90 return ValueError( 

91 f"Model '{ref}' is not installed. " 

92 "Install it with 'lilbee model pull <ref>' " 

93 "(or POST /api/models/pull) before assigning it to a role." 

94 ) 

95 

96 

97def validate_model_task_assignment(field_name: str, ref: str, *, allow_bypass: bool = True) -> str: 

98 """Check *ref* is assignable to *field_name*; return the canonical ref. 

99 

100 A current pick carries its own task, so it can be role-checked before it is 

101 installed, which is what the catalog UI offers. Anything else is checked 

102 against the installed manifest, the only other thing that can vouch for a 

103 model's role. Raises ``TaskMismatchError`` on role mismatch and 

104 ``ValueError`` when the model is neither a pick nor installed. 

105 """ 

106 if _skips_catalog_check(ref, allow_bypass=allow_bypass): 

107 return ref 

108 want = ModelTask(MODEL_FIELD_TO_TASK[field_name]) 

109 # The manifest is consulted first because it answers without touching the 

110 # network. This runs on the TUI main thread from /model and the model-bar 

111 # picker, where resolving picks would block the UI. 

112 installed_ref, installed_task = _installed_ref_and_task(ref) 

113 if installed_task is not None: 

114 if installed_task != want: 

115 raise TaskMismatchError(installed_ref, ModelTask(installed_task), want) 

116 return installed_ref 

117 entry = find_pick(ref) 

118 if entry is not None: 

119 return _canonical_pick_ref(ref, entry, want) 

120 raise _not_installed(installed_ref)