deepset-ai--haystack/haystack/components/rankers/meta_field_grouping_ranker.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from collections import defaultdict

from haystack import Document, component
from haystack.utils.misc import _deduplicate_documents


@component
class MetaFieldGroupingRanker:
    """
    Reorders the documents by grouping them based on metadata keys.

    The MetaFieldGroupingRanker can group documents by a primary metadata key `group_by`, and subgroup them with an optional
    secondary key, `subgroup_by`.
    Within each group or subgroup, it can also sort documents by a metadata key `sort_docs_by`.

    The output is a flat list of documents ordered by `group_by` and `subgroup_by` values.
    Any documents without a group are placed at the end of the list.

    The proper organization of documents helps improve the efficiency and performance of subsequent processing by an LLM.

    ### Usage example

    ```python
    from haystack.components.rankers import MetaFieldGroupingRanker
    from haystack.dataclasses import Document


    docs = [
        Document(content="Javascript is a popular programming language", meta={"group": "42", "split_id": 7, "subgroup": "subB"}),
        Document(content="Python is a popular programming language",meta={"group": "42", "split_id": 4, "subgroup": "subB"}),
        Document(content="A chromosome is a package of DNA", meta={"group": "314", "split_id": 2, "subgroup": "subC"}),
        Document(content="An octopus has three hearts", meta={"group": "11", "split_id": 2, "subgroup": "subD"}),
        Document(content="Java is a popular programming language", meta={"group": "42", "split_id": 3, "subgroup": "subB"})
    ]

    ranker = MetaFieldGroupingRanker(group_by="group",subgroup_by="subgroup", sort_docs_by="split_id")
    result = ranker.run(documents=docs)
    print(result["documents"])

    # >> [
    # >>  Document(id=d665bbc83e52c08c3d8275bccf4f22bf2bfee21c6e77d78794627637355b8ebc,
    # >>          content: 'Java is a popular programming language', meta: {'group': '42', 'split_id': 3, 'subgroup': 'subB'}),
    # >>  Document(id=a20b326f07382b3cbf2ce156092f7c93e8788df5d48f2986957dce2adb5fe3c2,
    # >>          content: 'Python is a popular programming language', meta: {'group': '42', 'split_id': 4, 'subgroup': 'subB'}),
    # >>  Document(id=ce12919795d22f6ca214d0f161cf870993889dcb146f3bb1b3e1ffdc95be960f,
    # >>          content: 'Javascript is a popular programming language', meta: {'group': '42', 'split_id': 7, 'subgroup': 'subB'}),
    # >>  Document(id=d9fc857046c904e5cf790b3969b971b1bbdb1b3037d50a20728fdbf82991aa94,
    # >>          content: 'A chromosome is a package of DNA', meta: {'group': '314', 'split_id': 2, 'subgroup': 'subC'}),
    # >>  Document(id=6d3b7bdc13d09aa01216471eb5fb0bfdc53c5f2f3e98ad125ff6b85d3106c9a3,
    # >>          content: 'An octopus has three hearts', meta: {'group': '11', 'split_id': 2, 'subgroup': 'subD'})
    ```
    """  # noqa: E501

    def __init__(self, group_by: str, subgroup_by: str | None = None, sort_docs_by: str | None = None) -> None:
        """
        Creates an instance of MetaFieldGroupingRanker.

        :param group_by: The metadata key to aggregate the documents by.
        :param subgroup_by: The metadata key to aggregate the documents within a group that was created by the
                            `group_by` key.
        :param sort_docs_by: Determines which metadata key is used to sort the documents. If not provided, the
                             documents within the groups or subgroups are not sorted and are kept in the same order as
                             they were inserted in the subgroups.

        """
        self.group_by = group_by
        self.sort_docs_by = sort_docs_by
        self.subgroup_by = subgroup_by

    @component.output_types(documents=list[Document])
    def run(self, documents: list[Document]) -> dict[str, list[Document]]:
        """
        Groups the provided list of documents based on the `group_by` parameter and optionally the `subgroup_by`.

        Before grouping, documents are deduplicated by their id, retaining only the document with the highest score
        if a score is present.

        The output is a list of documents reordered based on how they were grouped.

        :param documents: The list of documents to group.
        :returns:
            A dictionary with the following keys:
            - documents: The list of documents ordered by the `group_by` and `subgroup_by` metadata values.
        """

        if not documents:
            return {"documents": []}

        document_groups: dict[str, dict[str, list[Document]]] = defaultdict(lambda: defaultdict(list))
        no_group_docs = []

        deduplicated_documents = _deduplicate_documents(documents)
        for doc in deduplicated_documents:
            group_value = str(doc.meta.get(self.group_by, ""))

            # If no group value, add to no_group_docs and continue
            if not group_value:
                no_group_docs.append(doc)
                continue

            # Get subgroup value or use a default if not specified
            subgroup_value = "no_subgroup"
            if self.subgroup_by and self.subgroup_by in doc.meta:
                subgroup_value = str(doc.meta[self.subgroup_by])

            document_groups[group_value][subgroup_value].append(doc)

        # use a non-optional key for type checking; "" disables sorting.
        sort_field = self.sort_docs_by or ""

        ordered_docs = []
        for subgroups in document_groups.values():
            for docs in subgroups.values():
                if sort_field:
                    # Sort by the field value, placing documents with a missing value last.
                    # The (is_missing, value) tuple ensures that only actual field values are
                    # compared, making the sort work for numbers, strings, and other types.
                    docs.sort(key=lambda d: (d.meta.get(sort_field) is None, d.meta.get(sort_field)))
                ordered_docs.extend(docs)

        ordered_docs.extend(no_group_docs)

        return {"documents": ordered_docs}