Coverage for src/lilbee/catalog/compat.py: 100%

1"""Architecture compatibility classification for catalog entries.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING 

6 

7from huggingface_hub import hf_hub_url 

8from huggingface_hub.utils import HFValidationError 

9 

10from lilbee._generated.engine_archs import SUPPORTED_ARCHS 

11from lilbee.catalog.header_probe import GgufHeader, probe_header 

12from lilbee.catalog.refs import ( 

13 GGUF_SUFFIX, 

14 NATIVE_GGUF_REF_MIN_SLASHES, 

15 WILDCARD, 

16 gguf_filename_from_ref, 

17 hf_repo_from_ref, 

18) 

19from lilbee.catalog.types import ModelCompat 

20 

21if TYPE_CHECKING: 

22 from lilbee.catalog.hf_client import HfClient 

23 

24 

25def classify(architecture: str) -> ModelCompat: 

26 """Map a `general.architecture` string to a `ModelCompat` verdict.""" 

27 if not architecture: 

28 return ModelCompat.UNKNOWN 

29 return ModelCompat.SUPPORTED if architecture in SUPPORTED_ARCHS else ModelCompat.UNSUPPORTED 

30 

31 

32def resolve_arch_for_pull(ref: str, hf_client: HfClient) -> str: 

33 """Resolve general.architecture for *ref*: cache hit > Range-GET probe > empty (UNKNOWN). 

34 

35 ``probe_header`` returns empty fields on any failure (network, non-200, parse), 

36 so an empty result means "undetermined", never a real verdict. Only a non-empty 

37 arch is cached; otherwise a transient probe failure would be cached permanently 

38 and disable the unsupported-arch guard for this ref on every later pull. 

39 """ 

40 cached = hf_client.get_cached_arch(ref) 

41 if cached is not None: 

42 return cached 

43 url = _resolve_blob_url(ref) 

44 if not url: 

45 return "" 

46 arch = probe_header(url).architecture 

47 if arch: 

48 hf_client.cache_arch(ref, arch) 

49 return arch 

50 

51 

52def file_header(hf_repo: str, filename: str) -> GgufHeader: 

53 """The GGUF header of one repo file, empty when it cannot be read. 

54 

55 An unreadable header is not a verdict: every field stays empty, which reads 

56 as "a model of undetermined architecture" so an offline or gated repo 

57 resolves the way it did before the probe existed. 

58 """ 

59 try: 

60 url = hf_hub_url(hf_repo, filename) 

61 except (HFValidationError, ValueError): 

62 return GgufHeader() 

63 return probe_header(url) 

64 

65 

66def _resolve_blob_url(ref: str) -> str: 

67 """Return a probable .gguf blob URL for *ref*, or empty string if unresolvable. 

68 

69 lilbee's canonical native refs are slash-delimited ``<org>/<repo>/<file>.gguf`` 

70 (the filename may add subdirs for a quant), so the repo and filename are split 

71 on that shape; an ollama-style ``repo:tag`` ref is split on the colon. 

72 """ 

73 if ref.endswith(GGUF_SUFFIX) and ref.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES: 

74 repo, filename = hf_repo_from_ref(ref), gguf_filename_from_ref(ref) 

75 elif ":" in ref: 

76 repo, filename = ref.split(":", 1) 

77 else: 

78 repo, filename = ref, "" 

79 if not filename or WILDCARD in filename: 

80 return "" 

81 try: 

82 return hf_hub_url(repo, filename) 

83 except (HFValidationError, ValueError): 

84 return "" 

85 

86 

87class UnsupportedArchError(Exception): 

88 """Raised when a pull is attempted for a model whose architecture isn't supported.""" 

89 

90 def __init__(self, ref: str, architecture: str) -> None: 

91 self.ref = ref 

92 self.architecture = architecture 

93 super().__init__( 

94 f"Model {ref!r} uses architecture {architecture!r}, not supported by this lilbee build." 

95 ) 

96 

97 

98class UnsupportedQuantError(RuntimeError): 

99 """Raised when a GGUF's tensors carry a type the bundled engine cannot decode. 

100 

101 A RuntimeError so every pull surface renders it the way it already renders a 

102 failed pull. Unlike an unsupported architecture there is no preflight route 

103 and no list of accepted types to offer, so there is nothing structured for a 

104 client to read beyond the message. 

105 """ 

106 

107 def __init__(self, ref: str, quant: str) -> None: 

108 self.ref = ref 

109 self.quant = quant 

110 super().__init__( 

111 f"Model {ref!r} is quantized as {quant}, which this lilbee build cannot load. " 

112 "Pass --allow-unsupported to download it anyway." 

113 )