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

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 (when their ``type`` attribute, if any, is ``"text"``). 

64 Returns ``""`` for an empty/textless list. 

65 """ 

66 parts: list[str] = [] 

67 for block in blocks: 

68 if isinstance(block, str): 

69 parts.append(block) 

70 elif isinstance(block, dict): 

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

72 block.get("text"), str 

73 ): 

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

75 else: 

76 # A tool_result or thinking block object can expose .text too. 

77 # As for dicts, only an untyped or "text"-typed block counts. 

78 if getattr(block, "type", "text") != "text": 

79 continue 

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

81 if isinstance(text_attr, str): 81 ↛ 67line 81 didn't jump to line 67 because the condition on line 81 was always true

82 parts.append(text_attr) 

83 return "".join(parts) 

84 

85 

86def extract_json( 

87 text: str, 

88 expected_type: Optional[Type] = None, 

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

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

91 

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

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

94 

95 Args: 

96 text: Raw text potentially containing JSON. 

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

98 bracket extraction is ordered to prefer the matching type. 

99 None accepts either type. 

100 

101 Returns: 

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

103 """ 

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

105 return None 

106 

107 text = text.strip() 

108 text = _strip_code_fences(text) 

109 text = _remove_think_tags(text) 

110 

111 # Step 1: Try direct parse 

112 try: 

113 result = json.loads(text) 

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

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

116 return result 

117 # Type mismatch — fall through to bracket extraction 

118 except (json.JSONDecodeError, ValueError): 

119 pass 

120 

121 # Step 2: Bracket extraction ordered by expected_type 

122 if expected_type is list: 

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

124 elif expected_type is dict: 

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

126 else: 

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

128 

129 for open_char, close_char in bracket_pairs: 

130 extracted = _extract_by_brackets(text, open_char, close_char) 

131 if extracted is None: 

132 continue 

133 

134 # Try parsing the extracted substring 

135 try: 

136 result = json.loads(extracted) 

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

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

139 return result 

140 except (json.JSONDecodeError, ValueError): 

141 pass 

142 

143 # Try cleaning LLM artifacts and retrying 

144 cleaned = _clean_llm_json_artifacts(extracted) 

145 if cleaned != extracted: 

146 try: 

147 result = json.loads(cleaned) 

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

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

150 result, expected_type 

151 ): 

152 return result 

153 except (json.JSONDecodeError, ValueError): 

154 pass 

155 

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

157 return None 

158 

159 

160# --------------------------------------------------------------------------- 

161# Private helpers 

162# --------------------------------------------------------------------------- 

163 

164 

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

166 """Remove markdown code fences from text. 

167 

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

169 appearing mid-text with surrounding prose. 

170 """ 

171 if "```json" in text: 

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

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

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

175 elif "```" in text: 

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

177 if len(parts) >= 3: 

178 return parts[1].strip() 

179 return text 

180 

181 

182def _extract_by_brackets( 

183 text: str, open_char: str, close_char: str 

184) -> Optional[str]: 

185 """Extract substring between outermost matching brackets. 

186 

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

188 for bracket matching. 

189 """ 

190 start = text.find(open_char) 

191 end = text.rfind(close_char) 

192 if start >= 0 and end > start: 

193 return text[start : end + 1] 

194 return None 

195 

196 

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

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

199 

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

201 already malformed. These regexes cannot corrupt valid JSON. 

202 

203 Handles: 

204 - Trailing commas before ] or } 

205 - Inline // comments 

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

207 """ 

208 # Remove trailing commas 

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

210 # Remove // line comments 

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

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

213 text = re.sub( 

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

215 ) # between items: keep one comma 

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