Coverage for astrocyte/pipeline/_json

1"""Tolerant JSON parsing for LLM extraction outputs.

3gpt-4o-mini occasionally wraps its JSON in markdown fences or sprinkles

4prose around the object despite ``response_format={"type": "json_object"}``

5being set. The strict ``json.loads`` path then fails and the caller

6silently drops the section's facts / entities. This module gives the

7caller a layered fallback before giving up:

91. ``json.loads`` straight up — fast path, the common case.

102. Strip a ``` ```json ``` / ``` ``` ``` markdown fence wrapper.

113. Slice from the first ``{`` to the last ``}`` (drops surrounding prose).

13Returns the parsed object on success or ``None`` on total failure. The

14caller keeps its existing warn-and-return-empty fallback for ``None``.

16``looks_truncated`` is a sibling heuristic: when the LLM hit its

17``max_tokens`` budget mid-output, the JSON is unrecoverable and a retry

18with the same budget won't help either, so the caller should skip the

19retry path to avoid the latency cost.

20"""

22from __future__ import annotations

24import json

25from typing import Any

28def tolerant_json_loads(text: str) -> Any | None:

29 """Best-effort JSON parse. Returns ``None`` if all strategies fail."""

30 if not text:

31 return None

32 try:

33 return json.loads(text)

34 except json.JSONDecodeError:

35 pass

37 stripped = text.strip()

38 if stripped.startswith("```"):

39 # Drop the opening fence (possibly ```json) up to first newline.

40 newline = stripped.find("\n")

41 inner = stripped[newline + 1 :] if newline != -1 else stripped[3:]

42 # Drop trailing fence.

43 inner = inner.rstrip()

44 if inner.endswith("```"):

45 inner = inner[:-3].rstrip()

46 try:

47 return json.loads(inner)

48 except json.JSONDecodeError:

49 pass

51 start = text.find("{")

52 end = text.rfind("}")

53 if start != -1 and end > start:

54 try:

55 return json.loads(text[start : end + 1])

56 except json.JSONDecodeError:

57 pass

59 return None

62def looks_truncated(text: str) -> bool:

63 """Heuristic: did the LLM run out of budget mid-output?

65 Used to short-circuit the parse-failure retry path. Retrying when

66 the original response was budget-truncated wastes a round-trip — the

67 retry will hit the same cap. Counts are naive (don't track quoting),

68 which is fine: false positives just skip an occasional retry, false

69 negatives just spend an extra round-trip.

70 """

71 if not text:

72 return False

73 s = text.rstrip()

74 if not s:

75 return False

76 opens = s.count("{") + s.count("[")

77 # No braces at all → not a truncated JSON object. Probably a refusal

78 # or a prose answer; the retry path may recover with a stricter

79 # reminder, so don't short-circuit here.

80 if opens == 0:

81 return False

82 last = s[-1]

83 if last in (",", ":"):

84 return True

85 closes = s.count("}") + s.count("]")

86 if opens > closes:

87 return True

88 if last not in "}]\"'":

89 return True

90 return False

93def tolerant_json_loads_or_raise(text: str) -> Any:

94 """Like :func:`tolerant_json_loads` but raises ``json.JSONDecodeError``

95 when no strategy recovers a value.

97 Exists so call sites written as ``json.loads(completion.text)`` inside a

98 ``try/except json.JSONDecodeError`` block can upgrade to tolerant parsing

99 with a one-line swap, keeping their existing failure handling verbatim.

100 Providers with native constrained decoding (OpenAI json_object mode)

101 always emit bare JSON, so this is a provable no-op on that path; it only

102 widens acceptance for providers whose JSON mode is emulated (claude_cli)

103 or absent.

104 """

105 parsed = tolerant_json_loads(text)

106 if parsed is None:

107 raise json.JSONDecodeError("no JSON value recoverable", text[:80] or " ", 0)

108 return parsed

109

110

111#: Bounds the scan in :func:`first_json_value`: each candidate costs one

112#: ``raw_decode``, so a long prose reply full of brackets stays linear-ish.

113_MAX_CANDIDATES = 32

114

115

116def first_json_value(text: str, kind: type = object) -> Any | None:

117 """The first complete JSON value of type ``kind`` in ``text``, or ``None``.

118

119 Models wrap a JSON answer in a fence and then keep talking: ```` ```json

120 [] ``` ```` followed by "The new memory is a status update [...]". Slicing

121 from the first bracket to the *last* one swallows that prose whenever it

122 contains a bracket, and the whole answer is lost. ``raw_decode`` parses one

123 value and reports where it ended, so anything after the value is ignored

124 rather than fatal. Candidates are tried in order: the body of the first

125 fenced block, then each opening bracket in the text.

126 """

127 if not text:

128 return None

129 opener = "[" if kind is list else "{" if kind is dict else None

130 decoder = json.JSONDecoder()

131 candidates: list[str] = []

132 fence = text.find("```")

133 if fence != -1:

134 body_start = text.find("\n", fence)

135 body_end = text.find("```", body_start + 1) if body_start != -1 else -1

136 if body_end != -1:

137 candidates.append(text[body_start + 1 : body_end])

138 candidates.append(text)

139 for candidate in candidates:

140 tried = 0

141 idx = 0

142 while tried < _MAX_CANDIDATES:

143 if opener is None:

144 positions = [p for p in (candidate.find("[", idx), candidate.find("{", idx)) if p != -1]

145 pos = min(positions) if positions else -1

146 else:

147 pos = candidate.find(opener, idx)

148 if pos == -1:

149 break

150 tried += 1

151 try:

152 value, _end = decoder.raw_decode(candidate, pos)

153 except (json.JSONDecodeError, RecursionError):

154 # RecursionError: pathologically nested brackets. A parse

155 # failure like any other, not a crash of the caller.

156 idx = pos + 1

157 continue

158 if kind is object or isinstance(value, kind):

159 return value

160 idx = pos + 1

161 return None

Coverage for astrocyte/pipeline/_json_tolerant.py: 92%

86 statements