Add guardrails support to AmazonKnowledgeBasesRetriever

libs/aws/langchain_aws/retrievers/bedrock.py

@@ -93,6 +93,9 @@ class AmazonKnowledgeBasesRetriever(BaseRetriever):
 
             client: boto3 client for bedrock agent runtime.
 
+            guardrail_config: Configuration information for a guardrail that you want
+                to use in the request.
+
             retrieval_config: Optional configuration for retrieval specified as a
                 Python object (RetrievalConfig) or as a dictionary
 
@@ -124,6 +127,9 @@ class AmazonKnowledgeBasesRetriever(BaseRetriever):
     endpoint_url: Optional[str] = None
     config: Any = None
     client: Any = None
+    guardrail_config: Optional[Dict[str, Any]] = Field(
+        default=None, alias="guardrails"
+    )
     retrieval_config: Optional[Union[RetrievalConfig, Dict[str, Any]]] = None
     min_score_confidence: Annotated[
         Optional[float], Field(ge=0.0, le=1.0, default=None)
@@ -132,6 +138,8 @@ class AmazonKnowledgeBasesRetriever(BaseRetriever):
     @model_validator(mode="before")
     @classmethod
     def create_client(cls, values: Dict[str, Any]) -> Any:
+        if "guardrail_config" in values and "guardrails" not in values:
+            values["guardrails"] = values.pop("guardrail_config")
         if values.get("client") is None:
             values["client"] = create_aws_client(
                 region_name=values.get("region_name"),
@@ -198,10 +206,21 @@ def _get_retrieve_request(self, query: str) -> Dict[str, Any]:
             "retrievalQuery": {"text": query.strip()},
             "knowledgeBaseId": self.knowledge_base_id,
         }
+        if self.guardrail_config:
+            if not (self.guardrail_config.get("guardrailId")
+                    and self.guardrail_config.get("guardrailVersion")):
+                raise TypeError(
+                    "Guardrail configuration must be a dictionary with both 'guardrailId' "
+                    "and 'guardrailVersion' keys."
+                )
+            request["guardrailConfiguration"] = self.guardrail_config
         if self.retrieval_config:
-            request["retrievalConfiguration"] = self.retrieval_config.model_dump(
-                exclude_none=True, by_alias=True
-            )
+            if isinstance(self.retrieval_config, dict):
+                request["retrievalConfiguration"] = self.retrieval_config
+            else:
+                request["retrievalConfiguration"] = self.retrieval_config.model_dump(
+                    exclude_none=True, by_alias=True
+                )
         return request
 
     @staticmethod

libs/aws/tests/unit_tests/retrievers/test_bedrock.py

@@ -568,3 +568,81 @@ def validate_query_response_with_cutoff(documents: List[Document]):
         "location": "testLocation",
         "type": "TEXT",
     }
+
+
+def test_guardrail_config(mock_client):
+    retriever = AmazonKnowledgeBasesRetriever(
+        knowledge_base_id="test_kb_id",
+        client=mock_client,
+        guardrail_config={
+            "guardrailId": "test-guardrail-id",
+            "guardrailVersion": "test-guardrail-version",
+        },
+    )
+
+    mock_client.retrieve.return_value = {
+        "retrievalResults": [
+            {"content": {"text": "result1"}, "metadata": {"key": "value1"}},
+        ]
+    }
+
+    retriever.invoke("test query")
+
+    mock_client.retrieve.assert_called_once_with(
+        retrievalQuery={"text": "test query"},
+        knowledgeBaseId="test_kb_id",
+        guardrailConfiguration={
+            "guardrailId": "test-guardrail-id",
+            "guardrailVersion": "test-guardrail-version",
+        },
+    )
+
+
+def test_guardrail_config_validation(mock_client):
+    retriever = AmazonKnowledgeBasesRetriever(
+        knowledge_base_id="test_kb_id",
+        client=mock_client,
+        guardrail_config={
+            "guardrailVersion": "test-guardrail-version",
+        },
+    )
+
+    with pytest.raises(TypeError) as excinfo:
+        retriever.invoke("test query")
+
+    assert "Guardrail configuration must be a dictionary with both 'guardrailId'" in str(excinfo.value)
+
+
+def test_guardrail_config_with_retrieval_config(mock_client, mock_retriever_config):
+    retriever = AmazonKnowledgeBasesRetriever(
+        knowledge_base_id="test_kb_id",
+        client=mock_client,
+        guardrail_config={
+            "guardrailId": "test-guardrail-id",
+            "guardrailVersion": "test-guardrail-version",
+        },
+        retrieval_config=mock_retriever_config,
+    )
+
+    mock_client.retrieve.return_value = {
+        "retrievalResults": [
+            {"content": {"text": "result1"}, "metadata": {"key": "value1"}},
+        ]
+    }
+
+    retriever.invoke("test query")
+
+    mock_client.retrieve.assert_called_once_with(
+        retrievalQuery={"text": "test query"},
+        knowledgeBaseId="test_kb_id",
+        guardrailConfiguration={
+            "guardrailId": "test-guardrail-id",
+            "guardrailVersion": "test-guardrail-version",
+        },
+        retrievalConfiguration={
+            "vectorSearchConfiguration": {
+                "numberOfResults": 5,
+                "filter": {"in": {"key": "key", "value": ["value1", "value2"]}},
+            }
+        },
+    )