Added tests for knowlegebases retriever

libs/aws/langchain_aws/__init__.py

@@ -1,4 +1,11 @@
 from langchain_aws.llms import SagemakerEndpoint
-from langchain_aws.retrievers import AmazonKendraRetriever
+from langchain_aws.retrievers import (
+    AmazonKendraRetriever,
+    AmazonKnowledgeBasesRetriever,
+)
 
-__all__ = ["SagemakerEndpoint", "AmazonKendraRetriever"]
+__all__ = [
+    "SagemakerEndpoint",
+    "AmazonKendraRetriever",
+    "AmazonKnowledgeBasesRetriever",
+]

libs/aws/tests/integration_tests/retrievers/test_amazon_knowledgebases_retriever.py

@@ -0,0 +1,66 @@
+from unittest.mock import Mock
+
+import pytest
+from langchain_core.documents import Document
+
+from langchain_aws import AmazonKnowledgeBasesRetriever
+
+
+@pytest.fixture
+def mock_client() -> Mock:
+    return Mock()
+
+
+@pytest.fixture
+def retriever(mock_client: Mock) -> AmazonKnowledgeBasesRetriever:
+    return AmazonKnowledgeBasesRetriever(
+        knowledge_base_id="test-knowledge-base",
+        client=mock_client,
+        retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 4}},  # type: ignore[arg-type]
+    )
+
+
+def test_get_relevant_documents(retriever, mock_client) -> None:  # type: ignore[no-untyped-def]
+    response = {
+        "retrievalResults": [
+            {
+                "content": {"text": "This is the first result."},
+                "location": "location1",
+                "score": 0.9,
+            },
+            {
+                "content": {"text": "This is the second result."},
+                "location": "location2",
+                "score": 0.8,
+            },
+            {"content": {"text": "This is the third result."}, "location": "location3"},
+        ]
+    }
+    mock_client.retrieve.return_value = response
+
+    query = "test query"
+
+    expected_documents = [
+        Document(
+            page_content="This is the first result.",
+            metadata={"location": "location1", "score": 0.9},
+        ),
+        Document(
+            page_content="This is the second result.",
+            metadata={"location": "location2", "score": 0.8},
+        ),
+        Document(
+            page_content="This is the third result.",
+            metadata={"location": "location3", "score": 0.0},
+        ),
+    ]
+
+    documents = retriever.get_relevant_documents(query)
+
+    assert documents == expected_documents
+
+    mock_client.retrieve.assert_called_once_with(
+        retrievalQuery={"text": "test query"},
+        knowledgeBaseId="test-knowledge-base",
+        retrievalConfiguration={"vectorSearchConfiguration": {"numberOfResults": 4}},
+    )