Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/structured/parser.py: 13%

150 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

1"""Output parsing strategies for LLM-generated content. 

2 

3Provides a thin layer of JSON extraction utilities used by 

4:class:`~lexigram.ai.llm.structured.extractor.StructuredExtractor`. 

5 

6The key challenge is that LLMs do not always emit pure JSON: 

7 

8- They may wrap it in markdown code blocks (````json ... ````). 

9- They may include prose before or after the JSON object. 

10- They may emit partial or malformed JSON on the first attempt. 

11 

12:func:`extract_json_block` handles the common surface area — strip 

13fences, locate the outermost ``{`` / ``[``, and attempt a parse. 

14""" 

15 

16from __future__ import annotations 

17 

18import re 

19from typing import Any, cast 

20 

21from lexigram.ai.llm.thinking import normalize_thinking_text 

22from lexigram.serialization import JSONDecodeError as _JSONDecodeError 

23from lexigram.serialization import loads as _json_loads 

24 

25# --------------------------------------------------------------------------- 

26# Constants 

27# --------------------------------------------------------------------------- 

28 

29_FENCE_RE = re.compile( 

30 r"```(?:json|JSON)?\s*\n?(.*?)\n?\s*```", 

31 re.DOTALL, 

32) 

33"""Matches markdown-fenced JSON blocks.""" 

34 

35 

36# --------------------------------------------------------------------------- 

37# Internal helpers 

38# --------------------------------------------------------------------------- 

39 

40 

41def _extract_first_json_block(text: str) -> str | None: 

42 """Find the first complete, parseable JSON object or array using bracket counting. 

43 

44 Scans *text* left-to-right. For each ``{`` or ``[`` encountered it uses 

45 a depth counter — honouring string literals and backslash escapes — to 

46 locate the matching closing bracket. If the resulting candidate is valid 

47 JSON it is returned as a raw string; otherwise the scan resumes from the 

48 character *after* the closing bracket so that later valid blocks are still 

49 found. 

50 

51 This replaces the former ``re.compile(r"\\{.*\\}", re.DOTALL)`` approach 

52 which greedily matched from the **first** opening brace to the **last** 

53 closing brace, producing garbage for any input with multiple JSON objects 

54 or brace-containing string values. 

55 

56 Args: 

57 text: Raw text that may contain one or more embedded JSON blocks. 

58 

59 Returns: 

60 The first complete, parseable JSON block as a raw string, or ``None`` 

61 when no valid JSON block is found. 

62 """ 

63 i = 0 

64 length = len(text) 

65 

66 while i < length: 

67 ch = text[i] 

68 if ch not in ("{", "["): 

69 i += 1 

70 continue 

71 

72 open_char = ch 

73 close_char = "}" if ch == "{" else "]" 

74 

75 # Bracket-count from position i to find the matching close. 

76 depth = 0 

77 in_string = False 

78 escape_next = False 

79 end_idx: int | None = None 

80 

81 for j in range(i, length): 

82 c = text[j] 

83 

84 if escape_next: 

85 escape_next = False 

86 continue 

87 

88 if c == "\\" and in_string: 

89 escape_next = True 

90 continue 

91 

92 if c == '"': 

93 in_string = not in_string 

94 continue 

95 

96 if in_string: 

97 continue 

98 

99 if c == open_char: 

100 depth += 1 

101 elif c == close_char: 

102 depth -= 1 

103 if depth == 0: 

104 end_idx = j 

105 break 

106 

107 if end_idx is None: 

108 # No matching close bracket found — no more valid JSON possible. 

109 break 

110 

111 candidate = text[i : end_idx + 1] 

112 try: 

113 _json_loads(candidate) 

114 return candidate 

115 except (ValueError, _JSONDecodeError): 

116 # Not valid JSON (e.g. bare words like ``{some text}``). 

117 # Advance past this block and keep looking. 

118 i = end_idx + 1 

119 

120 return None 

121 

122 

123# --------------------------------------------------------------------------- 

124# Public helpers 

125# --------------------------------------------------------------------------- 

126 

127 

128def extract_json_block(text: str) -> Any: 

