Coverage for haystack/components/extractors/image/llm_document_content_extractor.py: 99%
156 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +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 functools import partial
10from typing import Any, Literal
12from jinja2 import meta
13from jinja2.sandbox import SandboxedEnvironment
15from haystack import Document, component, default_from_dict, default_to_dict, logging, tracing
16from haystack.components.converters.image.document_to_image import DocumentToImageContent
17from haystack.components.generators.chat.types import ChatGenerator
18from haystack.components.generators.utils import _trace_chat_generator_run
19from haystack.core.serialization import component_to_dict
20from haystack.dataclasses import ImageContent, TextContent
21from haystack.dataclasses.chat_message import ChatMessage
22from haystack.utils import deserialize_chatgenerator_inplace
23from haystack.utils.async_utils import _execute_component_async
24from haystack.utils.misc import _parse_dict_from_json
26logger = logging.getLogger(__name__)
29# Reserved key in the LLM JSON response that holds the main document text.
30DOCUMENT_CONTENT_KEY = "document_content"
33DEFAULT_PROMPT_TEMPLATE = """
34You are part of an information extraction pipeline that extracts the content of image-based documents.
36Extract the content from the provided image.
37You need to extract the content exactly.
38Format everything as markdown.
39Make sure to retain the reading order of the document.
41**Visual Elements**
42Do not extract figures, drawings, maps, graphs or any other visual elements.
43Instead, add a caption that describes briefly what you see in the visual element.
44You must describe each visual element.
45If you only see a visual element without other content, you must describe this visual element.
46Enclose each image caption with [img-caption][/img-caption]
48**Tables**
49Make sure to format the table in markdown.
50Add a short caption below the table that describes the table's content.
51Enclose each table caption with [table-caption][/table-caption].
52The caption must be placed below the extracted table.
54**Forms**
55Reproduce checkbox selections with markdown.
57Return a single JSON object. It must contain the key "document_content" with the extracted text as value.
59No markdown, no code fence, only raw JSON.
61Document:"""
64@component
65class LLMDocumentContentExtractor:
66 """
67 Extracts textual content and optionally metadata from image-based documents using a vision-enabled LLM.
69 One prompt and one LLM call per document. The component converts each document to an image via
70 DocumentToImageContent and sends it to the ChatGenerator. The prompt must not contain Jinja variables.
72 Response handling:
73 - If the LLM returns a **plain string** (non-JSON or not a JSON object), it is written to the document's content.
74 - If the LLM returns a **JSON object with only the key** `document_content`, that value is written to content.
75 - If the LLM returns a **JSON object with multiple keys**, the value of ``document_content`` (if present) is
76 written to content and all other keys are merged into the document's metadata.
78 The ChatGenerator can be configured to return JSON (e.g. ``response_format={"type": "json_object"}``
79 in ``generation_kwargs``).
81 Documents that fail extraction are returned in ``failed_documents`` with ``content_extraction_error`` in metadata.
83 ### Usage example
85 ```python
86 from haystack import Document
87 from haystack.components.generators.chat import OpenAIChatGenerator
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 = OpenAIChatGenerator(
100 generation_kwargs={
101 "response_format": {
102 "type": "json_schema",
103 "json_schema": {
104 "name": "entity_extraction",
105 "schema": {
106 "type": "object",
107 "properties": {
108 "document_content": {"type": "string"},
109 "author": {"type": "string"},
110 "date": {"type": "string"},
111 "document_type": {"type": "string"},
112 "title": {"type": "string"},
113 },
114 "additionalProperties": False,
115 },
116 },
117 }
118 }
119 )
121 extractor = LLMDocumentContentExtractor(
122 chat_generator=chat_generator,
123 file_path_meta_field="file_path",
124 raise_on_failure=False
125 )
127 documents = [
128 Document(content="", meta={"file_path": "test/test_files/images/image_metadata.png"}),
129 Document(content="", meta={"file_path": "test/test_files/images/apple.jpg", "page_number": 1})
130 ]
131 result = extractor.run(documents=documents)
132 updated_documents = result["documents"]
133 ```
134 """
136 def __init__(
137 self,
138 *,
139 chat_generator: ChatGenerator,
140 prompt: str = DEFAULT_PROMPT_TEMPLATE,
141 file_path_meta_field: str = "file_path",
142 root_path: str | None = None,
143 detail: Literal["auto", "high", "low"] | None = None,
144 size: tuple[int, int] | None = None,
145 raise_on_failure: bool = False,
146 max_workers: int = 3,
147 ) -> None:
148 """
149 Initialize the LLMDocumentContentExtractor component.
151 :param chat_generator: A ChatGenerator that supports vision input. Optionally configured for JSON
152 (e.g. ``response_format={"type": "json_object"}`` in ``generation_kwargs``).
153 :param prompt: Prompt for extraction. Must not contain Jinja variables.
154 :param file_path_meta_field: The metadata field in the Document that contains the file path to the image or PDF.
155 :param root_path: The root directory path where document files are located. If provided, file paths in
156 document metadata will be resolved relative to this path and are guaranteed to stay within it. If None,
157 file paths are treated as absolute paths with no containment check.
158 Security: this component reads the file referenced by `file_path_meta_field` from the host filesystem. If
159 document metadata may be influenced by untrusted input, set `root_path` to a dedicated data directory so
160 that path-traversal payloads (e.g. absolute paths or `../`) are rejected instead of read.
161 :param detail: Optional detail level of the image (only supported by OpenAI). Can be "auto", "high", or "low".
162 :param size: If provided, resizes the image to fit within (width, height) while keeping aspect ratio.
163 :param raise_on_failure: If True, exceptions from the LLM are raised. If False, failed documents are returned.
164 :param max_workers: Maximum number of threads for parallel LLM calls.
165 """
166 self._chat_generator = chat_generator
167 self.prompt = prompt
168 self.file_path_meta_field = file_path_meta_field
169 self.root_path = root_path or ""
170 self.detail = detail
171 self.size = size
172 LLMDocumentContentExtractor._validate_prompt_no_variables(prompt)
173 self.raise_on_failure = raise_on_failure
174 self.max_workers = max_workers
175 self._document_to_image_content = DocumentToImageContent(
176 file_path_meta_field=file_path_meta_field, root_path=root_path, detail=detail, size=size
177 )
179 def warm_up(self) -> None:
180 """
181 Warm up the underlying chat generator.
182 """
183 if hasattr(self._chat_generator, "warm_up"):
184 self._chat_generator.warm_up()
186 async def warm_up_async(self) -> None:
187 """
188 Warm up the underlying chat generator on the serving event loop.
189 """
190 if hasattr(self._chat_generator, "warm_up_async"):
191 await self._chat_generator.warm_up_async()
192 elif hasattr(self._chat_generator, "warm_up"):
193 self._chat_generator.warm_up()
195 def close(self) -> None:
196 """
197 Release the underlying chat generator's resources.
198 """
199 if hasattr(self._chat_generator, "close"):
200 self._chat_generator.close()
202 async def close_async(self) -> None:
203 """
204 Release the underlying chat generator's async resources.
205 """
206 if hasattr(self._chat_generator, "close_async"):
207 await self._chat_generator.close_async()
208 elif hasattr(self._chat_generator, "close"):
209 self._chat_generator.close()
211 def to_dict(self) -> dict[str, Any]:
212 """
213 Serializes the component to a dictionary.
215 :returns:
216 Dictionary with serialized data.
217 """
218 return default_to_dict(
219 self,
220 chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"),
221 prompt=self.prompt,
222 file_path_meta_field=self.file_path_meta_field,
223 root_path=self.root_path,
224 detail=self.detail,
225 size=self.size,
226 raise_on_failure=self.raise_on_failure,
227 max_workers=self.max_workers,
228 )
230 @classmethod
231 def from_dict(cls, data: dict[str, Any]) -> "LLMDocumentContentExtractor":
232 """
233 Deserializes the component from a dictionary.
235 :param data:
236 Dictionary with serialized data.
237 :returns:
238 An instance of the component.
239 """
240 init_params = data.get("init_parameters", {})
241 deserialize_chatgenerator_inplace(init_params, key="chat_generator")
243 return default_from_dict(cls, data)
245 @staticmethod
246 def _validate_prompt_no_variables(prompt: str) -> None:
247 ast = SandboxedEnvironment().parse(prompt)
248 template_variables = meta.find_undeclared_variables(ast)
249 variables = list(template_variables)
250 if variables:
251 raise ValueError(
252 f"The prompt must not have any variables, only instructions on how to extract the content of the "
253 f"the image-based document. Found {','.join(variables)} in the prompt."
254 )
256 @staticmethod
257 def _process_response(response_text: str) -> tuple[str | None, dict[str, Any], str | None]:
258 """
259 Parse LLM response. Returns (content, meta_updates, error).
261 - Plain string (non-JSON): use entire response as document content;
262 - Valid JSON object: use key ``document_content`` for Document.content and all other keys for Document.metadata;
263 - Valid JSON but not an object (e.g. array or primitive), report an error;
264 """
265 try:
266 parsed = _parse_dict_from_json(response_text, raise_on_failure=True)
267 except json.JSONDecodeError:
268 return response_text, {}, None
269 except ValueError:
270 return None, {}, "Response must be a JSON object, not an array or primitive."
272 content = parsed.get(DOCUMENT_CONTENT_KEY)
273 meta_updates = {k: v for k, v in parsed.items() if k != DOCUMENT_CONTENT_KEY}
274 return content, meta_updates, None
276 def _run_on_thread(
277 self, image_content: ImageContent | None, parent_span: tracing.Span | None = None
278 ) -> dict[str, Any]:
279 """
280 Execute the LLM inference in a separate thread for each document.
282 :param image_content: The image content for one document, or None if conversion failed.
283 :param parent_span: Span to nest the generator span under, captured on the calling thread.
284 :returns:
285 The LLM response if successful, or a dictionary with an "error" key on failure.
286 """
287 if image_content is None:
288 return {"error": "Document has no content, skipping LLM call."}
290 # the prompt is the same for all documents, so we can set it up once here for each document/thread
291 message = ChatMessage.from_user(content_parts=[TextContent(text=self.prompt), image_content])
293 try:
294 with _trace_chat_generator_run(
295 self._chat_generator, {"messages": [message]}, parent_span=parent_span
296 ) as span:
297 result = self._chat_generator.run(messages=[message])
298 span.set_content_tag("haystack.component.output", result)
299 except Exception as e:
300 if self.raise_on_failure:
301 raise e
302 logger.exception(
303 "LLM {class_name} execution failed. Skipping metadata extraction. Failed with exception '{error}'.",
304 class_name=self._chat_generator.__class__.__name__,
305 error=e,
306 )
307 result = {"error": "LLM failed with exception: " + str(e)}
309 return result
311 async def _run_async(
312 self, image_content: ImageContent | None, parent_span: tracing.Span | None = None
313 ) -> dict[str, Any]:
314 """
315 Execute the LLM inference asynchronously for each document.
317 :param image_content: The image content for one document, or None if conversion failed.
318 :param parent_span: Span to nest the generator span under, captured on the calling task.
319 :returns:
320 The LLM response if successful, or a dictionary with an "error" key on failure.
321 """
322 if image_content is None:
323 return {"error": "Document has no content, skipping LLM call."}
325 # the prompt is the same for all documents, so we can set it up once here for each document
326 message = ChatMessage.from_user(content_parts=[TextContent(text=self.prompt), image_content])
328 try:
329 with _trace_chat_generator_run(
330 self._chat_generator, {"messages": [message]}, parent_span=parent_span
331 ) as span:
332 result = await _execute_component_async(self._chat_generator, messages=[message])
333 span.set_content_tag("haystack.component.output", result)
334 except Exception as e:
335 if self.raise_on_failure:
336 raise e
337 logger.exception(
338 "LLM {class_name} execution failed. Skipping metadata extraction. Failed with exception '{error}'.",
339 class_name=self._chat_generator.__class__.__name__,
340 error=e,
341 )
342 result = {"error": "LLM failed with exception: " + str(e)}
344 return result
346 @staticmethod
347 def _process_llm_results(document: Document, result: dict[str, Any]) -> tuple[Document, bool]:
348 """
349 Process one document's LLM result using the unified response logic.
351 Returns (updated_document, True if success else False).
352 """
353 if "error" in result:
354 new_meta = {**document.meta, "extraction_error": result["error"]}
355 return replace(document, meta=new_meta), False
357 # remove potentially existing error metadata from previous runs
358 new_meta = {**document.meta}
359 new_meta.pop("extraction_error", None)
361 # process the LLM response considering the possible response formats
362 response_text = result["replies"][0].text
363 content, meta_updates, error = LLMDocumentContentExtractor._process_response(response_text)
365 if error:
366 new_meta["extraction_error"] = error
367 return replace(document, meta=new_meta), False
369 new_meta.update(meta_updates)
370 final_content = document.content if content is None else content
371 return replace(document, content=final_content, meta=new_meta), True
373 @component.output_types(documents=list[Document], failed_documents=list[Document])
374 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
375 """
376 Run extraction on image-based documents. One LLM call per document.
378 :param documents: A list of image-based documents to process. Each must have a valid file path in its metadata.
379 :returns:
380 A dictionary with "documents" (successfully processed) and "failed_documents" (with failure metadata).
381 """
382 if not documents:
383 return {"documents": [], "failed_documents": []}
385 self.warm_up()
387 image_contents = self._document_to_image_content.run(documents=documents)["image_contents"]
389 # Capture the current span here so worker threads nest their generator spans under the component span.
390 parent_span = tracing.tracer.current_span()
392 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
393 results = executor.map(partial(self._run_on_thread, parent_span=parent_span), image_contents)
395 successful_documents = []
396 failed_documents = []
397 for document, result in zip(documents, results, strict=True):
398 doc, success = self._process_llm_results(document, result)
399 if success:
400 successful_documents.append(doc)
401 else:
402 failed_documents.append(doc)
404 return {"documents": successful_documents, "failed_documents": failed_documents}
406 @component.output_types(documents=list[Document], failed_documents=list[Document])
407 async def run_async(self, documents: list[Document]) -> dict[str, list[Document]]:
408 """
409 Asynchronously run extraction on image-based documents. One LLM call per document.
411 This is the asynchronous version of the `run` method. It has the same parameters and return values
412 but can be used with `await` in an async code. LLM calls are made concurrently, bounded by `max_workers`.
413 If the chat generator only implements a synchronous `run` method, it is executed in a thread to avoid
414 blocking the event loop.
416 :param documents: A list of image-based documents to process. Each must have a valid file path in its metadata.
417 :returns:
418 A dictionary with "documents" (successfully processed) and "failed_documents" (with failure metadata).
419 """
420 if not documents:
421 return {"documents": [], "failed_documents": []}
423 await self.warm_up_async()
425 # Reading the files and rendering PDF pages is blocking work and `DocumentToImageContent` has no
426 # `run_async`, so it runs in a thread instead of on the event loop.
427 conversion_result = await _execute_component_async(self._document_to_image_content, documents=documents)
428 image_contents = conversion_result["image_contents"]
430 # Capture the current span here so concurrent tasks nest their generator spans under the component span.
431 parent_span = tracing.tracer.current_span()
433 # Run the LLM on each image content, bounding concurrency per task so max_workers is enforced.
434 sem = asyncio.Semaphore(max(1, self.max_workers))
436 async def _bounded_run(image_content: ImageContent | None) -> dict[str, Any]:
437 async with sem:
438 return await self._run_async(image_content, parent_span=parent_span)
440 results = await asyncio.gather(*[_bounded_run(image_content) for image_content in image_contents])
442 successful_documents = []
443 failed_documents = []
444 for document, result in zip(documents, results, strict=True):
445 doc, success = self._process_llm_results(document, result)
446 if success:
447 successful_documents.append(doc)
448 else:
449 failed_documents.append(doc)
451 return {"documents": successful_documents, "failed_documents": failed_documents}