Coverage for src/lilbee/core/llm_json.py: 100%

1"""Extract the first JSON value from an LLM reply. 

2 

3Models wrap JSON in prose, fences, and trailing commentary. A greedy 

4``\\{.*\\}`` span runs to the LAST brace in the reply, so any trailing text 

5containing one drops the whole value; a hand-rolled brace counter miscounts 

6braces inside string literals. ``raw_decode`` scanned from each opener lets the 

7stdlib own the parsing state, which is the only thing that gets both right. 

8""" 

9 

10from __future__ import annotations 

11 

12import json 

13from typing import TypeVar 

14 

15_DECODER = json.JSONDecoder() 

16 

17 

18def json_reply_format() -> dict[str, str]: 

19 """Provider option asking for a bare JSON value rather than JSON in prose. 

20 

21 llama.cpp's server and OpenAI-compatible remotes both honour it; providers 

22 that do not drop it rather than refusing (see the litellm adapter). Returns 

23 a fresh dict per call because it is handed to a third-party SDK that is free 

24 to mutate the request it is given. The scans below stay the fallback: a 

25 provider can ignore the request and a model can comply imperfectly, and 

26 neither should cost the caller its answer. 

27 """ 

28 return {"type": "json_object"} 

29 

30 

31_JsonT = TypeVar("_JsonT", dict, list) 

32 

33 

34def _first_json_value(text: str, opener: str, expected: type[_JsonT]) -> _JsonT | None: 

35 """The first *expected*-typed JSON value starting at an *opener*, or None.""" 

36 start = text.find(opener) 

37 while start >= 0: 

38 try: 

39 parsed, _ = _DECODER.raw_decode(text, start) 

40 except json.JSONDecodeError: 

41 start = text.find(opener, start + 1) 

42 continue 

43 return parsed if isinstance(parsed, expected) else None 

44 return None 

45 

46 

47def first_json_object(text: str) -> dict | None: 

48 """The first JSON object in *text*, or None.""" 

49 return _first_json_value(text, "{", dict) 

50 

51 

52def first_json_array(text: str) -> list | None: 

53 """The first JSON array in *text*, or None.""" 

54 return _first_json_value(text, "[", list)