129 """Extract and parse a JSON value from an LLM response string. 

130 

131 Tries the following strategies in order: 

132 

133 0. Strip inline thinking/reasoning preamble (Qwen3, Gemma-4, DeepSeek-R1, etc.). 

134 1. Strip markdown fences (````` ```json … ``` `````) and parse directly. 

135 2. Parse the text as-is (the model returned raw JSON). 

136 3. Locate the first ``{`` or ``[`` and parse from there. 

137 

138 Args: 

139 text: Raw LLM response text. 

140 

141 Returns: 

142 The parsed Python value (dict, list, str, int, …). 

143 

144 Raises: 

145 ValueError: When no valid JSON can be extracted from the text. 

146 """ 

147 stripped = text.strip() 

148 

149 # Step 0: strip inline thinking/reasoning preamble before attempting JSON parse 

150 stripped, _ = normalize_thinking_text(stripped) 

151 

152 # Strategy 1: markdown fence 

153 fence_match = _FENCE_RE.search(stripped) 

154 if fence_match: 

155 candidate = fence_match.group(1).strip() 

156 try: 

157 return _json_loads(candidate) 

158 except (ValueError, _JSONDecodeError): 

159 pass # Fall through to other strategies 

160 

161 # Strategy 2: direct parse 

162 try: 

163 return _json_loads(stripped) 

164 except (ValueError, _JSONDecodeError): 

165 pass 

166 

167 # Strategy 3: bracket-counting extraction — correctly handles nested 

168 # structures, multiple JSON objects, and braces inside string values. 

169 block = _extract_first_json_block(stripped) 

170 if block is not None: 

171 return _json_loads(block) 

172 

173 raise ValueError(f"No valid JSON found in response: {text[:200]!r}") 

174 

175 

176def validate_against_model(data: Any, output_model: type[Any]) -> Any: 

177 """Validate *data* and construct an instance of *output_model*. 

178 

179 Supports: 

180 

181 - **Pydantic v2** ``BaseModel`` subclasses — uses ``model_validate()``. 

182 - **Pydantic v1** ``BaseModel`` subclasses — uses ``parse_obj()``. 

183 - **Dataclasses** — constructs via ``output_model(**data)`` after basic 

184 key filtering. 

185 - **Dict-only** round-trip — when a plain ``dict`` is acceptable, returns 

186 *data* after schema check if *output_model* is ``dict``. 

187 

188 Args: 

189 data: Parsed JSON value (usually a dict). 

190 output_model: Target model class. 

191 

192 Returns: 

193 A validated instance of *output_model*. 

194 

195 Raises: 

196 TypeError: When *data* is not a ``dict`` and the model requires one. 

197 ValueError: When validation fails for any reason. 

198 """ 

199 if output_model is dict: 

200 if not isinstance(data, dict): 

201 raise TypeError(f"Expected dict, got {type(data).__name__}") 

202 return data 

203 

204 # Try Pydantic v2 first (model_validate is Pydantic v2 API) 

205 model_validate = getattr(output_model, "model_validate", None) 

206 if callable(model_validate): 

207 try: 

208 return model_validate(data) 

209 except Exception as exc: 

210 raise ValueError(f"Pydantic validation failed: {exc}") from exc 

211 

212 # Try Pydantic v1 (parse_obj) 

213 parse_obj = getattr(output_model, "parse_obj", None) 

214 if callable(parse_obj): 

215 try: 

216 return parse_obj(data) 

217 except Exception as exc: 

218 raise ValueError(f"Pydantic v1 validation failed: {exc}") from exc 

219 

220 # Fallback: plain constructor with dict unpacking 

221 if not isinstance(data, dict): 

222 raise TypeError( 

223 f"Cannot construct {output_model.__name__!r} from {type(data).__name__}; " 

224 "expected a dict" 

225 ) 

226 try: 

227 return output_model(**data) 

228 except Exception as exc: 

229 raise ValueError( 

230 f"Construction of {output_model.__name__!r} failed: {exc}" 

231 ) from exc 

232 

233 

234def build_json_schema(output_model: type[Any]) -> dict[str, Any]: 

235 """Generate a JSON Schema dict from *output_model*. 

236 

237 Supports Pydantic v2 (``model_json_schema``), Pydantic v1 

238 (``schema``), and dataclasses (via ``dataclasses.fields``). 

239 

240 Args: 

241 output_model: Model class to inspect. 

242 

243 Returns: 

244 JSON Schema dict describing the model's fields. 

245 """ 

246 # Pydantic v2 

247 model_json_schema = getattr(output_model, "model_json_schema", None) 

