Coverage for haystack/components/extractors/llm_metadata_extractor.py: 93%
166 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 copy
6import json
7from asyncio import Semaphore, gather
8from concurrent.futures import ThreadPoolExecutor
9from dataclasses import replace
10from typing import Any
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.builders import PromptBuilder
17from haystack.components.generators.chat.types import ChatGenerator
18from haystack.components.generators.utils import _trace_chat_generator_run
19from haystack.components.preprocessors import DocumentSplitter
20from haystack.core.serialization import component_to_dict
21from haystack.dataclasses import ChatMessage
22from haystack.utils import deserialize_chatgenerator_inplace, expand_page_range
23from haystack.utils.async_utils import _execute_component_async
24from haystack.utils.misc import _parse_dict_from_json
26logger = logging.getLogger(__name__)
29@component
30class LLMMetadataExtractor:
31 """
32 Extracts metadata from documents using a Large Language Model (LLM).
34 The metadata is extracted by providing a prompt to an LLM that generates the metadata.
36 This component expects as input a list of documents and a prompt. The prompt must have exactly one variable, called
37 `document`, that points to a single document in the list of documents. So to access the content of the document,
38 you can use `{{ document.content }}` in the prompt.
40 The component will run the LLM on each document in the list and extract metadata from the document. The metadata
41 will be added to the document's metadata field. If the LLM fails to extract metadata from a document, the document
42 will be added to the `failed_documents` list. The failed documents will have the keys `metadata_extraction_error` and
43 `metadata_extraction_response` in their metadata. These documents can be re-run with another extractor to
44 extract metadata by using the `metadata_extraction_response` and `metadata_extraction_error` in the prompt.
46 ```python
47 from haystack import Document
48 from haystack.components.extractors.llm_metadata_extractor import LLMMetadataExtractor
49 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
51 NER_PROMPT = '''
52 -Goal-
53 Given text and a list of entity types, identify all entities of those types from the text.
55 -Steps-
56 1. Identify all entities. For each identified entity, extract the following information:
57 - entity: Name of the entity
58 - entity_type: One of the following types: [organization, product, service, industry]
59 Format each entity as a JSON like: {"entity": <entity_name>, "entity_type": <entity_type>}
61 2. Return output in a single list with all the entities identified in steps 1.
63 -Examples-
64 ######################
65 Example 1:
66 entity_types: [organization, person, partnership, financial metric, product, service, industry, investment strategy, market trend]
67 text: Another area of strength is our co-brand issuance. Visa is the primary network partner for eight of the top
68 10 co-brand partnerships in the US today and we are pleased that Visa has finalized a multi-year extension of
69 our successful credit co-branded partnership with Alaska Airlines, a portfolio that benefits from a loyal customer
70 base and high cross-border usage.
71 We have also had significant co-brand momentum in CEMEA. First, we launched a new co-brand card in partnership
72 with Qatar Airways, British Airways and the National Bank of Kuwait. Second, we expanded our strong global
73 Marriott relationship to launch Qatar's first hospitality co-branded card with Qatar Islamic Bank. Across the
74 United Arab Emirates, we now have exclusive agreements with all the leading airlines marked by a recent
75 agreement with Emirates Skywards.
76 And we also signed an inaugural Airline co-brand agreement in Morocco with Royal Air Maroc. Now newer digital
77 issuers are equally
78 ------------------------
79 output:
80 {"entities": [{"entity": "Visa", "entity_type": "company"}, {"entity": "Alaska Airlines", "entity_type": "company"}, {"entity": "Qatar Airways", "entity_type": "company"}, {"entity": "British Airways", "entity_type": "company"}, {"entity": "National Bank of Kuwait", "entity_type": "company"}, {"entity": "Marriott", "entity_type": "company"}, {"entity": "Qatar Islamic Bank", "entity_type": "company"}, {"entity": "Emirates Skywards", "entity_type": "company"}, {"entity": "Royal Air Maroc", "entity_type": "company"}]}
81 #############################
82 -Real Data-
83 ######################
84 entity_types: [company, organization, person, country, product, service]
85 text: {{ document.content }}
86 ######################
87 output:
88 '''
90 docs = [
91 Document(content="deepset was founded in 2018 in Berlin, and is known for its Haystack framework"),
92 Document(content="Hugging Face is a company that was founded in New York, USA and is known for its Transformers library")
93 ]
95 chat_generator = OpenAIResponsesChatGenerator(
96 generation_kwargs={
97 "text": {
98 "format": {
99 "type": "json_schema",
100 "name": "entity_extraction",
101 "strict": True,
102 "schema": {
103 "type": "object",
104 "properties": {
105 "entities": {
106 "type": "array",
107 "items": {
108 "type": "object",
109 "properties": {
110 "entity": {"type": "string"},
111 "entity_type": {"type": "string"}
112 },
113 "required": ["entity", "entity_type"],
114 "additionalProperties": False
115 }
116 }
117 },
118 "required": ["entities"],
119 "additionalProperties": False
120 }
121 }
122 },
123 },
124 max_retries=1,
125 timeout=60.0,
126 )
128 extractor = LLMMetadataExtractor(
129 prompt=NER_PROMPT,
130 chat_generator=chat_generator,
131 expected_keys=["entities"],
132 raise_on_failure=False,
133 )
135 extractor.run(documents=docs)
136 # >> {'documents': [
137 # Document(id=.., content: 'deepset was founded in 2018 in Berlin, and is known for its Haystack framework',
138 # meta: {'entities': [{'entity': 'deepset', 'entity_type': 'company'}, {'entity': 'Berlin', 'entity_type': 'city'},
139 # {'entity': 'Haystack', 'entity_type': 'product'}]}),
140 # Document(id=.., content: 'Hugging Face is a company that was founded in New York, USA and is known for its Transformers library',
141 # meta: {'entities': [
142 # {'entity': 'Hugging Face', 'entity_type': 'company'}, {'entity': 'New York', 'entity_type': 'city'},
143 # {'entity': 'USA', 'entity_type': 'country'}, {'entity': 'Transformers', 'entity_type': 'product'}
144 # ]})
145 # ]
146 # 'failed_documents': []
147 # }
148 # >>
149 ```
150 """ # noqa: E501
152 def __init__(
153 self,
154 prompt: str,
155 chat_generator: ChatGenerator,
156 expected_keys: list[str] | None = None,
157 page_range: list[str | int] | None = None,
158 raise_on_failure: bool = False,
159 max_workers: int = 3,
160 ) -> None:
161 """
162 Initializes the LLMMetadataExtractor.
164 :param prompt: The prompt to be used for the LLM. It must contain exactly one variable, called `document`,
165 which points to a single document in the list of documents. For example, to access the content of the
166 document, use `{{ document.content }}` in the prompt.
167 :param chat_generator: a ChatGenerator instance which represents the LLM. In order for the component to work,
168 the LLM should be configured to return a JSON object. For example, when using the
169 OpenAIResponsesChatGenerator, you should pass `{"text": {"format": {"type": "json_schema", ...}}}` in the
170 `generation_kwargs`.
171 :param expected_keys: The keys expected in the JSON output from the LLM.
172 :param page_range: A range of pages to extract metadata from. For example, page_range=['1', '3'] will extract
173 metadata from the first and third pages of each document. It also accepts printable range strings, e.g.:
174 ['1-3', '5', '8', '10-12'] will extract metadata from pages 1, 2, 3, 5, 8, 10,11, 12.
175 If None, metadata will be extracted from the entire document for each document in the documents list.
176 This parameter is optional and can be overridden in the `run` method.
177 :param raise_on_failure: Whether to raise an error on failure during the execution of the Generator or
178 validation of the JSON output.
179 :param max_workers: The maximum number of workers to use in the thread pool executor.
180 This parameter is used limit the maximum number of requests that should be allowed to run concurrently
181 when using the `run_async` method.
182 """
183 self.prompt = prompt
184 ast = SandboxedEnvironment().parse(prompt)
185 template_variables = meta.find_undeclared_variables(ast)
186 variables = list(template_variables)
187 if variables != ["document"]:
188 raise ValueError(
189 f"Prompt must have exactly one variable called 'document'. "
190 f"Found {','.join(variables) or 'no variables'} in the prompt."
191 )
192 self.builder = PromptBuilder(prompt, required_variables=variables)
193 self.raise_on_failure = raise_on_failure
194 self.expected_keys = expected_keys or []
195 self.splitter = DocumentSplitter(split_by="page", split_length=1)
196 self.expanded_range = expand_page_range(page_range) if page_range else None
197 self.max_workers = max_workers
198 self._chat_generator = chat_generator
200 def warm_up(self) -> None:
201 """
202 Warm up the underlying chat generator and splitter.
203 """
204 for inner in (self._chat_generator, self.splitter):
205 if hasattr(inner, "warm_up"):
206 inner.warm_up()
208 async def warm_up_async(self) -> None:
209 """
210 Warm up the underlying chat generator and splitter on the serving event loop.
211 """
212 for inner in (self._chat_generator, self.splitter):
213 if hasattr(inner, "warm_up_async"):
214 await inner.warm_up_async()
215 elif hasattr(inner, "warm_up"):
216 inner.warm_up()
218 def close(self) -> None:
219 """
220 Release the underlying chat generator's and splitter's resources.
221 """
222 for inner in (self._chat_generator, self.splitter):
223 if hasattr(inner, "close"):
224 inner.close()
226 async def close_async(self) -> None:
227 """
228 Release the underlying chat generator's and splitter's async resources.
229 """
230 for inner in (self._chat_generator, self.splitter):
231 if hasattr(inner, "close_async"):
232 await inner.close_async()
233 elif hasattr(inner, "close"):
234 inner.close()
236 def to_dict(self) -> dict[str, Any]:
237 """
238 Serializes the component to a dictionary.
240 :returns:
241 Dictionary with serialized data.
242 """
244 return default_to_dict(
245 self,
246 prompt=self.prompt,
247 chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"),
248 expected_keys=self.expected_keys,
249 page_range=self.expanded_range,
250 raise_on_failure=self.raise_on_failure,
251 max_workers=self.max_workers,
252 )
254 @classmethod
255 def from_dict(cls, data: dict[str, Any]) -> "LLMMetadataExtractor":
256 """
257 Deserializes the component from a dictionary.
259 :param data:
260 Dictionary with serialized data.
261 :returns:
262 An instance of the component.
263 """
265 deserialize_chatgenerator_inplace(data["init_parameters"], key="chat_generator")
266 return default_from_dict(cls, data)
268 def _prepare_prompts(
269 self, documents: list[Document], expanded_range: list[int] | None = None
270 ) -> list[ChatMessage | None]:
271 all_prompts: list[ChatMessage | None] = []
272 for document in documents:
273 if not document.content:
274 logger.warning("Document {doc_id} has no content. Skipping metadata extraction.", doc_id=document.id)
275 all_prompts.append(None)
276 continue
278 if expanded_range:
279 doc_copy = copy.deepcopy(document)
280 pages = self.splitter.run(documents=[doc_copy])
281 content = ""
282 for idx, page in enumerate(pages["documents"]):
283 if idx + 1 in expanded_range and page.content is not None:
284 content += page.content
285 doc_copy = replace(doc_copy, content=content)
286 else:
287 doc_copy = document
289 prompt_with_doc = self.builder.run(template=self.prompt, template_variables={"document": doc_copy})
291 # build a ChatMessage with the prompt
292 message = ChatMessage.from_user(prompt_with_doc["prompt"])
293 all_prompts.append(message)
295 return all_prompts
297 @staticmethod
298 def _fail(document: Document, error: str, response: ChatMessage | None = None) -> tuple[Document, bool]:
299 """Return a copy of ``document`` with the failure metadata set, flagged as failed."""
300 new_meta = {**document.meta, "metadata_extraction_error": error, "metadata_extraction_response": response}
301 return replace(document, meta=new_meta), False
303 def _run_on_thread(
304 self, document: Document, prompt: ChatMessage | None, parent_span: tracing.Span | None = None
305 ) -> tuple[Document, bool]:
306 if prompt is None:
307 return self._fail(document, "Document has no content, skipping LLM call.")
309 try:
310 with _trace_chat_generator_run(
311 self._chat_generator, {"messages": [prompt]}, parent_span=parent_span
312 ) as span:
313 result = self._chat_generator.run(messages=[prompt])
314 span.set_content_tag("haystack.component.output", result)
315 except Exception as e:
316 if self.raise_on_failure:
317 raise e
318 logger.exception(
319 "LLM {class_name} execution failed. Skipping metadata extraction. Failed with exception '{error}'.",
320 class_name=self._chat_generator.__class__.__name__,
321 error=e,
322 )
323 return self._fail(document, "LLM failed with exception: " + str(e))
324 return self._process_reply(document, result["replies"][0])
326 async def _run_async(
327 self, document: Document, prompt: ChatMessage | None, parent_span: tracing.Span | None = None
328 ) -> tuple[Document, bool]:
329 if prompt is None:
330 return self._fail(document, "Document has no content, skipping LLM call.")
332 try:
333 with _trace_chat_generator_run(
334 self._chat_generator, {"messages": [prompt]}, parent_span=parent_span
335 ) as span:
336 result = await _execute_component_async(self._chat_generator, messages=[prompt])
337 span.set_content_tag("haystack.component.output", result)
338 except Exception as e:
339 if self.raise_on_failure:
340 raise e
341 logger.exception(
342 "LLM {class_name} execution failed. Skipping metadata extraction. Failed with exception '{error}'.",
343 class_name=self._chat_generator.__class__.__name__,
344 error=e,
345 )
346 return self._fail(document, "LLM failed with exception: " + str(e))
347 return self._process_reply(document, result["replies"][0])
349 def _process_reply(self, document: Document, reply: ChatMessage) -> tuple[Document, bool]:
350 """
351 Parse one document's LLM reply into metadata.
353 Returns (updated_document, True if success else False).
354 """
355 try:
356 parsed_metadata = _parse_dict_from_json(
357 reply.text or "", expected_keys=self.expected_keys, raise_on_failure=True
358 )
359 except (ValueError, json.JSONDecodeError) as e:
360 logger.warning(
361 "Response from the LLM is not valid JSON or missing expected keys. Received output: {response}",
362 response=reply.text,
363 )
364 if self.raise_on_failure:
365 raise
366 return self._fail(document, "Response is not valid JSON or missing keys. Error: " + str(e), response=reply)
368 new_meta = {**document.meta, **parsed_metadata}
369 # Remove metadata_extraction_error and metadata_extraction_response if present from previous runs.
370 new_meta.pop("metadata_extraction_error", None)
371 new_meta.pop("metadata_extraction_response", None)
372 return replace(document, meta=new_meta), True
374 @component.output_types(documents=list[Document], failed_documents=list[Document])
375 def run(self, documents: list[Document], page_range: list[str | int] | None = None) -> dict[str, Any]:
376 """
377 Extract metadata from documents using a Large Language Model.
379 If `page_range` is provided, the metadata will be extracted from the specified range of pages. This component
380 will split the documents into pages and extract metadata from the specified range of pages. The metadata will be
381 extracted from the entire document if `page_range` is not provided.
383 The original documents will be returned updated with the extracted metadata.
385 :param documents: List of documents to extract metadata from.
386 :param page_range: A range of pages to extract metadata from. For example, page_range=['1', '3'] will extract
387 metadata from the first and third pages of each document. It also accepts printable range
388 strings, e.g.: ['1-3', '5', '8', '10-12'] will extract metadata from pages 1, 2, 3, 5, 8, 10,
389 11, 12.
390 If None, metadata will be extracted from the entire document for each document in the
391 documents list.
392 :returns:
393 A dictionary with the keys:
394 - "documents": A list of documents that were successfully updated with the extracted metadata.
395 - "failed_documents": A list of documents that failed to extract metadata. These documents will have
396 "metadata_extraction_error" and "metadata_extraction_response" in their metadata. These documents can be
397 re-run with the extractor to extract metadata.
398 """
399 if len(documents) == 0:
400 logger.warning("No documents provided. Skipping metadata extraction.")
401 return {"documents": [], "failed_documents": []}
403 self.warm_up()
405 expanded_range = self.expanded_range
406 if page_range:
407 expanded_range = expand_page_range(page_range)
409 # Create ChatMessage prompts for each document
410 all_prompts = self._prepare_prompts(documents=documents, expanded_range=expanded_range)
412 # Capture the current span here so worker threads nest their generator spans under the component span.
413 parent_span = tracing.tracer.current_span()
415 # Run the LLM on each prompt
416 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
417 results = executor.map(
418 lambda document, prompt: self._run_on_thread(document, prompt, parent_span=parent_span),
419 documents,
420 all_prompts,
421 )
423 successful_documents = []
424 failed_documents = []
425 for doc, success in results:
426 if success:
427 successful_documents.append(doc)
428 else:
429 failed_documents.append(doc)
431 return {"documents": successful_documents, "failed_documents": failed_documents}
433 @component.output_types(documents=list[Document], failed_documents=list[Document])
434 async def run_async(self, documents: list[Document], page_range: list[str | int] | None = None) -> dict[str, Any]:
435 """
436 Asynchronously extract metadata from documents using a Large Language Model.
438 If `page_range` is provided, the metadata will be extracted from the specified range of pages. This component
439 will split the documents into pages and extract metadata from the specified range of pages. The metadata will be
440 extracted from the entire document if `page_range` is not provided.
442 The original documents will be returned updated with the extracted metadata.
444 This is the asynchronous version of the `run` method. It has the same parameters
445 and return values but can be used with `await` in an async code.
447 :param documents: List of documents to extract metadata from.
448 :param page_range: A range of pages to extract metadata from. For example, page_range=['1', '3'] will extract
449 metadata from the first and third pages of each document. It also accepts printable range
450 strings, e.g.: ['1-3', '5', '8', '10-12'] will extract metadata from pages 1, 2, 3, 5, 8, 10,
451 11, 12.
452 If None, metadata will be extracted from the entire document for each document in the
453 documents list.
454 :returns:
455 A dictionary with the keys:
456 - "documents": A list of documents that were successfully updated with the extracted metadata.
457 - "failed_documents": A list of documents that failed to extract metadata. These documents will have
458 "metadata_extraction_error" and "metadata_extraction_response" in their metadata. These documents can be
459 re-run with the extractor to extract metadata.
460 """
461 if len(documents) == 0:
462 logger.warning("No documents provided. Skipping metadata extraction.")
463 return {"documents": [], "failed_documents": []}
465 await self.warm_up_async()
467 expanded_range = self.expanded_range
468 if page_range:
469 expanded_range = expand_page_range(page_range)
471 # Create ChatMessage prompts for each document
472 all_prompts = self._prepare_prompts(documents=documents, expanded_range=expanded_range)
474 # Capture the current span here so concurrent tasks nest their generator spans under the component span.
475 parent_span = tracing.tracer.current_span()
477 # Run the LLM on each prompt, bounding concurrency per task so max_workers is enforced.
478 sem = Semaphore(max(1, self.max_workers))
480 async def _bounded_run(document: Document, prompt: ChatMessage | None) -> tuple[Document, bool]:
481 async with sem:
482 return await self._run_async(document, prompt, parent_span=parent_span)
484 results = await gather(
485 *[_bounded_run(document, prompt) for document, prompt in zip(documents, all_prompts, strict=True)]
486 )
488 successful_documents = []
489 failed_documents = []
490 for doc, success in results:
491 if success:
492 successful_documents.append(doc)
493 else:
494 failed_documents.append(doc)
496 return {"documents": successful_documents, "failed_documents": failed_documents}