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