248 if callable(model_json_schema): 

249 return cast("dict[str, Any]", model_json_schema()) 

250 

251 # Pydantic v1 

252 schema_fn = getattr(output_model, "schema", None) 

253 if callable(schema_fn): 

254 return cast("dict[str, Any]", schema_fn()) 

255 

256 # Dataclass 

257 import dataclasses 

258 

259 if dataclasses.is_dataclass(output_model): 

260 properties: dict[str, Any] = {} 

261 required: list[str] = [] 

262 for f in dataclasses.fields(output_model): 

263 properties[f.name] = {"type": "string"} # simplified 

264 if ( 

265 f.default is dataclasses.MISSING 

266 and f.default_factory is dataclasses.MISSING 

267 ): 

268 required.append(f.name) 

269 schema: dict[str, Any] = { 

270 "type": "object", 

271 "properties": properties, 

272 } 

273 if required: 

274 schema["required"] = required 

275 return schema 

276 

277 # Unknown model — return a permissive schema 

278 return {"type": "object"} 

279 

280 

281class StructuredOutputParser: 

282 """Schema-aware parser that validates LLM responses against a model. 

283 

284 Wraps :func:`extract_json_block`, :func:`validate_against_model`, and 

285 :func:`build_json_schema` into a convenient class-based API. 

286 

287 Args: 

288 output_model: Model class for validation. 

289 strict: Whether to enforce strict validation (default ``True``). 

290 """ 

291 

292 def __init__(self, output_model: type[Any], *, strict: bool = True) -> None: 

293 """Initialise with model class.""" 

294 self._output_model = output_model 

295 self._strict = strict 

296 

297 def parse(self, completion: Any) -> Any: 

298 """Parse and validate a completion into an output_model instance. 

299 

300 Args: 

301 completion: Completion object with ``.content`` attribute, or a string. 

302 

303 Returns: 

304 Validated model instance. 

305 

306 Raises: 

307 ParseError: When JSON cannot be extracted. 

308 SchemaValidationError: When validation fails. 

309 """ 

310 from lexigram.ai.llm.structured.exceptions import ( 

311 ParseError, 

312 SchemaValidationError, 

313 ) 

314 

315 content = ( 

316 completion.content if hasattr(completion, "content") else str(completion) 

317 ) 

318 try: 

319 parsed = extract_json_block(content) 

320 except ValueError as exc: 

321 raise ParseError(str(exc)) from exc 

322 try: 

323 return validate_against_model(parsed, self._output_model) 

324 except (TypeError, ValueError) as exc: 

325 raise SchemaValidationError(str(exc)) from exc 

326 

327 def parse_array(self, completion: Any) -> list[Any]: 

328 """Parse and validate an array of output_model instances. 

329 

330 Args: 

331 completion: Completion object with ``.content`` attribute. 

332 

333 Returns: 

334 List of validated model instances. 

335 

336 Raises: 

337 ParseError: When JSON is not an array. 

338 SchemaValidationError: When validation fails. 

339 """ 

340 from lexigram.ai.llm.structured.exceptions import ( 

341 ParseError, 

342 SchemaValidationError, 

343 ) 

344 

345 content = ( 

346 completion.content if hasattr(completion, "content") else str(completion) 

347 ) 

348 try: 

349 parsed = extract_json_block(content) 

350 except ValueError as exc: 

351 raise ParseError(str(exc)) from exc 

352 if not isinstance(parsed, list): 

353 raise ParseError(f"Expected JSON array, got {type(parsed).__name__}") 

354 results = [] 

355 for item in parsed: 

356 try: 

357 results.append(validate_against_model(item, self._output_model)) 

358 except (TypeError, ValueError) as exc: 

359 raise SchemaValidationError(str(exc)) from exc 

360 return results 

361 

362 def get_json_schema(self) -> dict[str, Any]: 

363 """Return JSON Schema dict for the output model.""" 

364 return build_json_schema(self._output_model) 

365 

366 def get_schema_prompt(self) -> str: 

367 """Return a human-readable schema prompt string.""" 

368 from lexigram.serialization import dumps_str 

369 

370 schema = self.get_json_schema() 

371 return dumps_str(schema, indent=2) 

372 

373 

374__all__ = [ 

375 "StructuredOutputParser", 

376 "build_json_schema", 

377 "extract_json_block", 

378 "validate_against_model", 

379]