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

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

2 

3The architecture is read by walking the metadata KV table directly (see 

4``_parse_arch``) rather than via gguf-py's ``GGUFReader``, which needs the whole 

5KV table -- including a model's multi-megabyte tokenizer arrays -- present, and 

6so chokes on a Range-GET-truncated header. 

7""" 

8 

9from __future__ import annotations 

10 

11import logging 

12import struct 

13from http import HTTPStatus 

14 

15import httpx 

16from gguf import GGUFValueType, ReaderField 

17 

18log = logging.getLogger(__name__) 

19 

20GGUF_HEADER_PROBE_BYTES = 65536 

21GGUF_MAGIC = b"GGUF" 

22GGUF_ARCH_KEY = "general.architecture" 

23_PROBE_TIMEOUT_S = 10.0 

24 

25 

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

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

28 try: 

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

30 # follow_redirects: a HuggingFace ``/resolve/`` URL for an LFS-backed GGUF 

31 # 302-redirects to the CDN; without following it the probe reads the tiny 

32 # redirect body, the GGUF magic check fails, and every arch verdict is 

33 # silently UNKNOWN (so the unsupported-arch guard never fires). 

34 resp = httpx.get(blob_url, headers=headers, timeout=_PROBE_TIMEOUT_S, follow_redirects=True) 

35 if resp.status_code >= HTTPStatus.BAD_REQUEST: 

36 return "" 

37 return _parse_arch(resp.content) 

38 except httpx.HTTPError as exc: 

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

40 return "" 

41 

42 

43# GGUF metadata value-type tags (gguf spec) and the fixed byte sizes of the 

44# scalar ones. ARRAY/STRING carry their own length prefixes. 

45_GGUF_TYPE_STRING = 8 

46_GGUF_TYPE_ARRAY = 9 

47_GGUF_SCALAR_SIZES = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8} 

48# magic(4) + version(4) + tensor_count(8) + kv_count(8) 

49_GGUF_HEADER_FIXED = 24 

50 

51 

52class _TruncatedHeaderError(Exception): 

53 """The probe window ended before the bytes a parse step needed.""" 

54 

55 

56class _HeaderCursor: 

57 """Little-endian reader over the probed GGUF header bytes.""" 

58 

59 def __init__(self, blob: bytes) -> None: 

60 self._blob = blob 

61 self._pos = 0 

62 

63 def take(self, n: int) -> bytes: 

64 end = self._pos + n 

65 if n < 0 or end > len(self._blob): 

66 raise _TruncatedHeaderError 

67 chunk = self._blob[self._pos : end] 

68 self._pos = end 

69 return chunk 

70 

71 def u32(self) -> int: 

72 return int(struct.unpack_from("<I", self.take(4))[0]) 

73 

74 def u64(self) -> int: 

75 return int(struct.unpack_from("<Q", self.take(8))[0]) 

76 

77 def gguf_string(self) -> bytes: 

78 return self.take(self.u64()) 

79 

80 

81def _skip_value(cur: _HeaderCursor, value_type: int) -> None: 

82 """Advance *cur* past one metadata value of *value_type*.""" 

83 size = _GGUF_SCALAR_SIZES.get(value_type) 

84 if size is not None: 

85 cur.take(size) 

86 return 

87 if value_type == _GGUF_TYPE_STRING: 

88 cur.gguf_string() 

89 return 

90 if value_type == _GGUF_TYPE_ARRAY: 

91 elem_type = cur.u32() 

92 count = cur.u64() 

93 elem_size = _GGUF_SCALAR_SIZES.get(elem_type) 

94 if elem_size is not None: 

95 cur.take(elem_size * count) 

96 elif elem_type == _GGUF_TYPE_STRING: 

97 for _ in range(count): 

98 cur.gguf_string() 

99 else: 

100 raise _TruncatedHeaderError # nested/unknown array element: stop parsing 

101 return 

102 raise _TruncatedHeaderError # unknown value type: stop parsing 

103 

104 

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

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

107 

108 Walks the metadata KV table directly and returns as soon as it reaches 

109 ``general.architecture``, which GGUF writers emit among the first entries. 

110 This deliberately avoids gguf-py's ``GGUFReader``, which parses the entire KV 

111 table up front: a real model's multi-megabyte tokenizer arrays run past the 

112 Range-GET probe window, so ``GGUFReader`` raises on the truncation before any 

113 field can be read. Returns empty string on a malformed or too-short header. 

114 """ 

115 if len(blob) < _GGUF_HEADER_FIXED or blob[:4] != GGUF_MAGIC: 

116 return "" 

117 cur = _HeaderCursor(blob) 

118 arch_key = GGUF_ARCH_KEY.encode("utf-8") 

119 try: 

120 cur.take(4) # magic 

121 cur.u32() # version 

122 cur.u64() # tensor_count (the tensor-info table is never read) 

123 kv_count = cur.u64() 

124 for _ in range(kv_count): 

125 key = cur.gguf_string() 

126 value_type = cur.u32() 

127 if key == arch_key: 

128 if value_type != _GGUF_TYPE_STRING: 

129 return "" 

130 return cur.gguf_string().decode("utf-8", errors="replace") 

131 _skip_value(cur, value_type) 

132 except _TruncatedHeaderError: 

133 return "" 

134 return "" 

135 

136 

137def gguf_scalar_str(field: ReaderField | None) -> str | None: 

138 """Render a scalar GGUF metadata field as a string, or ``None``. 

139 

140 Mirrors what ``Llama(vocab_only=True).metadata`` produced: STRING fields 

141 decode their bytes; numeric scalars (UINT*/INT*/FLOAT*/BOOL) stringify 

142 their single value. ARRAY fields and empty fields return ``None`` (callers 

143 here only read scalars). This is the one place the gguf-py ReaderField shape 

144 is decoded, so a gguf-py layout change breaks here, not at every call site. 

145 """ 

146 if field is None or not field.types or not field.data: 

147 return None 

148 value_type = field.types[-1] 

149 part = field.parts[field.data[0]] 

150 if value_type == GGUFValueType.STRING: 

151 return bytes(part).decode("utf-8", errors="replace") 

152 if value_type == GGUFValueType.ARRAY: 

153 return None 

154 scalar = part.tolist() if hasattr(part, "tolist") else part 

155 if isinstance(scalar, (list, tuple)): 

156 scalar = scalar[0] if scalar else None 

157 return None if scalar is None else str(scalar)