Coverage for src/local_deep_research/utilities/json_utils.py: 92%

1""" 

2Centralized utilities for extracting and parsing JSON from LLM responses. 

3 

4Provides two public functions: 

5- get_llm_response_text: Extract text from LLM response objects 

6- extract_json: Parse JSON from LLM-generated text with robust cleaning 

7""" 

8 

9import json 

10import re 

11from typing import Optional, Type, Union 

12 

13from loguru import logger 

14 

15from .search_utilities import remove_think_tags as _remove_think_tags 

16 

17 

18def get_llm_response_text(response) -> str: 

19 """Extract text content from an LLM response object. 

20 

21 Handles LangChain AIMessage (.content), plain text responses (.text), 

22 list-type content blocks (Anthropic extended-thinking / tool-use, where 

23 ``.content`` is a list of blocks), and arbitrary objects (via str()). 

24 Removes <think> tags from the output. 

25 

26 Args: 

27 response: LLM response object, string, or None. 

28 

29 Returns: 

30 Extracted text with think tags removed. Empty string for None input. 

31 """ 

32 if response is None: 

33 return "" 

34 if hasattr(response, "content") and response.content is not None: 

35 raw = response.content 

36 elif hasattr(response, "text") and response.text is not None: 

37 raw = response.text 

38 else: 

39 raw = str(response) 

40 if isinstance(raw, list): 

41 # Anthropic-style content blocks, e.g. 

42 # [{"type": "text", "text": "..."}, {"type": "tool_use", ...}]. 

43 # Extract and join the text so downstream string ops get clean text 

44 # rather than the list's Python repr (the #4615 data-corruption bug). 

45 raw = _coerce_content_blocks(raw) 

46 if not isinstance(raw, str): 

47 raw = str(raw) 

48 return _remove_think_tags(raw) 

49 

50 

51def _coerce_content_blocks(blocks: list) -> str: 

52 """Join the text from a list of LLM content blocks into a plain string. 

53 

54 Some providers (notably Anthropic, and some LangChain paths) return 

55 ``message.content`` as a list of blocks rather than a string — for 

56 example ``[{"type": "text", "text": "Paris"}, {"type": "tool_use", ...}]``. 

57 Keep only the textual parts; skip non-text blocks (tool_use, thinking, 

58 etc.) so callers that then ``.strip()`` / ``.split()`` operate on real 

59 text instead of the list's ``repr()``. 

60 

61 Accepts blocks that are plain strings, dicts with a ``"text"`` key (when 

62 ``type`` is ``"text"`` or unset), or objects exposing a ``.text`` str 

63 attribute. Returns ``""`` for an empty/textless list. 

64 """ 

65 parts: list[str] = [] 

66 for block in blocks: 

67 if isinstance(block, str): 67 ↛ 68line 67 didn't jump to line 68 because the condition on line 67 was never true

68 parts.append(block) 

69 elif isinstance(block, dict): 69 ↛ 75line 69 didn't jump to line 75 because the condition on line 69 was always true

70 if block.get("type", "text") == "text" and isinstance( 

71 block.get("text"), str 

72 ): 

73 parts.append(block["text"]) 

74 else: 

75 text_attr = getattr(block, "text", None) 

76 if isinstance(text_attr, str): 

77 parts.append(text_attr) 

78 return "".join(parts) 

79 

80 

81def extract_json( 

82 text: str, 

83 expected_type: Optional[Type] = None, 

84) -> Optional[Union[dict, list]]: 

85 """Extract and parse JSON from LLM-generated text. 

86 

87 Applies a cleaning pipeline to handle common LLM output patterns: 

88 code fences, think tags, prose surrounding JSON, and minor artifacts. 

89 

90 Args: 

91 text: Raw text potentially containing JSON. 

92 expected_type: Expected JSON type (dict or list). If specified, 

93 bracket extraction is ordered to prefer the matching type. 

94 None accepts either type. 

95 

96 Returns: 

97 Parsed dict or list, or None if no valid JSON found. 

98 """ 

99 if not text or not text.strip(): 

100 return None 

101 

102 text = text.strip() 

103 text = _strip_code_fences(text) 

