deepset-ai--haystack/test/components/retrievers/test_sentence_window_retriever.py

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

import random
import re
from unittest.mock import ANY

import pytest

from haystack import DeserializationError, Document, Pipeline
from haystack.components.preprocessors import DocumentSplitter
from haystack.components.retrievers import InMemoryBM25Retriever
from haystack.components.retrievers.sentence_window_retriever import SentenceWindowRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore


class TestSentenceWindowRetriever:
    def test_init_default(self, in_memory_doc_store):
        retriever = SentenceWindowRetriever(in_memory_doc_store)
        assert retriever.window_size == 3

    def test_init_with_parameters(self, in_memory_doc_store):
        retriever = SentenceWindowRetriever(in_memory_doc_store, window_size=5)
        assert retriever.window_size == 5

    def test_init_with_invalid_window_size_parameter(self, in_memory_doc_store):
        with pytest.raises(ValueError):
            SentenceWindowRetriever(in_memory_doc_store, window_size=-2)

    def test_merge_documents(self):
        docs = [
            {
                "id": "doc_0",
                "content": "This is a text with some words. There is a ",
                "source_id": "c5d7c632affc486d0cfe7b3c0f4dc1d3896ea720da2b538d6d10b104a3df5f99",
                "page_number": 1,
                "split_id": 0,
                "split_idx_start": 0,
                "_split_overlap": [{"doc_id": "doc_1", "range": (0, 23)}],
            },
            {
                "id": "doc_1",
                "content": "some words. There is a second sentence. And there is ",
                "source_id": "c5d7c632affc486d0cfe7b3c0f4dc1d3896ea720da2b538d6d10b104a3df5f99",
                "page_number": 1,
                "split_id": 1,
                "split_idx_start": 20,
                "_split_overlap": [{"doc_id": "doc_0", "range": (20, 43)}, {"doc_id": "doc_2", "range": (0, 29)}],
            },
            {
                "id": "doc_2",
                "content": "second sentence. And there is also a third sentence",
                "source_id": "c5d7c632affc486d0cfe7b3c0f4dc1d3896ea720da2b538d6d10b104a3df5f99",
                "page_number": 1,
                "split_id": 2,
                "split_idx_start": 43,
                "_split_overlap": [{"doc_id": "doc_1", "range": (23, 52)}],
            },
        ]
        merged_text = SentenceWindowRetriever.merge_documents_text([Document.from_dict(doc) for doc in docs])
        expected = "This is a text with some words. There is a second sentence. And there is also a third sentence"
        assert merged_text == expected

    def test_to_dict(self, in_memory_doc_store):
        window_retriever = SentenceWindowRetriever(in_memory_doc_store)
        data = window_retriever.to_dict()

        assert data == {
            "type": "haystack.components.retrievers.sentence_window_retriever.SentenceWindowRetriever",
            "init_parameters": {
                "document_store": {
                    "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore",
                    "init_parameters": {
                        "bm25_algorithm": "BM25L",
                        "bm25_parameters": {},
                        "bm25_tokenization_regex": "(?u)\\b\\w+\\b",
                        "embedding_similarity_function": "dot_product",
                        "index": ANY,
                        "shared": True,
                        "return_embedding": True,
                    },
                },
                "window_size": 3,
                "source_id_meta_field": "source_id",
                "split_id_meta_field": "split_id",
                "raise_on_missing_meta_fields": True,
            },
        }

    def test_from_dict(self):
        data = {
            "type": "haystack.components.retrievers.sentence_window_retriever.SentenceWindowRetriever",
            "init_parameters": {
                "document_store": {
                    "type": "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore",
                    "init_parameters": {},
                },
                "window_size": 5,
                "source_id_meta_field": "source_id_test",
                "split_id_meta_field": "split_id_test",
                "raise_on_missing_meta_fields": False,
            },
        }
        component = SentenceWindowRetriever.from_dict(data)
        assert isinstance(component.document_store, InMemoryDocumentStore)
        assert component.window_size == 5
        assert component.source_id_meta_field == "source_id_test"
        assert component.split_id_meta_field == "split_id_test"
        assert not component.raise_on_missing_meta_fields

    def test_from_dict_without_docstore(self):
        data = {
            "type": "haystack.components.retrievers.sentence_window_retriever.SentenceWindowRetriever",
            "init_parameters": {},
        }
        with pytest.raises(TypeError, match="missing 1 required positional argument: 'document_store'"):
            SentenceWindowRetriever.from_dict(data)

    def test_from_dict_without_docstore_type(self):
        data = {"type": "SentenceWindowRetriever", "init_parameters": {"document_store": {"init_parameters": {}}}}
        with pytest.raises(DeserializationError):
            SentenceWindowRetriever.from_dict(data)

    def test_from_dict_non_existing_docstore(self):
        data = {
            "type": "SentenceWindowRetriever",
            "init_parameters": {"document_store": {"type": "Nonexisting.Docstore", "init_parameters": {}}},
        }
        with pytest.raises(DeserializationError):
            SentenceWindowRetriever.from_dict(data)

    def test_document_without_split_id(self, in_memory_doc_store):
        docs = [
            Document(content="This is a text with some words. There is a ", meta={"id": "doc_0"}),
            Document(content="some words. There is a second sentence. And there is ", meta={"id": "doc_1"}),
        ]
        with pytest.raises(ValueError, match="The retrieved documents must have 'split_id_test' in their metadata."):
            retriever = SentenceWindowRetriever(
                document_store=in_memory_doc_store, window_size=3, split_id_meta_field="split_id_test"
            )
            retriever.run(retrieved_documents=docs)

    def test_document_without_source_id(self, in_memory_doc_store):
        docs = [
            Document(content="This is a text with some words. There is a ", meta={"id": "doc_0", "split_id": 0}),
            Document(
                content="some words. There is a second sentence. And there is ",
                meta={"id": "doc_1", "split_id": 1, "source_id_test": "source1"},
            ),
        ]
        with pytest.raises(ValueError, match="The retrieved documents must have 'source_id_test' in their metadata."):
            retriever = SentenceWindowRetriever(
                document_store=in_memory_doc_store, window_size=3, source_id_meta_field="source_id_test"
            )
            retriever.run(retrieved_documents=docs)

    def test_document_without_all_source_ids(self, in_memory_doc_store):
        docs = [
            Document(
                content="These are words from the first section",
                meta={"id": "doc_1", "split_id": 0, "section_id": "section1"},
            ),
            Document(
                content="These are words from the second section, but missing section_id",
                meta={"id": "doc_0", "split_id": 0},
            ),
        ]
        with pytest.raises(
            ValueError, match=re.escape("The retrieved documents must have '['id', 'section_id']' in their metadata.")
        ):
            retriever = SentenceWindowRetriever(
                document_store=in_memory_doc_store, window_size=3, source_id_meta_field=["id", "section_id"]
            )
            retriever.run(retrieved_documents=docs)

    def test_run_invalid_window_size(self, in_memory_doc_store):
        docs = [Document(content="This is a text with some words. There is a ", meta={"id": "doc_0", "split_id": 0})]
        with pytest.raises(ValueError):
            retriever = SentenceWindowRetriever(document_store=in_memory_doc_store, window_size=0)
            retriever.run(retrieved_documents=docs)

    def test_constructor_parameter_does_not_change(self, in_memory_doc_store):
        retriever = SentenceWindowRetriever(in_memory_doc_store, window_size=5)
        assert retriever.window_size == 5

        doc = {
            "id": "doc_0",
            "content": "This is a text with some words. There is a ",
            "source_id": "c5d7c632affc486d0cfe7b3c0f4dc1d3896ea720da2b538d6d10b104a3df5f99",
            "page_number": 1,
            "split_id": 0,
            "split_idx_start": 0,
            "_split_overlap": [{"doc_id": "doc_1", "range": (0, 23)}],
        }

        retriever.run(retrieved_documents=[Document.from_dict(doc)], window_size=1)
        assert retriever.window_size == 5

    def test_context_documents_returned_are_ordered_by_split_idx_start(self, in_memory_doc_store):
        docs = []
        accumulated_length = 0
        for sent in range(10):
            content = f"Sentence {sent}."
            docs.append(
                Document(
                    content=content,
                    meta={
                        "id": f"doc_{sent}",
                        "split_idx_start": accumulated_length,
                        "source_id": "source1",
                        "split_id": sent,
                    },
                )
            )
            accumulated_length += len(content)

        random.shuffle(docs)

        in_memory_doc_store.write_documents(docs)
        retriever = SentenceWindowRetriever(document_store=in_memory_doc_store, window_size=3)

        # run the retriever with a document whose content = "Sentence 4."
        result = retriever.run(retrieved_documents=[doc for doc in docs if doc.content == "Sentence 4."])

        # assert that the context documents are in the correct order
        assert len(result["context_documents"]) == 7
        assert [doc.meta["split_idx_start"] for doc in result["context_documents"]] == [11, 22, 33, 44, 55, 66, 77]

    def test_run_custom_fields(self, in_memory_doc_store):
        docs = []
        accumulated_length = 0
        for sent in range(10):
            content = f"Sentence {sent}."
            docs.append(
                Document(
                    content=content,
                    meta={
                        "id": f"doc_{sent}",
                        # Missing split_idx_start
                        "source_id_test": "source1",
                        "split_id_test": sent,
                    },
                )
            )
            accumulated_length += len(content)

        random.shuffle(docs)

        in_memory_doc_store.write_documents(docs)
        retriever = SentenceWindowRetriever(
            document_store=in_memory_doc_store,
            window_size=3,
            source_id_meta_field="source_id_test",
            split_id_meta_field="split_id_test",
        )

        # run the retriever with a document whose content = "Sentence 4."
        result = retriever.run(retrieved_documents=[doc for doc in docs if doc.content == "Sentence 4."])
        assert len(result["context_documents"]) == 7

    def test_run_with_multiple_source_ids(self, in_memory_doc_store):
        docs = [
            Document(content="This is the first chunk.", meta={"section": "1", "split_id": 0, "source_id": "source1"}),
            Document(content="This is the second chunk.", meta={"section": "1", "split_id": 1, "source_id": "source1"}),
            Document(content="This is the third chunk.", meta={"section": "1", "split_id": 2, "source_id": "source1"}),
            Document(
                content="This is a chunk from section 2.", meta={"section": "2", "split_id": 3, "source_id": "source1"}
            ),
        ]
        in_memory_doc_store.write_documents(docs)

        retriever = SentenceWindowRetriever(
            document_store=in_memory_doc_store, window_size=5, source_id_meta_field=["section", "source_id"]
        )
        result = retriever.run(
            retrieved_documents=[
                Document(
                    content="This is the second chunk.", meta={"section": "1", "split_id": 1, "source_id": "source1"}
                )
            ]
        )

        assert len(result["context_windows"]) == 1
        assert len(result["context_documents"]) == 3
        assert all(doc.meta["section"] == "1" for doc in result["context_documents"])

    @pytest.mark.integration
    def test_run_with_pipeline(self, in_memory_doc_store):
        splitter = DocumentSplitter(split_length=1, split_overlap=0, split_by="period")
        text = (
            "This is a text with some words. There is a second sentence. And there is also a third sentence. "
            "It also contains a fourth sentence. And a fifth sentence. And a sixth sentence. And a seventh sentence"
        )
        doc = Document(content=text)
        docs = splitter.run([doc])
        in_memory_doc_store.write_documents(docs["documents"])

        pipe = Pipeline()
        pipe.add_component("bm25_retriever", InMemoryBM25Retriever(in_memory_doc_store, top_k=1))
        pipe.add_component(
            "sentence_window_retriever", SentenceWindowRetriever(document_store=in_memory_doc_store, window_size=2)
        )
        pipe.connect("bm25_retriever", "sentence_window_retriever")
        result = pipe.run({"bm25_retriever": {"query": "third"}})

        assert result["sentence_window_retriever"]["context_windows"] == [
            "This is a text with some words. There is a second sentence. And there is also a third sentence. "
            "It also contains a fourth sentence. And a fifth sentence."
        ]
        assert len(result["sentence_window_retriever"]["context_documents"]) == 5

        result = pipe.run({"bm25_retriever": {"query": "third"}, "sentence_window_retriever": {"window_size": 1}})
        assert result["sentence_window_retriever"]["context_windows"] == [
            " There is a second sentence. And there is also a third sentence. It also contains a fourth sentence."
        ]
        assert len(result["sentence_window_retriever"]["context_documents"]) == 3

    @pytest.mark.integration
    def test_serialization_deserialization_in_pipeline(self, in_memory_doc_store):
        pipe = Pipeline()
        pipe.add_component("bm25_retriever", InMemoryBM25Retriever(in_memory_doc_store, top_k=1))
        pipe.add_component(
            "sentence_window_retriever", SentenceWindowRetriever(document_store=in_memory_doc_store, window_size=2)
        )
        pipe.connect("bm25_retriever", "sentence_window_retriever")

        serialized = pipe.to_dict()
        deserialized = Pipeline.from_dict(serialized)

        assert deserialized == pipe