retrievers.py — langchain Source File

Source Code

"""Integration tests for retrievers."""

from abc import abstractmethod
from typing import Any

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[str, Any]:
        """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:
        """Return retriever fixture."""
        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:

            ```python
// ... (123 more lines)