104 text = _remove_think_tags(text) 

105 

106 # Step 1: Try direct parse 

107 try: 

108 result = json.loads(text) 

109 if isinstance(result, (dict, list)): 

110 if expected_type is None or isinstance(result, expected_type): 

111 return result 

112 # Type mismatch — fall through to bracket extraction 

113 except (json.JSONDecodeError, ValueError): 

114 pass 

115 

116 # Step 2: Bracket extraction ordered by expected_type 

117 if expected_type is list: 

118 bracket_pairs = [("[", "]"), ("{", "}")] 

119 elif expected_type is dict: 

120 bracket_pairs = [("{", "}"), ("[", "]")] 

121 else: 

122 bracket_pairs = [("{", "}"), ("[", "]")] 

123 

124 for open_char, close_char in bracket_pairs: 

125 extracted = _extract_by_brackets(text, open_char, close_char) 

126 if extracted is None: 

127 continue 

128 

129 # Try parsing the extracted substring 

130 try: 

131 result = json.loads(extracted) 

132 if isinstance(result, (dict, list)): 132 ↛ 139line 132 didn't jump to line 139 because the condition on line 132 was always true

133 if expected_type is None or isinstance(result, expected_type): 

134 return result 

135 except (json.JSONDecodeError, ValueError): 

136 pass 

137 

138 # Try cleaning LLM artifacts and retrying 

139 cleaned = _clean_llm_json_artifacts(extracted) 

140 if cleaned != extracted: 

141 try: 

142 result = json.loads(cleaned) 

143 if isinstance(result, (dict, list)): 143 ↛ 124line 143 didn't jump to line 124 because the condition on line 143 was always true

144 if expected_type is None or isinstance( 144 ↛ 124line 144 didn't jump to line 124 because the condition on line 144 was always true

145 result, expected_type 

146 ): 

147 return result 

148 except (json.JSONDecodeError, ValueError): 

149 pass 

150 

151 logger.debug("No valid JSON found in text") 

152 return None 

153 

154 

155# --------------------------------------------------------------------------- 

156# Private helpers 

157# --------------------------------------------------------------------------- 

158 

159 

160def _strip_code_fences(text: str) -> str: 

161 """Remove markdown code fences from text. 

162 

163 Uses split-based extraction (not startswith) to handle fences 

164 appearing mid-text with surrounding prose. 

165 """ 

166 if "```json" in text: 

167 parts = text.split("```json") 

168 if len(parts) > 1: 168 ↛ 174line 168 didn't jump to line 174 because the condition on line 168 was always true

169 return parts[1].split("```")[0].strip() 

170 elif "```" in text: 

171 parts = text.split("```") 

172 if len(parts) >= 3: 

173 return parts[1].strip() 

174 return text 

175 

176 

177def _extract_by_brackets( 

178 text: str, open_char: str, close_char: str 

179) -> Optional[str]: 

180 """Extract substring between outermost matching brackets. 

181 

182 Uses find()/rfind() which is equivalent to re.search with re.DOTALL 

183 for bracket matching. 

184 """ 

185 start = text.find(open_char) 

186 end = text.rfind(close_char) 

187 if start >= 0 and end > start: 

188 return text[start : end + 1] 

189 return None 

190 

191 

192def _clean_llm_json_artifacts(text: str) -> str: 

193 """Clean common LLM JSON artifacts from malformed JSON text. 

194 

195 Only called after json.loads has already failed, so the text is 

196 already malformed. These regexes cannot corrupt valid JSON. 

197 

198 Handles: 

199 - Trailing commas before ] or } 

200 - Inline // comments 

201 - Ellipsis entries (... or "...") 

202 """ 

203 # Remove trailing commas 

204 text = re.sub(r",\s*([}\]])", r"\1", text) 

205 # Remove // line comments 

206 text = re.sub(r"//[^\n]*", "", text) 

207 # Remove ellipsis entries, preserving comma separator when between items 

208 text = re.sub( 

209 r',\s*"?\.\.\.+"?\s*,', ",", text 

210 ) # between items: keep one comma 

211 return re.sub(r',?\s*"?\.\.\.+"?\s*', "", text) # trailing/leading ellipsis