Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/structured/extractor.py: 17%
145 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Structured output extraction from LLM responses.
3:class:`StructuredExtractor` bridges the gap between raw LLM text and
4typed, validated application objects. It wraps any :class:`LLMClientProtocol`,
5instructs the model to respond as JSON, parses the response, and
6validates it against a caller-provided model class.
8Example::
10 from pydantic import BaseModel
11 from typing import Literal
13 class Sentiment(BaseModel):
14 label: Literal["positive", "negative", "neutral"]
15 confidence: float
16 reasoning: str
18 extractor = StructuredExtractor(llm_client)
19 result = await extractor.extract(
20 prompt="Analyze the sentiment of: 'Great product!'",
21 output_model=Sentiment,
22 )
23 if result.is_ok():
24 analysis = result.unwrap()
25 print(analysis.label, analysis.confidence)
26 else:
27 error = result.unwrap_err()
28 # handle ExtractionParseError / ExtractionValidationError / ExtractionMaxRetriesError
29"""
31from __future__ import annotations
33from typing import Any, TypeVar
35from lexigram.ai.llm.exceptions import (
36 ExtractionMaxRetriesError,
37 ExtractionParseError,
38 ExtractionValidationError,
39 LLMError,
40)
41from lexigram.ai.llm.structured.parser import (
42 build_json_schema,
43 extract_json_block,
44 validate_against_model,
45)
46from lexigram.contracts.ai.llm import ChatMessage, LLMClientProtocol, Role
47from lexigram.logging import (
48 get_logger,
49)
50from lexigram.result import Err, Ok, Result
51from lexigram.serialization import JSONDecodeError, dumps_str, loads
53logger = get_logger(__name__)
55T = TypeVar("T")
57# ---------------------------------------------------------------------------
58# System prompt template injected before user content
59# ---------------------------------------------------------------------------
61_EXTRACTION_SYSTEM_PROMPT = """\
62You are a precise data extraction assistant.
64Your task is to analyse the user's input and return a **single JSON object** \
65that exactly matches the schema provided below.
67Schema:
68{schema}
70Rules:
71- Output ONLY valid JSON. No markdown, no prose, no code fences.
72- Every required field must be present.
73- Do NOT add extra fields not in the schema.
74- If a value cannot be determined from the input, use a sensible default \
75 (empty string, 0, false, etc.) rather than omitting the field.
76"""
79class JSONExtractor:
80 """Extract and parse JSON from LLM responses."""
82 @staticmethod
83 def extract(text: str, multiple: bool = False) -> Any:
84 """Extract JSON from text."""
85 import re
87 # Try to extract from code blocks first
88 json_blocks = re.findall(r"```(?:json)?\s*(.*?)```", text, re.DOTALL)
90 if json_blocks:
91 if multiple:
92 results = []
93 for block in json_blocks:
94 try:
95 results.append(loads(block.strip()))
96 except (ValueError, TypeError, JSONDecodeError):
97 continue
98 if results:
99 return results
100 else:
101 for block in json_blocks:
102 try:
103 return loads(block.strip())
104 except (ValueError, TypeError, JSONDecodeError):
105 continue
107 # Try parsing entire text
108 try:
109 result = loads(text.strip())
110 return [result] if multiple else result
111 except (ValueError, TypeError, JSONDecodeError):
112 pass
114 # Try to find JSON objects in text using bracket-counting scanner
115 from lexigram.ai.llm.extraction._json_scanner import extract_json_objects
117 matches = extract_json_objects(text)
119 if matches:
120 if multiple:
121 results = []
122 for match in matches:
123 try:
124 results.append(loads(match))
125 except (ValueError, TypeError, JSONDecodeError):
126 continue
127 if results:
128 return results
129 else:
130 for match in matches:
131 try:
132 return loads(match)
133 except (ValueError, TypeError, JSONDecodeError):
134 pass
136 from lexigram.ai.llm.structured.exceptions import ParseError
138 raise ParseError("Cannot extract JSON from text")
140 @staticmethod
141 def extract_array(text: str) -> list[Any]:
142 """Extract JSON array from text."""
143 from lexigram.ai.llm.structured.exceptions import ParseError
145 result = JSONExtractor.extract(text)
146 if not isinstance(result, list):
147 raise ParseError(f"Expected JSON array, got {type(result)}")
148 return result
151class StructuredExtractor:
152 """Extracts validated, typed data from an LLM's text response.
154 The extractor wraps a :class:`~lexigram.contracts.ai.protocols.LLMClientProtocol`
155 (any object implementing :meth:`complete`), builds a JSON-focused system
156 prompt derived from *output_model*'s schema, calls the LLM, and validates
157 the result.
159 Failed parse or validation attempts consume retry budget. All outcomes
160 are surfaced as :class:`~lexigram.result.Result` — the
161 caller never receives a raw exception.
163 Args:
164 client: Any ``LLMClientProtocol`` implementation.
165 default_model: Optional model name to use when :meth:`extract` is
166 called without an explicit ``model`` argument.
167 default_max_retries: Default number of extra attempts after the first
168 failure. Individual :meth:`extract` calls can
169 override this.
171 Example::
173 extractor = StructuredExtractor(openai_client, default_model="gpt-4o")
174 result = await extractor.extract(
175 prompt="Extract key fields from this invoice: ...",
176 output_model=Invoice,
177 )
178 """
180 def __init__(
181 self,
182 client: LLMClientProtocol,
183 *,
184 default_model: str | None = None,
185 default_max_retries: int = 2,
186 ) -> None:
187 """Initialise the extractor.
189 Args:
190 client: LLM client to use for completions.
191 default_model: Default model name; forwarded as ``model`` kwarg
192 to the client when no ``model`` is given in
193 :meth:`extract`.
194 default_max_retries: Default retry budget per extraction request.
195 """
196 self._client = client
197 self._default_model = default_model
198 self._default_max_retries = default_max_retries
200 async def extract(
201 self,
202 prompt: str | list[Any],
203 output_model: type[T],
204 *,
205 max_retries: int | None = None,
206 model: str | None = None,
207 **kwargs: Any,
208 ) -> Result[
209 T,
210 ExtractionParseError
211 | ExtractionValidationError
212 | ExtractionMaxRetriesError
213 | LLMError,
214 ]:
215 """Extract a validated instance of *output_model* from an LLM response.
217 Args:
218 prompt: User prompt text, or a list of chat message dicts/objects.
219 A text string is converted to a single ``user`` message.
220 output_model: Model class to validate against. Pydantic models,
221 dataclasses, and classes with a ``model_validate``
222 class method are all supported.
223 max_retries: Override the instance-level retry budget for this
224 call. ``0`` means a single attempt with no retries.
225 model: Override the instance-level default model for this call.
226 **kwargs: Additional keyword arguments forwarded verbatim to the
227 underlying :meth:`~LLMClientProtocol.complete` call.
229 Returns:
230 ``Ok(T)`` — a fully validated *output_model* instance.
231 ``Err(ExtractionParseError)`` — the last response was not valid JSON.
232 ``Err(ExtractionValidationError)`` — JSON was valid but did not satisfy
233 the model's schema.
234 ``Err(ExtractionMaxRetriesError)`` — all attempts were exhausted.
235 ``Err(LLMError)`` — the underlying client returned a recoverable
236 provider/model failure.
237 """
238 retries = max_retries if max_retries is not None else self._default_max_retries
239 if retries < 0:
240 raise ValueError("max_retries must be greater than or equal to 0")
241 effective_model = model or self._default_model
243 schema = build_json_schema(output_model)
244 schema_str = dumps_str(schema, indent=2)
245 system_content = _EXTRACTION_SYSTEM_PROMPT.format(schema=schema_str)
247 messages = self._build_messages(prompt, system_content)
248 if effective_model is not None:
249 kwargs.setdefault("model", effective_model)
251 last_error: ExtractionParseError | ExtractionValidationError | None = None
253 for attempt in range(retries + 1):
254 if attempt > 0:
255 logger.debug(
256 "extraction_retry",
257 attempt=attempt,
258 model=effective_model,
259 output_model=output_model.__name__,
260 )
262 completion_result = await self._client.complete(messages, **kwargs)
263 if completion_result.is_err():
264 return Err(completion_result.unwrap_err()) # type: ignore[arg-type]
266 completion = completion_result.unwrap()
267 raw_text = _get_content(completion)
269 # --- Parse ---
270 try:
271 parsed = extract_json_block(raw_text)
272 except ValueError as exc:
273 last_error = ExtractionParseError(
274 f"Failed to parse JSON from LLM response (attempt {attempt + 1}): {exc}"
275 )
276 logger.debug(
277 "extraction_parse_failed",
278 attempt=attempt + 1,
279 error=str(exc),
280 raw_text=raw_text[:200],
281 )
282 continue
284 # --- Validate ---
285 try:
286 instance = validate_against_model(parsed, output_model)
287 except (TypeError, ValueError) as exc:
288 last_error = ExtractionValidationError(
289 f"Model validation failed (attempt {attempt + 1}): {exc}"
290 )
291 logger.debug(
292 "extraction_validation_failed",
293 attempt=attempt + 1,
294 error=str(exc),
295 )
296 continue
298 logger.debug(
299 "extraction_succeeded",
300 attempt=attempt + 1,
301 output_model=output_model.__name__,
302 )
303 return Ok(instance)
305 # All attempts exhausted
306 if last_error is not None:
307 return Err(
308 ExtractionMaxRetriesError(
309 f"Extraction failed after {retries + 1} attempt(s). "
310 f"Last error: {last_error}"
311 )
312 )
314 # Should not be reached
315 return Err(
316 ExtractionMaxRetriesError("Extraction failed with no recorded error.")
317 )
319 # ------------------------------------------------------------------
320 # Helpers
321 # ------------------------------------------------------------------
323 @staticmethod
324 def _build_messages(
325 prompt: str | list[Any],
326 system_content: str,
327 ) -> list[ChatMessage]:
328 """Build a message list with the extraction system prompt prepended.
330 Args:
331 prompt: User prompt string or pre-built message list.
332 system_content: The JSON-extraction system prompt.
334 Returns:
335 List of ``ChatMessage`` objects.
336 """
337 if isinstance(prompt, str):
338 return [
339 ChatMessage(role=Role.SYSTEM, content=system_content),
340 ChatMessage(role=Role.USER, content=prompt),
341 ]
343 messages: list[ChatMessage] = []
344 has_system = False
345 for msg in prompt:
346 role = (
347 msg.get("role") if isinstance(msg, dict) else getattr(msg, "role", None)
348 )
349 content = (
350 msg.get("content", "")
351 if isinstance(msg, dict)
352 else getattr(msg, "content", "")
353 )
354 if role == "system":
355 has_system = True
356 combined = f"{content}\n\n---\n\n{system_content}"
357 messages.append(ChatMessage(role=Role.SYSTEM, content=combined))
358 elif isinstance(content, list):
359 messages.append(ChatMessage(role=str(role), content=content))
360 else:
361 messages.append(ChatMessage(role=str(role), content=str(content)))
363 if not has_system:
364 messages.insert(0, ChatMessage(role=Role.SYSTEM, content=system_content))
366 return messages
369def _get_content(completion: Any) -> str:
370 """Extract text content from a Completion-like object.
372 Accepts:
373 - Objects with a ``.content`` attribute (``lexigram-ai-llm`` Completion).
374 - Plain strings.
375 - Dicts with a ``content`` key.
377 Args:
378 completion: Completion object or raw value.
380 Returns:
381 The text content as a string.
383 Raises:
384 ValueError: When content cannot be found.
385 """
386 if isinstance(completion, str):
387 return completion
388 if isinstance(completion, dict):
389 content = completion.get("content") or completion.get("text", "")
390 return str(content)
391 content = getattr(completion, "content", None)
392 if content is not None:
393 return str(content)
394 raise ValueError(
395 f"Cannot extract text content from completion of type {type(completion).__name__!r}"
396 )
399class StructuredOutputExtractor:
400 """Utility class that exposes the bracket-counting JSON block extractor.
402 A lightweight, argument-free companion to :class:`StructuredExtractor`
403 designed for direct use and unit-testing of the JSON extraction step in
404 isolation. The extraction logic itself lives in
405 :func:`~lexigram.ai.llm.structured.parser._extract_first_json_block`;
406 this class adds markdown-fence stripping on top and provides a stable
407 surface for callers that only need the raw JSON string (not a parsed
408 Python value).
410 Example::
412 extractor = StructuredOutputExtractor()
413 raw = extractor._extract_json_block(llm_response_text)
414 if raw is not None:
415 data = json.loads(raw)
416 """
418 def _extract_json_block(self, text: str) -> str | None:
419 """Extract the first complete JSON object or array from *text*.
421 Uses bracket counting (via
422 :func:`~lexigram.ai.llm.structured.parser._extract_first_json_block`)
423 which correctly handles:
425 - Nested objects and arrays.
426 - Multiple JSON objects in the same string (returns the first valid one).
427 - Braces / brackets inside quoted string values.
428 - Markdown code fences (````json … ````).
430 Args:
431 text: Raw text that may contain a JSON object or array.
433 Returns:
434 The first complete, parseable JSON block as a raw string, or
435 ``None`` when no valid JSON block is found.
436 """
437 from lexigram.ai.llm.structured.parser import _extract_first_json_block
439 # Strip markdown code fences before scanning so the bracket counter
440 # is not confused by the fence syntax.
441 stripped = text.strip()
442 if stripped.startswith("```"):
443 lines = stripped.split("\n")
444 # Drop the opening fence line (```json or ```) …
445 inner_lines = lines[1:]
446 # … and the closing ``` if present.
447 if inner_lines and inner_lines[-1].strip() == "```":
448 inner_lines = inner_lines[:-1]
449 stripped = "\n".join(inner_lines)
451 return _extract_first_json_block(stripped)
454# Backward-compat alias
455__all__ = ["JSONExtractor", "StructuredExtractor", "StructuredOutputExtractor"]