langchain/libs/standard-tests/langchain_tests/integration_tests/retrievers.py

from abc import abstractmethod

import pytest
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever

from langchain_tests.base import BaseStandardTests


class RetrieversIntegrationTests(BaseStandardTests):
    """Base class for retrievers integration tests."""

    @property
    @abstractmethod
    def retriever_constructor(self) -> type[BaseRetriever]:
        """A ``BaseRetriever`` subclass to be tested."""
        ...

    @property
    def retriever_constructor_params(self) -> dict:
        """Returns a dictionary of parameters to pass to the retriever constructor."""
        return {}

    @property
    @abstractmethod
    def retriever_query_example(self) -> str:
        """Returns a str representing the ``query`` of an example retriever call."""

    @property
    def num_results_arg_name(self) -> str:
        """Returns the name of the parameter for the number of results returned.

        Usually something like ``k`` or ``top_k``.

        """
        return "k"

    @pytest.fixture
    def retriever(self) -> BaseRetriever:
        """:private:"""
        return self.retriever_constructor(**self.retriever_constructor_params)

    def test_k_constructor_param(self) -> None:
        """Test the number of results constructor parameter.

        Test that the retriever constructor accepts a parameter representing
        the number of documents to return.

        By default, the parameter tested is named ``k``, but it can be overridden by
        setting the ``num_results_arg_name`` property.

        .. note::
            If the retriever doesn't support configuring the number of results returned
            via the constructor, this test can be skipped using a pytest ``xfail`` on
            the test class:

            .. code-block:: python

                @pytest.mark.xfail(
                    reason="This retriever doesn't support setting "
                    "the number of results via the constructor."
                )
                def test_k_constructor_param(self) -> None:
                    raise NotImplementedError

        .. dropdown:: Troubleshooting

            If this test fails, the retriever constructor does not accept a number
            of results parameter, or the retriever does not return the correct number
            of documents ( of the one set in ``num_results_arg_name``) when it is
            set.

            For example, a retriever like

            .. code-block:: python

                    MyRetriever(k=3).invoke("query")

            should return 3 documents when invoked with a query.

        """
        params = {
            k: v
            for k, v in self.retriever_constructor_params.items()
            if k != self.num_results_arg_name
        }
        params_3 = {**params, self.num_results_arg_name: 3}
        retriever_3 = self.retriever_constructor(**params_3)
        result_3 = retriever_3.invoke(self.retriever_query_example)
        assert len(result_3) == 3
        assert all(isinstance(doc, Document) for doc in result_3)

        params_1 = {**params, self.num_results_arg_name: 1}
        retriever_1 = self.retriever_constructor(**params_1)
        result_1 = retriever_1.invoke(self.retriever_query_example)
        assert len(result_1) == 1
        assert all(isinstance(doc, Document) for doc in result_1)

    def test_invoke_with_k_kwarg(self, retriever: BaseRetriever) -> None:
        """Test the number of results parameter in ``invoke()``.

        Test that the invoke method accepts a parameter representing
        the number of documents to return.

        By default, the parameter is named ``, but it can be overridden by
        setting the ``num_results_arg_name`` property.

        .. note::
            If the retriever doesn't support configuring the number of results returned
            via the invoke method, this test can be skipped using a pytest ``xfail`` on
            the test class:

            .. code-block:: python

                @pytest.mark.xfail(
                    reason="This retriever doesn't support setting "
                    "the number of results in the invoke method."
                )
                def test_invoke_with_k_kwarg(self) -> None:
                    raise NotImplementedError

        .. dropdown:: Troubleshooting

            If this test fails, the retriever's invoke method does not accept a number
            of results parameter, or the retriever does not return the correct number
            of documents (``k`` of the one set in ``num_results_arg_name``) when it is
            set.

            For example, a retriever like

            .. code-block:: python

                MyRetriever().invoke("query", k=3)

            should return 3 documents when invoked with a query.

        """
        result_1 = retriever.invoke(
            self.retriever_query_example, None, **{self.num_results_arg_name: 1}
        )
        assert len(result_1) == 1
        assert all(isinstance(doc, Document) for doc in result_1)

        result_3 = retriever.invoke(
            self.retriever_query_example, None, **{self.num_results_arg_name: 3}
        )
        assert len(result_3) == 3
        assert all(isinstance(doc, Document) for doc in result_3)

    def test_invoke_returns_documents(self, retriever: BaseRetriever) -> None:
        """If invoked with the example params, the retriever should return a list of
        Documents.

        .. dropdown:: Troubleshooting

            If this test fails, the retriever's invoke method does not return a list of
            ``langchain_core.document.Document`` objects. Please confirm that your
            ``_get_relevant_documents`` method returns a list of ``Document`` objects.
        """
        result = retriever.invoke(self.retriever_query_example)

        assert isinstance(result, list)
        assert all(isinstance(doc, Document) for doc in result)

    async def test_ainvoke_returns_documents(self, retriever: BaseRetriever) -> None:
        """If ainvoked with the example params, the retriever should return a list of
        Documents.

        See :meth:`test_invoke_returns_documents` for more information on
        troubleshooting.
        """
        result = await retriever.ainvoke(self.retriever_query_example)

        assert isinstance(result, list)
        assert all(isinstance(doc, Document) for doc in result)