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

1"""Range-GET a GGUF blob's header and extract general.architecture via gguf-py.""" 

2 

3from __future__ import annotations 

4 

5import contextlib 

6import logging 

7import struct 

8import tempfile 

9from http import HTTPStatus 

10from pathlib import Path 

11 

12import httpx 

13from gguf import GGUFReader, GGUFValueType 

14 

15log = logging.getLogger(__name__) 

16 

17GGUF_HEADER_PROBE_BYTES = 65536 

18GGUF_MAGIC = b"GGUF" 

19GGUF_ARCH_KEY = "general.architecture" 

20_PROBE_TIMEOUT_S = 10.0 

21_TENSOR_COUNT_OFFSET = 8 

22_TENSOR_COUNT_SIZE = 8 

23 

24 

25def probe_architecture(blob_url: str) -> str: 

26 """Return general.architecture from the GGUF header, or empty string on any failure.""" 

27 try: 

28 headers = {"Range": f"bytes=0-{GGUF_HEADER_PROBE_BYTES - 1}"} 

29 resp = httpx.get(blob_url, headers=headers, timeout=_PROBE_TIMEOUT_S) 

30 if resp.status_code >= HTTPStatus.BAD_REQUEST: 

31 return "" 

32 return _parse_arch(resp.content) 

33 except httpx.HTTPError as exc: 

34 log.debug("GGUF header probe failed for %s: %s", blob_url, exc) 

35 return "" 

36 

37 

38def _parse_arch(blob: bytes) -> str: 

39 """Extract general.architecture from a (possibly truncated) GGUF header. 

40 

41 Patches the tensor_count field to zero so ``gguf-py``'s ``GGUFReader`` 

42 skips the tensor info table that would otherwise require the full 

43 file. Only the KV table needs to be intact for architecture lookup. 

44 """ 

45 if len(blob) < _TENSOR_COUNT_OFFSET + _TENSOR_COUNT_SIZE or blob[:4] != GGUF_MAGIC: 

46 return "" 

47 patched = bytearray(blob) 

48 patched[_TENSOR_COUNT_OFFSET : _TENSOR_COUNT_OFFSET + _TENSOR_COUNT_SIZE] = struct.pack("<Q", 0) 

49 with tempfile.NamedTemporaryFile(suffix=".gguf", delete=False) as f: 

50 f.write(bytes(patched)) 

51 tmp = Path(f.name) 

52 try: 

53 return _read_architecture(tmp) 

54 except (ValueError, struct.error, IndexError, OSError, UnicodeDecodeError) as exc: 

55 log.debug("GGUFReader parse failed: %s", exc) 

56 return "" 

57 finally: 

58 # GGUFReader memory-maps the file; on Windows the mapping must be released 

59 # before the file can be deleted (a held mapping raises PermissionError / 

60 # WinError 32, which missing_ok does not cover). _read_architecture scopes 

61 # the reader so it is freed on return; suppress is a best-effort backstop. 

62 with contextlib.suppress(OSError): 

63 tmp.unlink() 

64 

65 

66def _read_architecture(path: Path) -> str: 

67 """Return general.architecture from the GGUF file at *path*, or empty string. 

68 

69 Kept separate so the ``GGUFReader`` (and its memory map of *path*) is released 

70 when this returns, letting the caller delete the temp file on Windows. 

71 """ 

72 reader = GGUFReader(str(path)) 

73 field = reader.fields.get(GGUF_ARCH_KEY) 

74 if field is None or not field.data: 

75 return "" 

76 if not field.types or field.types[-1] != GGUFValueType.STRING: 

77 return "" 

78 return bytes(field.parts[field.data[0]]).decode("utf-8", errors="replace")