Function test_generate_answer

sdk/python/tests/unit/test_rag_retriever.py:415–456

Test generating an answer using the RAG retriever.

(rag_retriever)

# End-to-end functionality test
def test_generate_answer(rag_retriever):
    """Test generating an answer using the RAG retriever."""
    # Mock the retrieve method using patch
    with patch.object(
        rag_retriever,
        "retrieve",
        return_value=(
            np.array([[[0.1] * 8, [0.2] * 8]]),  # 8-dimensional embeddings
            np.array([[1, 2]]),
            [
                {
                    "text": ["context1", "context2"],
                    "id": ["doc1", "doc2"],
                    "title": ["Doc 1", "Doc 2"],
                }
            ],
        ),
    ) as mock_retrieve:
        # Mock the generator model's generate method
        rag_retriever.generator_model.generate = Mock(
            return_value=torch.tensor([[1, 2, 3]])
        )

        # Generate an answer
        answer = rag_retriever.generate_answer(
            "test query", top_k=2, max_new_tokens=100
        )

        # Verify the answer
        assert isinstance(answer, str)
        assert len(answer) > 0

        # Verify that retrieve was called with correct parameters
        mock_retrieve.assert_called_once()
        call_args = mock_retrieve.call_args[1]
        assert call_args["n_docs"] == 2
        assert call_args["query"] == "test query"

        # Verify that generate was called with correct parameters
        rag_retriever.generator_model.generate.assert_called_once()
        call_args = rag_retriever.generator_model.generate.call_args[1]
        assert call_args["max_new_tokens"] == 100