bedrock: add unit test for retriever (#21485)

This was implemented in
https://github.com/langchain-ai/langchain/pull/21349 but dropped before
merge.
pull/21496/head
ccurme 2 weeks ago committed by GitHub
parent 4035a1d234
commit 3bb9bec314
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -0,0 +1,68 @@
from typing import List
from unittest.mock import MagicMock
import pytest
from langchain_core.documents import Document
from langchain_community.retrievers import AmazonKnowledgeBasesRetriever
@pytest.fixture
def mock_client() -> MagicMock:
return MagicMock()
@pytest.fixture
def mock_retriever_config() -> dict:
return {"vectorSearchConfiguration": {"numberOfResults": 4}}
@pytest.fixture
def amazon_retriever(
mock_client: MagicMock, mock_retriever_config: dict
) -> AmazonKnowledgeBasesRetriever:
return AmazonKnowledgeBasesRetriever(
knowledge_base_id="test_kb_id",
retrieval_config=mock_retriever_config,
client=mock_client,
)
def test_create_client(amazon_retriever: AmazonKnowledgeBasesRetriever) -> None:
with pytest.raises(ImportError):
amazon_retriever.create_client({})
def test_get_relevant_documents(
amazon_retriever: AmazonKnowledgeBasesRetriever, mock_client: MagicMock
) -> None:
query: str = "test query"
mock_client.retrieve.return_value = {
"retrievalResults": [
{"content": {"text": "result1"}, "metadata": {"key": "value1"}},
{
"content": {"text": "result2"},
"metadata": {"key": "value2"},
"score": 1,
"location": "testLocation",
},
{"content": {"text": "result3"}},
]
}
documents: List[Document] = amazon_retriever._get_relevant_documents(
query,
run_manager=None, # type: ignore
)
assert len(documents) == 3
assert isinstance(documents[0], Document)
assert documents[0].page_content == "result1"
assert documents[0].metadata == {"score": 0, "source_metadata": {"key": "value1"}}
assert documents[1].page_content == "result2"
assert documents[1].metadata == {
"score": 1,
"source_metadata": {"key": "value2"},
"location": "testLocation",
}
assert documents[2].page_content == "result3"
assert documents[2].metadata == {"score": 0}
Loading…
Cancel
Save