chore: import upstream snapshot with attribution

haystack/components/extractors/llm_metadata_extractor.py

@@ -0,0 +1,474 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+import json
+from asyncio import Semaphore, gather
+from collections.abc import Iterable
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import replace
+from typing import Any
+
+from jinja2 import meta
+from jinja2.sandbox import SandboxedEnvironment
+
+from haystack import Document, component, default_from_dict, default_to_dict, logging
+from haystack.components.builders import PromptBuilder
+from haystack.components.generators.chat.types import ChatGenerator
+from haystack.components.preprocessors import DocumentSplitter
+from haystack.core.serialization import component_to_dict
+from haystack.dataclasses import ChatMessage
+from haystack.utils import deserialize_chatgenerator_inplace, expand_page_range
+from haystack.utils.async_utils import _execute_component_async
+from haystack.utils.misc import _parse_dict_from_json
+
+logger = logging.getLogger(__name__)
+
+
+@component
+class LLMMetadataExtractor:
+    """
+    Extracts metadata from documents using a Large Language Model (LLM).
+
+    The metadata is extracted by providing a prompt to an LLM that generates the metadata.
+
+    This component expects as input a list of documents and a prompt. The prompt must have exactly one variable, called
+    `document`, that points to a single document in the list of documents. So to access the content of the document,
+    you can use `{{ document.content }}` in the prompt.
+
+    The component will run the LLM on each document in the list and extract metadata from the document. The metadata
+    will be added to the document's metadata field. If the LLM fails to extract metadata from a document, the document
+    will be added to the `failed_documents` list. The failed documents will have the keys `metadata_extraction_error` and
+    `metadata_extraction_response` in their metadata. These documents can be re-run with another extractor to
+    extract metadata by using the `metadata_extraction_response` and `metadata_extraction_error` in the prompt.
+
+    ```python
+    from haystack import Document
+    from haystack.components.extractors.llm_metadata_extractor import LLMMetadataExtractor
+    from haystack.components.generators.chat import OpenAIChatGenerator
+
+    NER_PROMPT = '''
+    -Goal-
+    Given text and a list of entity types, identify all entities of those types from the text.
+
+    -Steps-
+    1. Identify all entities. For each identified entity, extract the following information:
+    - entity: Name of the entity
+    - entity_type: One of the following types: [organization, product, service, industry]
+    Format each entity as a JSON like: {"entity": <entity_name>, "entity_type": <entity_type>}
+
+    2. Return output in a single list with all the entities identified in steps 1.
+
+    -Examples-
+    ######################
+    Example 1:
+    entity_types: [organization, person, partnership, financial metric, product, service, industry, investment strategy, market trend]
+    text: Another area of strength is our co-brand issuance. Visa is the primary network partner for eight of the top
+    10 co-brand partnerships in the US today and we are pleased that Visa has finalized a multi-year extension of
+    our successful credit co-branded partnership with Alaska Airlines, a portfolio that benefits from a loyal customer
+    base and high cross-border usage.
+    We have also had significant co-brand momentum in CEMEA. First, we launched a new co-brand card in partnership
+    with Qatar Airways, British Airways and the National Bank of Kuwait. Second, we expanded our strong global
+    Marriott relationship to launch Qatar's first hospitality co-branded card with Qatar Islamic Bank. Across the
+    United Arab Emirates, we now have exclusive agreements with all the leading airlines marked by a recent
+    agreement with Emirates Skywards.
+    And we also signed an inaugural Airline co-brand agreement in Morocco with Royal Air Maroc. Now newer digital
+    issuers are equally
+    ------------------------
+    output:
+    {"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"}]}
+    #############################
+    -Real Data-
+    ######################
+    entity_types: [company, organization, person, country, product, service]
+    text: {{ document.content }}
+    ######################
+    output:
+    '''
+
+    docs = [
+        Document(content="deepset was founded in 2018 in Berlin, and is known for its Haystack framework"),
+        Document(content="Hugging Face is a company that was founded in New York, USA and is known for its Transformers library")
+    ]
+
+    chat_generator = OpenAIChatGenerator(
+        generation_kwargs={
+            "max_completion_tokens": 500,
+            "temperature": 0.0,
+            "seed": 0,
+            "response_format": {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "entity_extraction",
+                    "schema": {
+                        "type": "object",
+                        "properties": {
+                            "entities": {
+                                "type": "array",
+                                "items": {
+                                    "type": "object",
+                                    "properties": {
+                                        "entity": {"type": "string"},
+                                        "entity_type": {"type": "string"}
+                                    },
+                                    "required": ["entity", "entity_type"],
+                                    "additionalProperties": False
+                                }
+                            }
+                        },
+                        "required": ["entities"],
+                        "additionalProperties": False
+                    }
+                }
+            },
+        },
+        max_retries=1,
+        timeout=60.0,
+    )
+
+    extractor = LLMMetadataExtractor(
+        prompt=NER_PROMPT,
+        chat_generator=chat_generator,
+        expected_keys=["entities"],
+        raise_on_failure=False,
+    )
+
+    extractor.run(documents=docs)
+    # >> {'documents': [
+    #     Document(id=.., content: 'deepset was founded in 2018 in Berlin, and is known for its Haystack framework',
+    #     meta: {'entities': [{'entity': 'deepset', 'entity_type': 'company'}, {'entity': 'Berlin', 'entity_type': 'city'},
+    #           {'entity': 'Haystack', 'entity_type': 'product'}]}),
+    #     Document(id=.., content: 'Hugging Face is a company that was founded in New York, USA and is known for its Transformers library',
+    #     meta: {'entities': [
+    #             {'entity': 'Hugging Face', 'entity_type': 'company'}, {'entity': 'New York', 'entity_type': 'city'},
+    #             {'entity': 'USA', 'entity_type': 'country'}, {'entity': 'Transformers', 'entity_type': 'product'}
+    #             ]})
+    #        ]
+    #     'failed_documents': []
+    #    }
+    # >>
+    ```
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        prompt: str,
+        chat_generator: ChatGenerator,
+        expected_keys: list[str] | None = None,
+        page_range: list[str | int] | None = None,
+        raise_on_failure: bool = False,
+        max_workers: int = 3,
+    ) -> None:
+        """
+        Initializes the LLMMetadataExtractor.
+
+        :param prompt: The prompt to be used for the LLM. It must contain exactly one variable, called `document`,
+            which points to a single document in the list of documents. For example, to access the content of the
+            document, use `{{ document.content }}` in the prompt.
+        :param chat_generator: a ChatGenerator instance which represents the LLM. In order for the component to work,
+            the LLM should be configured to return a JSON object. For example, when using the OpenAIChatGenerator, you
+            should pass `{"response_format": {"type": "json_object"}}` in the `generation_kwargs`.
+        :param expected_keys: The keys expected in the JSON output from the LLM.
+        :param page_range: A range of pages to extract metadata from. For example, page_range=['1', '3'] will extract
+            metadata from the first and third pages of each document. It also accepts printable range strings, e.g.:
+            ['1-3', '5', '8', '10-12'] will extract metadata from pages 1, 2, 3, 5, 8, 10,11, 12.
+            If None, metadata will be extracted from the entire document for each document in the documents list.
+            This parameter is optional and can be overridden in the `run` method.
+        :param raise_on_failure: Whether to raise an error on failure during the execution of the Generator or
+            validation of the JSON output.
+        :param max_workers: The maximum number of workers to use in the thread pool executor.
+            This parameter is used limit the maximum number of requests that should be allowed to run concurrently
+            when using the `run_async` method.
+        """
+        self.prompt = prompt
+        ast = SandboxedEnvironment().parse(prompt)
+        template_variables = meta.find_undeclared_variables(ast)
+        variables = list(template_variables)
+        if variables != ["document"]:
+            raise ValueError(
+                f"Prompt must have exactly one variable called 'document'. "
+                f"Found {','.join(variables) or 'no variables'} in the prompt."
+            )
+        self.builder = PromptBuilder(prompt, required_variables=variables)
+        self.raise_on_failure = raise_on_failure
+        self.expected_keys = expected_keys or []
+        self.splitter = DocumentSplitter(split_by="page", split_length=1)
+        self.expanded_range = expand_page_range(page_range) if page_range else None
+        self.max_workers = max_workers
+        self._chat_generator = chat_generator
+
+    def warm_up(self) -> None:
+        """
+        Warm up the underlying chat generator and splitter.
+        """
+        for inner in (self._chat_generator, self.splitter):
+            if hasattr(inner, "warm_up"):
+                inner.warm_up()
+
+    async def warm_up_async(self) -> None:
+        """
+        Warm up the underlying chat generator and splitter on the serving event loop.
+        """
+        for inner in (self._chat_generator, self.splitter):
+            if hasattr(inner, "warm_up_async"):
+                await inner.warm_up_async()
+            elif hasattr(inner, "warm_up"):
+                inner.warm_up()
+
+    def close(self) -> None:
+        """
+        Release the underlying chat generator's and splitter's resources.
+        """
+        for inner in (self._chat_generator, self.splitter):
+            if hasattr(inner, "close"):
+                inner.close()
+
+    async def close_async(self) -> None:
+        """
+        Release the underlying chat generator's and splitter's async resources.
+        """
+        for inner in (self._chat_generator, self.splitter):
+            if hasattr(inner, "close_async"):
+                await inner.close_async()
+            elif hasattr(inner, "close"):
+                inner.close()
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serializes the component to a dictionary.
+
+        :returns:
+            Dictionary with serialized data.
+        """
+
+        return default_to_dict(
+            self,
+            prompt=self.prompt,
+            chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"),
+            expected_keys=self.expected_keys,
+            page_range=self.expanded_range,
+            raise_on_failure=self.raise_on_failure,
+            max_workers=self.max_workers,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "LLMMetadataExtractor":
+        """
+        Deserializes the component from a dictionary.
+
+        :param data:
+            Dictionary with serialized data.
+        :returns:
+            An instance of the component.
+        """
+
+        deserialize_chatgenerator_inplace(data["init_parameters"], key="chat_generator")
+        return default_from_dict(cls, data)
+
+    def _extract_metadata(self, llm_answer: str) -> dict[str, Any]:
+        try:
+            parsed_metadata = _parse_dict_from_json(llm_answer, expected_keys=self.expected_keys, raise_on_failure=True)
+        except (ValueError, json.JSONDecodeError) as e:
+            logger.warning(
+                "Response from the LLM is not valid JSON or missing expected keys. Received output: {response}",
+                response=llm_answer,
+            )
+            if self.raise_on_failure:
+                raise e
+            return {"error": "Response is not valid JSON or missing keys. Error: " + str(e)}
+
+        return parsed_metadata
+
+    def _prepare_prompts(
+        self, documents: list[Document], expanded_range: list[int] | None = None
+    ) -> list[ChatMessage | None]:
+        all_prompts: list[ChatMessage | None] = []
+        for document in documents:
+            if not document.content:
+                logger.warning("Document {doc_id} has no content. Skipping metadata extraction.", doc_id=document.id)
+                all_prompts.append(None)
+                continue
+
+            if expanded_range:
+                doc_copy = copy.deepcopy(document)
+                pages = self.splitter.run(documents=[doc_copy])
+                content = ""
+                for idx, page in enumerate(pages["documents"]):
+                    if idx + 1 in expanded_range and page.content is not None:
+                        content += page.content
+                doc_copy = replace(doc_copy, content=content)
+            else:
+                doc_copy = document
+
+            prompt_with_doc = self.builder.run(template=self.prompt, template_variables={"document": doc_copy})
+
+            # build a ChatMessage with the prompt
+            message = ChatMessage.from_user(prompt_with_doc["prompt"])
+            all_prompts.append(message)
+
+        return all_prompts
+
+    def _run_on_thread(self, prompt: ChatMessage | None) -> dict[str, Any]:
+        # If prompt is None, return an error dictionary
+        if prompt is None:
+            return {"error": "Document has no content, skipping LLM call."}
+
+        try:
+            result = self._chat_generator.run(messages=[prompt])
+        except Exception as e:
+            if self.raise_on_failure:
+                raise e
+            logger.exception(
+                "LLM {class_name} execution failed. Skipping metadata extraction. Failed with exception '{error}'.",
+                class_name=self._chat_generator.__class__.__name__,
+                error=e,
+            )
+            result = {"error": "LLM failed with exception: " + str(e)}
+        return result
+
+    async def _run_async(self, prompt: ChatMessage | None) -> dict[str, Any]:
+        # If prompt is None, return an error dictionary
+        if prompt is None:
+            return {"error": "Document has no content, skipping LLM call."}
+
+        try:
+            result = await _execute_component_async(self._chat_generator, messages=[prompt])
+        except Exception as e:
+            if self.raise_on_failure:
+                raise e
+            logger.exception(
+                "LLM {class_name} execution failed. Skipping metadata extraction. Failed with exception '{error}'.",
+                class_name=self._chat_generator.__class__.__name__,
+                error=e,
+            )
+            result = {"error": "LLM failed with exception: " + str(e)}
+        return result
+
+    def _process_results(
+        self, documents: list[Document], results: Iterable[dict[str, Any]]
+    ) -> tuple[list[Document], list[Document]]:
+        successful_documents = []
+        failed_documents = []
+        for document, result in zip(documents, results, strict=True):
+            new_meta = {**document.meta}
+            if "error" in result:
+                new_meta["metadata_extraction_error"] = result["error"]
+                new_meta["metadata_extraction_response"] = None
+                failed_documents.append(replace(document, meta=new_meta))
+                continue
+
+            parsed_metadata = self._extract_metadata(result["replies"][0].text)
+            if "error" in parsed_metadata:
+                new_meta["metadata_extraction_error"] = parsed_metadata["error"]
+                new_meta["metadata_extraction_response"] = result["replies"][0]
+                failed_documents.append(replace(document, meta=new_meta))
+                continue
+
+            for key in parsed_metadata:
+                new_meta[key] = parsed_metadata[key]
+                # Remove metadata_extraction_error and metadata_extraction_response if present from previous runs
+                new_meta.pop("metadata_extraction_error", None)
+                new_meta.pop("metadata_extraction_response", None)
+            successful_documents.append(replace(document, meta=new_meta))
+        return successful_documents, failed_documents
+
+    @component.output_types(documents=list[Document], failed_documents=list[Document])
+    def run(self, documents: list[Document], page_range: list[str | int] | None = None) -> dict[str, Any]:
+        """
+        Extract metadata from documents using a Large Language Model.
+
+        If `page_range` is provided, the metadata will be extracted from the specified range of pages. This component
+        will split the documents into pages and extract metadata from the specified range of pages. The metadata will be
+        extracted from the entire document if `page_range` is not provided.
+
+        The original documents will be returned  updated with the extracted metadata.
+
+        :param documents: List of documents to extract metadata from.
+        :param page_range: A range of pages to extract metadata from. For example, page_range=['1', '3'] will extract
+                           metadata from the first and third pages of each document. It also accepts printable range
+                           strings, e.g.: ['1-3', '5', '8', '10-12'] will extract metadata from pages 1, 2, 3, 5, 8, 10,
+                           11, 12.
+                           If None, metadata will be extracted from the entire document for each document in the
+                           documents list.
+        :returns:
+            A dictionary with the keys:
+            - "documents": A list of documents that were successfully updated with the extracted metadata.
+            - "failed_documents": A list of documents that failed to extract metadata. These documents will have
+            "metadata_extraction_error" and "metadata_extraction_response" in their metadata. These documents can be
+            re-run with the extractor to extract metadata.
+        """
+        if len(documents) == 0:
+            logger.warning("No documents provided. Skipping metadata extraction.")
+            return {"documents": [], "failed_documents": []}
+
+        self.warm_up()
+
+        expanded_range = self.expanded_range
+        if page_range:
+            expanded_range = expand_page_range(page_range)
+
+        # Create ChatMessage prompts for each document
+        all_prompts = self._prepare_prompts(documents=documents, expanded_range=expanded_range)
+
+        # Run the LLM on each prompt
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            results = executor.map(self._run_on_thread, all_prompts)
+
+        successful_documents, failed_documents = self._process_results(documents, results)
+
+        return {"documents": successful_documents, "failed_documents": failed_documents}
+
+    @component.output_types(documents=list[Document], failed_documents=list[Document])
+    async def run_async(self, documents: list[Document], page_range: list[str | int] | None = None) -> dict[str, Any]:
+        """
+        Asynchronously extract metadata from documents using a Large Language Model.
+
+        If `page_range` is provided, the metadata will be extracted from the specified range of pages. This component
+        will split the documents into pages and extract metadata from the specified range of pages. The metadata will be
+        extracted from the entire document if `page_range` is not provided.
+
+        The original documents will be returned  updated with the extracted metadata.
+
+        This is the asynchronous version of the `run` method. It has the same parameters
+        and return values but can be used with `await` in an async code.
+
+        :param documents: List of documents to extract metadata from.
+        :param page_range: A range of pages to extract metadata from. For example, page_range=['1', '3'] will extract
+                           metadata from the first and third pages of each document. It also accepts printable range
+                           strings, e.g.: ['1-3', '5', '8', '10-12'] will extract metadata from pages 1, 2, 3, 5, 8, 10,
+                           11, 12.
+                           If None, metadata will be extracted from the entire document for each document in the
+                           documents list.
+        :returns:
+            A dictionary with the keys:
+            - "documents": A list of documents that were successfully updated with the extracted metadata.
+            - "failed_documents": A list of documents that failed to extract metadata. These documents will have
+            "metadata_extraction_error" and "metadata_extraction_response" in their metadata. These documents can be
+            re-run with the extractor to extract metadata.
+        """
+        if len(documents) == 0:
+            logger.warning("No documents provided. Skipping metadata extraction.")
+            return {"documents": [], "failed_documents": []}
+
+        await self.warm_up_async()
+
+        expanded_range = self.expanded_range
+        if page_range:
+            expanded_range = expand_page_range(page_range)
+
+        # Create ChatMessage prompts for each document
+        all_prompts = self._prepare_prompts(documents=documents, expanded_range=expanded_range)
+
+        # Run the LLM on each prompt, bounding concurrency per task so max_workers is enforced.
+        sem = Semaphore(max(1, self.max_workers))
+
+        async def _bounded_run(prompt: ChatMessage | None) -> dict[str, Any]:
+            async with sem:
+                return await self._run_async(prompt)
+
+        results = await gather(*[_bounded_run(prompt) for prompt in all_prompts])
+
+        successful_documents, failed_documents = self._process_results(documents, results)
+
+        return {"documents": successful_documents, "failed_documents": failed_documents}