chore: import upstream snapshot with attribution

haystack/components/builders/answer_builder.py

@@ -0,0 +1,314 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from dataclasses import replace
+from typing import Any
+
+from haystack import Document, GeneratedAnswer, component, logging
+from haystack.dataclasses.chat_message import ChatMessage
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_REFERENCE_PATTERN = r"\[(\d+)\]"
+EXPANDED_REFERENCE_PATTERN = r"\[(\d+(?:[,-]\d+)*)\]"
+
+
+@component
+class AnswerBuilder:
+    """
+    Converts a query and Generator replies into a `GeneratedAnswer` object.
+
+    AnswerBuilder parses Generator replies using custom regular expressions.
+    Check out the usage example below to see how it works.
+    Optionally, it can also take documents and metadata from the Generator to add to the `GeneratedAnswer` object.
+    AnswerBuilder works with both non-chat and chat Generators.
+
+    ### Usage example
+
+    ```python
+    from haystack.components.builders import AnswerBuilder
+
+    builder = AnswerBuilder(pattern="Answer: (.*)")
+    builder.run(query="What's the answer?", replies=["This is an argument. Answer: This is the answer."])
+    ```
+
+    ### Usage example with documents and reference pattern
+
+    ```python
+    from haystack import Document
+    from haystack.components.builders import AnswerBuilder
+
+    replies = ["The capital of France is Paris [2]."]
+
+    docs = [
+        Document(content="Berlin is the capital of Germany."),
+        Document(content="Paris is the capital of France."),
+        Document(content="Rome is the capital of Italy."),
+    ]
+
+    builder = AnswerBuilder(reference_pattern="\\[(\\d+)\\]", return_only_referenced_documents=False)
+    result = builder.run(query="What is the capital of France?", replies=replies, documents=docs)["answers"][0]
+
+    print(f"Answer: {result.data}")
+    print("References:")
+    for doc in result.documents:
+        if doc.meta["referenced"]:
+            print(f"[{doc.meta['source_index']}] {doc.content}")
+    print("Other sources:")
+    for doc in result.documents:
+        if not doc.meta["referenced"]:
+            print(f"[{doc.meta['source_index']}] {doc.content}")
+
+    # >> Answer: The capital of France is Paris
+    # >> References:
+    # >> [2] Paris is the capital of France.
+    # >> Other sources:
+    # >> [1] Berlin is the capital of Germany.
+    # >> [3] Rome is the capital of Italy.
+    ```
+    """
+
+    def __init__(
+        self,
+        pattern: str | None = None,
+        reference_pattern: str | None = None,
+        last_message_only: bool = False,
+        *,
+        return_only_referenced_documents: bool = True,
+        expand_reference_ranges: bool = False,
+    ) -> None:
+        """
+        Creates an instance of the AnswerBuilder component.
+
+        :param pattern:
+            The regular expression pattern to extract the answer text from the Generator.
+            If not specified, the entire response is used as the answer.
+            The regular expression can have one capture group at most.
+            If present, the capture group text
+            is used as the answer. If no capture group is present, the whole match is used as the answer.
+            Examples:
+                `[^\\n]+$` finds "this is an answer" in a string "this is an argument.\\nthis is an answer".
+                `Answer: (.*)` finds "this is an answer" in a string "this is an argument. Answer: this is an answer".
+
+        :param reference_pattern:
+            The regular expression pattern used for parsing the document references.
+            If not specified, no parsing is done, and all documents are returned.
+            References need to be specified as indices of the input documents and start at [1].
+            Example: `\\[(\\d+)\\]` finds "1" in a string "this is an answer[1]".
+            If this parameter is provided, documents metadata will contain a "referenced" key with a boolean value.
+
+        :param last_message_only:
+           If False (default value), all messages are used as the answer.
+           If True, only the last message is used as the answer.
+
+        :param return_only_referenced_documents:
+            To be used in conjunction with `reference_pattern`.
+            If True (default value), only the documents that were actually referenced in `replies` are returned.
+            If False, all documents are returned.
+            If `reference_pattern` is not provided, this parameter has no effect, and all documents are returned.
+        :param expand_reference_ranges:
+            If True, reference ranges like `[6-10]` are expanded to documents 6 through 10.
+            Defaults to False for backwards compatibility.
+            When enabled with the default `reference_pattern`, a broader pattern is used automatically.
+        """
+        if pattern:
+            AnswerBuilder._check_num_groups_in_regex(pattern)
+
+        self.pattern = pattern
+        self.reference_pattern = reference_pattern
+        self.last_message_only = last_message_only
+        self.return_only_referenced_documents = return_only_referenced_documents
+        self.expand_reference_ranges = expand_reference_ranges
+
+    @component.output_types(answers=list[GeneratedAnswer])
+    def run(
+        self,
+        query: str,
+        replies: list[str] | list[ChatMessage],
+        meta: list[dict[str, Any]] | None = None,
+        documents: list[Document] | None = None,
+        pattern: str | None = None,
+        reference_pattern: str | None = None,
+        expand_reference_ranges: bool | None = None,
+    ) -> dict[str, Any]:
+        """
+        Turns the output of a Generator into `GeneratedAnswer` objects using regular expressions.
+
+        :param query:
+            The input query used as the Generator prompt.
+        :param replies:
+            The output of the Generator. Can be a list of strings or a list of `ChatMessage` objects.
+        :param meta:
+            The metadata returned by the Generator. If not specified, the generated answer will contain no metadata.
+        :param documents:
+            The documents used as the Generator inputs. If specified, they are added to
+            the `GeneratedAnswer` objects.
+            The Document copies inside the returned `GeneratedAnswer.documents` each include a "source_index" key,
+            representing the document's 1-based position in the input list. The original input documents are
+            not modified.
+            When `reference_pattern` is provided:
+            - "referenced" key is added to the Document copies inside `GeneratedAnswer.documents`, indicating if
+            the document was referenced in the output.
+            - `return_only_referenced_documents` init parameter controls if all or only referenced documents are
+            returned.
+        :param pattern:
+            The regular expression pattern to extract the answer text from the Generator.
+            If not specified, the entire response is used as the answer.
+            The regular expression can have one capture group at most.
+            If present, the capture group text
+            is used as the answer. If no capture group is present, the whole match is used as the answer.
+                Examples:
+                    `[^\\n]+$` finds "this is an answer" in a string "this is an argument.\\nthis is an answer".
+                    `Answer: (.*)` finds "this is an answer" in a string
+                    "this is an argument. Answer: this is an answer".
+        :param reference_pattern:
+            The regular expression pattern used for parsing the document references.
+            If not specified, no parsing is done, and all documents are returned.
+            References need to be specified as indices of the input documents and start at [1].
+            Example: `\\[(\\d+)\\]` finds "1" in a string "this is an answer[1]".
+        :param expand_reference_ranges:
+            If True, reference ranges like `[6-10]` are expanded to documents 6 through 10.
+            If not specified, the value from the component initialization is used.
+
+        :returns: A dictionary with the following keys:
+            - `answers`: The answers received from the output of the Generator.
+        """
+        if not meta:
+            meta = [{}] * len(replies)
+        elif len(replies) != len(meta):
+            raise ValueError(f"Number of replies ({len(replies)}), and metadata ({len(meta)}) must match.")
+
+        if pattern:
+            AnswerBuilder._check_num_groups_in_regex(pattern)
+
+        pattern = pattern or self.pattern
+        reference_pattern = reference_pattern or self.reference_pattern
+        expand_reference_ranges = (
+            self.expand_reference_ranges if expand_reference_ranges is None else expand_reference_ranges
+        )
+        reference_pattern = AnswerBuilder._resolve_reference_pattern(
+            reference_pattern=reference_pattern, expand_reference_ranges=expand_reference_ranges
+        )
+
+        replies_to_iterate = replies[-1:] if self.last_message_only and replies else replies
+        meta_to_iterate = meta[-1:] if self.last_message_only and meta else meta
+
+        all_answers = []
+        for reply, given_metadata in zip(replies_to_iterate, meta_to_iterate, strict=True):
+            # Extract content from ChatMessage objects if reply is a ChatMessages, else use the string as is
+            extracted_reply = reply.text or "" if isinstance(reply, ChatMessage) else str(reply)
+            extracted_metadata = reply.meta if isinstance(reply, ChatMessage) else {}
+
+            extracted_metadata = {**extracted_metadata, **given_metadata}
+            extracted_metadata["all_messages"] = replies
+
+            referenced_docs = []
+            if documents:
+                referenced_idxs = (
+                    AnswerBuilder._extract_reference_idxs(
+                        extracted_reply,
+                        reference_pattern,
+                        expand_ranges=expand_reference_ranges,
+                        num_documents=len(documents),
+                    )
+                    if reference_pattern
+                    else set()
+                )
+                doc_idxs = (
+                    referenced_idxs
+                    if reference_pattern and self.return_only_referenced_documents
+                    else set(range(len(documents)))
+                )
+
+                for idx in doc_idxs:
+                    try:
+                        doc = documents[idx]
+                    except IndexError:
+                        logger.warning(
+                            "Document index '{index}' referenced in Generator output is out of range. ", index=idx + 1
+                        )
+                        continue
+
+                    doc_meta: dict[str, Any] = dict(doc.meta or {})
+                    doc_meta["source_index"] = idx + 1
+                    if reference_pattern:
+                        doc_meta["referenced"] = idx in referenced_idxs
+                    referenced_docs.append(replace(doc, meta=doc_meta))
+
+            answer_string = AnswerBuilder._extract_answer_string(extracted_reply, pattern)
+            answer = GeneratedAnswer(
+                data=answer_string, query=query, documents=referenced_docs, meta=extracted_metadata
+            )
+            all_answers.append(answer)
+
+        return {"answers": all_answers}
+
+    @staticmethod
+    def _extract_answer_string(reply: str, pattern: str | None = None) -> str:
+        """
+        Extract the answer string from the generator output using the specified pattern.
+
+        If no pattern is specified, the whole string is used as the answer.
+
+        :param reply:
+            The output of the Generator. A string.
+        :param pattern:
+            The regular expression pattern to use to extract the answer text from the generator output.
+        """
+        if pattern is None:
+            return reply
+
+        if match := re.search(pattern, reply):
+            # No capture group in pattern -> use the whole match as answer
+            if not match.lastindex:
+                return match.group(0)
+            # One capture group in pattern -> use the capture group as answer
+            return match.group(1)
+        return ""
+
+    @staticmethod
+    def _resolve_reference_pattern(reference_pattern: str | None, expand_reference_ranges: bool) -> str | None:
+        if not reference_pattern or not expand_reference_ranges:
+            return reference_pattern
+        if reference_pattern == DEFAULT_REFERENCE_PATTERN:
+            return EXPANDED_REFERENCE_PATTERN
+        return reference_pattern
+
+    @staticmethod
+    def _extract_reference_idxs(
+        reply: str, reference_pattern: str, expand_ranges: bool = False, num_documents: int | None = None
+    ) -> set[int]:
+        matches = re.findall(reference_pattern, reply)
+        idxs: set[int] = set()
+        for match in matches:
+            if expand_ranges:
+                for part in match.split(","):
+                    part = part.strip()
+                    if not part:
+                        continue
+                    if "-" in part:
+                        start_str, end_str = part.split("-", 1)
+                        start, end = int(start_str), int(end_str)
+                        if start > end:
+                            continue
+                        # Clamp the range end to the number of documents to avoid materializing a huge
+                        # set from an out-of-range citation like `[1-999999999]` in the Generator output.
+                        if num_documents is not None:
+                            end = min(end, num_documents)
+                        idxs.update(range(start - 1, end))
+                    else:
+                        idxs.add(int(part) - 1)
+            else:
+                idxs.add(int(match) - 1)
+        return idxs
+
+    @staticmethod
+    def _check_num_groups_in_regex(pattern: str) -> None:
+        num_groups = re.compile(pattern).groups
+        if num_groups > 1:
+            raise ValueError(
+                f"Pattern '{pattern}' contains multiple capture groups. "
+                f"Please specify a pattern with at most one capture group."
+            )