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 probe_architecture 

12from lilbee.catalog.refs import ( 

13 NATIVE_GGUF_REF_MIN_SLASHES, 

14 gguf_filename_from_ref, 

15 hf_repo_from_ref, 

16) 

17from lilbee.catalog.types import ModelCompat 

18 

19if TYPE_CHECKING: 

20 from lilbee.catalog.hf_client import HfClient 

21 

22 

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

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

25 if not architecture: 

26 return ModelCompat.UNKNOWN 

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

28 

29 

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

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

32 

33 ``probe_architecture`` returns ``""`` on any failure (network, non-200, parse), 

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

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

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

37 """ 

38 cached = hf_client.get_cached_arch(ref) 

39 if cached is not None: 

40 return cached 

41 url = _resolve_blob_url(ref) 

42 if not url: 

43 return "" 

44 arch = probe_architecture(url) 

45 if arch: 

46 hf_client.cache_arch(ref, arch) 

47 return arch 

48 

49 

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

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

52 

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

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

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

56 """ 

57 if ref.endswith(".gguf") and ref.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES: 

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

59 elif ":" in ref: 

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

61 else: 

62 repo, filename = ref, "" 

63 if not filename or "*" in filename: 

64 return "" 

65 try: 

66 return hf_hub_url(repo, filename) 

67 except (HFValidationError, ValueError): 

68 return "" 

69 

70 

71class UnsupportedArchError(Exception): 

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

73 

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

75 self.ref = ref 

76 self.architecture = architecture 

77 super().__init__( 

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

79 )