Coverage for haystack/components/extractors/image/llm_document_content_extractor.py: 100%
153 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import asyncio
6import json
7from concurrent.futures import ThreadPoolExecutor
8from dataclasses import replace
9from typing import Any, Literal
11from jinja2 import meta
12from jinja2.sandbox import SandboxedEnvironment
14from haystack import Document, component, default_from_dict, default_to_dict, logging, tracing
15from haystack.components.converters.image.document_to_image import DocumentToImageContent
16from haystack.components.generators.chat.types import ChatGenerator
17from haystack.components.generators.utils import _trace_chat_generator_run
18from haystack.core.serialization import component_to_dict
19from haystack.dataclasses import ImageContent, TextContent
20from haystack.dataclasses.chat_message import ChatMessage
21from haystack.utils import deserialize_chatgenerator_inplace
22from haystack.utils.async_utils import _execute_component_async
23from haystack.utils.misc import _parse_dict_from_json
25logger = logging.getLogger(__name__)
28# Reserved key in the LLM JSON response that holds the main document text.
29DOCUMENT_CONTENT_KEY = "document_content"
32DEFAULT_PROMPT_TEMPLATE = """
33You are part of an information extraction pipeline that extracts the content of image-based documents.
35Extract the content from the provided image.
36You need to extract the content exactly.
37Format everything as markdown.
38Make sure to retain the reading order of the document.
40**Visual Elements**
41Do not extract figures, drawings, maps, graphs or any other visual elements.
42Instead, add a caption that describes briefly what you see in the visual element.
43You must describe each visual element.
44If you only see a visual element without other content, you must describe this visual element.
45Enclose each image caption with [img-caption][/img-caption]
47**Tables**
48Make sure to format the table in markdown.
49Add a short caption below the table that describes the table's content.
50Enclose each table caption with [table-caption][/table-caption].
51The caption must be placed below the extracted table.
53**Forms**
54Reproduce checkbox selections with markdown.
56Return a single JSON object. It must contain the key "document_content" with the extracted text as value.
58No markdown, no code fence, only raw JSON.
60Document:"""
63@component
64class LLMDocumentContentExtractor:
65 """
66 Extracts textual content and optionally metadata from image-based documents using a vision-enabled LLM.
68 One prompt and one LLM call per document. The component converts each document to an image via
69 DocumentToImageContent and sends it to the ChatGenerator. The prompt must not contain Jinja variables.
71 Response handling:
72 - If the LLM returns a **plain string** (non-JSON or not a JSON object), it is written to the document's content.
73 - If the LLM returns a **JSON object with only the key** `document_content`, that value is written to content.
74 - If the LLM returns a **JSON object with multiple keys**, the value of `document_content` (if present) is
75 written to content and all other keys are merged into the document's metadata.
77 The ChatGenerator can be configured to return JSON. For example, with the OpenAIResponsesChatGenerator, pass
78 `{"text": {"format": {"type": "json_schema", ...}}}` in `generation_kwargs` as shown in the usage example below.
80 Documents that fail extraction are returned in `failed_documents` with the error in their `extraction_error`
81 metadata field.
83 ### Usage example
85 ```python
86 from haystack import Document
87 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
88 from haystack.components.extractors.image import LLMDocumentContentExtractor
90 prompt = \"\"\"
91 Extract the content from the provided image.
92 Format everything as markdown. Return only the extracted content as a JSON object with the key 'document_content'.
93 No markdown, no code fence, only raw JSON.
95 Extract metadata about the image like source of the image, date of creation, etc. if you can.
96 Return this metadata as additional key-value pairs in the same JSON object.
97 \"\"\"
99 chat_generator = OpenAIResponsesChatGenerator(
100 generation_kwargs={
101 "text": {
102 "format": {
103 "type": "json_schema",
104 "name": "entity_extraction",
105 "strict": False,
106 "schema": {
107 "type": "object",
108 "properties": {
109 "document_content": {"type": "string"},
110 "author": {"type": "string"},
111 "date": {"type": "string"},
112 "document_type": {"type": "string"},
113 "title": {"type": "string"},
114 },
115 "additionalProperties": False,
116 },
117 },
118 }
119 }
120 )
122 extractor = LLMDocumentContentExtractor(
123 chat_generator=chat_generator,
124 file_path_meta_field="file_path",
125 raise_on_failure=False
126 )
128 documents = [
129 Document(content="", meta={"file_path": "test/test_files/images/image_metadata.png"}),
130 Document(content="", meta={"file_path": "test/test_files/images/apple.jpg", "page_number": 1})
131 ]
132 result = extractor.run(documents=documents)
133 updated_documents = result["documents"]
134 ```
135 """
137 def __init__(
138 self,
139 *,
140 chat_generator: ChatGenerator,
141 prompt: str = DEFAULT_PROMPT_TEMPLATE,
142 file_path_meta_field: str = "file_path",
143 root_path: str | None = None,
144 detail: Literal["auto", "high", "low"] | None = None,
145 size: tuple[int, int] | None = None,
146 raise_on_failure: bool = False,
147 max_workers: int = 3,
148 ) -> None:
149 """
150 Initialize the LLMDocumentContentExtractor component.
152 :param chat_generator: A ChatGenerator that supports vision input. Optionally configured for JSON
153 (e.g. `{"text": {"format": {"type": "json_schema", ...}}}` in `generation_kwargs` when using the
154 OpenAIResponsesChatGenerator).
155 :param prompt: Prompt for extraction. Must not contain Jinja variables.
156 :param file_path_meta_field: The metadata field in the Document that contains the file path to the image or PDF.
157 :param root_path: The root directory path where document files are located. If provided, file paths in
158 document metadata will be resolved relative to this path and are guaranteed to stay within it. If None,
159 file paths are treated as absolute paths with no containment check.
160 Security: this component reads the file referenced by `file_path_meta_field` from the host filesystem. If
161 document metadata may be influenced by untrusted input, set `root_path` to a dedicated data directory so
162 that path-traversal payloads (e.g. absolute paths or `../`) are rejected instead of read.
163 :param detail: Optional detail level of the image (only supported by OpenAI). Can be "auto", "high", or "low".
164 :param size: If provided, resizes the image to fit within (width, height) while keeping aspect ratio.
165 :param raise_on_failure: If True, exceptions from the LLM are raised. If False, documents that fail extraction
166 are returned in `failed_documents` with the error in their `extraction_error` metadata field.
167 :param max_workers: Maximum number of threads for parallel LLM calls.
168 """
169 self._chat_generator = chat_generator
170 self.prompt = prompt
171 self.file_path_meta_field = file_path_meta_field
172 self.root_path = root_path or ""
173 self.detail = detail
174 self.size = size
175 LLMDocumentContentExtractor._validate_prompt_no_variables(prompt)
176 self.raise_on_failure = raise_on_failure
177 self.max_workers = max_workers
178 self._document_to_image_content = DocumentToImageContent(
179 file_path_meta_field=file_path_meta_field, root_path=root_path, detail=detail, size=size
180 )
182 def warm_up(self) -> None:
183 """
184 Warm up the underlying chat generator.
185 """
186 if hasattr(self._chat_generator, "warm_up"):
187 self._chat_generator.warm_up()
189 async def warm_up_async(self) -> None:
190 """
191 Warm up the underlying chat generator on the serving event loop.
192 """
193 if hasattr(self._chat_generator, "warm_up_async"):
194 await self._chat_generator.warm_up_async()
195 elif hasattr(self._chat_generator, "warm_up"):
196 self._chat_generator.warm_up()
198 def close(self) -> None:
199 """
200 Release the underlying chat generator's resources.
201 """
202 if hasattr(self._chat_generator, "close"):
203 self._chat_generator.close()
205 async def close_async(self) -> None:
206 """
207 Release the underlying chat generator's async resources.
208 """
209 if hasattr(self._chat_generator, "close_async"):
210 await self._chat_generator.close_async()
211 elif hasattr(self._chat_generator, "close"):
212 self._chat_generator.close()
214 def to_dict(self) -> dict[str, Any]:
215 """
216 Serializes the component to a dictionary.
218 :returns:
219 Dictionary with serialized data.
220 """
221 return default_to_dict(
222 self,
223 chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"),
224 prompt=self.prompt,
225 file_path_meta_field=self.file_path_meta_field,
226 root_path=self.root_path,
227 detail=self.detail,
228 size=self.size,
229 raise_on_failure=self.raise_on_failure,
230 max_workers=self.max_workers,
231 )
233 @classmethod
234 def from_dict(cls, data: dict[str, Any]) -> "LLMDocumentContentExtractor":
235 """
236 Deserializes the component from a dictionary.
238 :param data:
239 Dictionary with serialized data.
240 :returns:
241 An instance of the component.
242 """
243 init_params = data.get("init_parameters", {})
244 deserialize_chatgenerator_inplace(init_params, key="chat_generator")
246 return default_from_dict(cls, data)
248 @staticmethod
249 def _validate_prompt_no_variables(prompt: str) -> None:
250 ast = SandboxedEnvironment().parse(prompt)
251 template_variables = meta.find_undeclared_variables(ast)
252 variables = list(template_variables)
253 if variables:
254 raise ValueError(
255 f"The prompt must not have any variables, only instructions on how to extract the content of the "
256 f"the image-based document. Found {','.join(variables)} in the prompt."
257 )
259 @staticmethod
260 def _process_response(response_text: str) -> tuple[str | None, dict[str, Any], str | None]:
261 """
262 Parse LLM response. Returns (content, meta_updates, error).
264 - Plain string (non-JSON): use entire response as document content;
265 - Valid JSON object: use key `document_content` for Document.content and all other keys for Document.metadata;
266 - Valid JSON but not an object (e.g. array or primitive), report an error;
267 """
268 try:
269 parsed = _parse_dict_from_json(response_text, raise_on_failure=True)
270 except json.JSONDecodeError:
271 return response_text, {}, None
272 except ValueError:
273 return None, {}, "Response must be a JSON object, not an array or primitive."
275 content = parsed.get(DOCUMENT_CONTENT_KEY)
276 meta_updates = {k: v for k, v in parsed.items() if k != DOCUMENT_CONTENT_KEY}
277 return content, meta_updates, None
279 @staticmethod
280 def _fail(document: Document, error: str) -> tuple[Document, bool]:
281 """Return a copy of `document` with `extraction_error` set, flagged as failed."""
282 return replace(document, meta={**document.meta, "extraction_error": error}), False
284 def _run_on_thread(
285 self, document: Document, image_content: ImageContent | None, parent_span: tracing.Span | None = None
286 ) -> tuple[Document, bool]:
287 """
288 Execute the LLM inference in a separate thread for each document.
290 :param document: The document to extract content for.
291 :param image_content: The image content for the document, or None if conversion failed.
292 :param parent_span: Span to nest the generator span under, captured on the calling thread.
293 :returns:
294 The updated document and True on success, or the document with failure metadata and False.
295 """
296 if image_content is None:
297 return self._fail(document, "Document could not be converted to an image, skipping LLM call.")
299 # the prompt is the same for all documents, so we can set it up once here for each document/thread
300 message = ChatMessage.from_user(content_parts=[TextContent(text=self.prompt), image_content])
302 try:
303 with _trace_chat_generator_run(
304 self._chat_generator, {"messages": [message]}, parent_span=parent_span
305 ) as span:
306 result = self._chat_generator.run(messages=[message])
307 span.set_content_tag("haystack.component.output", result)
308 except Exception as e:
309 if self.raise_on_failure:
310 raise e
311 logger.exception(
312 "LLMDocumentContentExtractor failed during chat generation. Returning the document in "
313 "failed_documents. Error: {error}",
314 error=e,
315 )
316 return self._fail(document, "LLM failed with exception: " + str(e))
318 return self._process_llm_results(document, result["replies"][0])
320 async def _run_async(
321 self, document: Document, image_content: ImageContent | None, parent_span: tracing.Span | None = None
322 ) -> tuple[Document, bool]:
323 """
324 Execute the LLM inference asynchronously for each document.
326 :param document: The document to extract content for.
327 :param image_content: The image content for the document, or None if conversion failed.
328 :param parent_span: Span to nest the generator span under, captured on the calling task.
329 :returns:
330 The updated document and True on success, or the document with failure metadata and False.
331 """
332 if image_content is None:
333 return self._fail(document, "Document could not be converted to an image, skipping LLM call.")
335 # the prompt is the same for all documents, so we can set it up once here for each document
336 message = ChatMessage.from_user(content_parts=[TextContent(text=self.prompt), image_content])
338 try:
339 with _trace_chat_generator_run(
340 self._chat_generator, {"messages": [message]}, parent_span=parent_span
341 ) as span:
342 result = await _execute_component_async(self._chat_generator, messages=[message])
343 span.set_content_tag("haystack.component.output", result)
344 except Exception as e:
345 if self.raise_on_failure:
346 raise e
347 logger.exception(
348 "LLMDocumentContentExtractor failed during chat generation. Returning the document in "
349 "failed_documents. Error: {error}",
350 error=e,
351 )
352 return self._fail(document, "LLM failed with exception: " + str(e))
354 return self._process_llm_results(document, result["replies"][0])
356 @staticmethod
357 def _process_llm_results(document: Document, reply: ChatMessage) -> tuple[Document, bool]:
358 """
359 Process one document's LLM reply using the unified response logic.
361 Returns (updated_document, True if success else False).
362 """
363 # remove potentially existing error metadata from previous runs
364 new_meta = {**document.meta}
365 new_meta.pop("extraction_error", None)
367 # process the LLM response considering the possible response formats
368 content, meta_updates, error = LLMDocumentContentExtractor._process_response(reply.text or "")
370 if error:
371 logger.warning(
372 "LLMDocumentContentExtractor couldn't use the ChatGenerator response for document {document_id}. "
373 "Returning the document in failed_documents. Error: {error}",
374 document_id=document.id,
375 error=error,
376 )
377 new_meta["extraction_error"] = error
378 return replace(document, meta=new_meta), False
380 new_meta.update(meta_updates)
381 final_content = document.content if content is None else content
382 return replace(document, content=final_content, meta=new_meta), True
384 @component.output_types(documents=list[Document], failed_documents=list[Document])
385 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
386 """
387 Run extraction on image-based documents. One LLM call per document.
389 :param documents: A list of image-based documents to process. Each must have a valid file path in its metadata.
390 :returns:
391 A dictionary with "documents" (successfully processed) and "failed_documents" (with failure metadata).
392 """
393 if not documents:
394 return {"documents": [], "failed_documents": []}
396 self.warm_up()
398 image_contents = self._document_to_image_content.run(documents=documents)["image_contents"]
400 # Capture the current span here so worker threads nest their generator spans under the component span.
401 parent_span = tracing.tracer.current_span()
403 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
404 results = executor.map(
405 lambda document, image_content: self._run_on_thread(document, image_content, parent_span=parent_span),
406 documents,
407 image_contents,
408 )
410 successful_documents = []
411 failed_documents = []
412 for doc, success in results:
413 if success:
414 successful_documents.append(doc)
415 else:
416 failed_documents.append(doc)
418 return {"documents": successful_documents, "failed_documents": failed_documents}
420 @component.output_types(documents=list[Document], failed_documents=list[Document])
421 async def run_async(self, documents: list[Document]) -> dict[str, list[Document]]:
422 """
423 Asynchronously run extraction on image-based documents. One LLM call per document.
425 This is the asynchronous version of the `run` method. It has the same parameters and return values
426 but can be used with `await` in an async code. LLM calls are made concurrently, bounded by `max_workers`.
427 If the chat generator only implements a synchronous `run` method, it is executed in a thread to avoid
428 blocking the event loop.
430 :param documents: A list of image-based documents to process. Each must have a valid file path in its metadata.
431 :returns:
432 A dictionary with "documents" (successfully processed) and "failed_documents" (with failure metadata).
433 """
434 if not documents:
435 return {"documents": [], "failed_documents": []}
437 await self.warm_up_async()
439 # Reading the files and rendering PDF pages is blocking work and `DocumentToImageContent` has no
440 # `run_async`, so it runs in a thread instead of on the event loop.
441 conversion_result = await _execute_component_async(self._document_to_image_content, documents=documents)
442 image_contents = conversion_result["image_contents"]
444 # Capture the current span here so concurrent tasks nest their generator spans under the component span.
445 parent_span = tracing.tracer.current_span()
447 # Run the LLM on each image content, bounding concurrency per task so max_workers is enforced.
448 sem = asyncio.Semaphore(max(1, self.max_workers))
450 async def _bounded_run(document: Document, image_content: ImageContent | None) -> tuple[Document, bool]:
451 async with sem:
452 return await self._run_async(document, image_content, parent_span=parent_span)
454 results = await asyncio.gather(
455 *[
456 _bounded_run(document, image_content)
457 for document, image_content in zip(documents, image_contents, strict=True)
458 ]
459 )
461 successful_documents = []
462 failed_documents = []
463 for doc, success in results:
464 if success:
465 successful_documents.append(doc)
466 else:
467 failed_documents.append(doc)
469 return {"documents": successful_documents, "failed_documents": failed_documents}