Coverage for src/lilbee/vision.py: 100%

1"""Vision-model OCR helpers: prompt resolution and OpenAI-compatible image messages. 

2 

3PDF rasterisation and the page loop now live inside xberg (the registered 

4lilbee-vision OCR backend); this module only builds the single-image request the 

5provider's ``vision_ocr`` sends to the vision server. 

6""" 

7 

8OCR_PROMPT = ( 

9 "Extract ALL text from this page as clean markdown. " 

10 "Preserve table structure using markdown table syntax. " 

11 "Include all rows, columns, headers, and page text exactly as shown." 

12) 

13 

14# Lowercase family token (as in the model ref or GGUF path) -> the model's 

15# documented OCR prompt. Unlisted models fall back to OCR_PROMPT. 

16_NATIVE_OCR_PROMPTS: tuple[tuple[str, str], ...] = ( 

17 ("deepseek-ocr", "<|grounding|>Convert the document to markdown."), 

18 ("glm-ocr", "OCR"), 

19) 

20 

21 

22def resolve_ocr_prompt(model_ref: str) -> str: 

23 """Return *model_ref*'s native OCR prompt, or the generic one if it has none.""" 

24 needle = model_ref.lower() 

25 for family, prompt in _NATIVE_OCR_PROMPTS: 

26 if family in needle: 

27 return prompt 

28 return OCR_PROMPT 

29 

30 

31def _png_to_data_url(png_bytes: bytes) -> str: 

32 """Convert raw PNG bytes to a base64 data URL for OpenAI-compatible messages.""" 

33 import base64 

34 

35 b64 = base64.b64encode(png_bytes).decode("ascii") 

36 return f"data:image/png;base64,{b64}" 

37 

38 

39def build_vision_messages(prompt: str, png_bytes: bytes) -> list[dict]: 

40 """Build OpenAI-compatible messages with image content for vision models. 

41 

42 Uses the multipart content format expected by llama.cpp's mtmd pipeline. 

43 """ 

44 return [ 

45 { 

46 "role": "user", 

47 "content": [ 

48 {"type": "image_url", "image_url": {"url": _png_to_data_url(png_bytes)}}, 

49 {"type": "text", "text": prompt}, 

50 ], 

51 } 

52 ]