Add support for KB filtering capability to Bedrock Agent Target

CHANGELOG.md

@@ -35,3 +35,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Initial release
+
+### Unreleased
+- add support for KB filtering capability to Bedrock Agent Target.

src/agenteval/targets/bedrock_agent/target.py

@@ -18,6 +18,7 @@ def __init__(
         bedrock_agent_alias_id: str,
         bedrock_session_attributes: Optional[dict] = {},
         bedrock_prompt_session_attributes: Optional[dict] = {},
+        knowledge_base_configurations: Optional[dict] = {},
         **kwargs
     ):
         """Initialize the target.
@@ -32,6 +33,8 @@ def __init__(
         self._bedrock_agent_id = bedrock_agent_id
         self._bedrock_agent_alias_id = bedrock_agent_alias_id
         self._session_state = {}
+        if knowledge_base_configurations:
+            self._session_state["knowledgeBaseConfigurations"] = knowledge_base_configurations
         if bedrock_session_attributes:
             self._session_state["sessionAttributes"] = bedrock_session_attributes
         if bedrock_prompt_session_attributes:

tests/src/agenteval/targets/bedrock_agent/test_target.py

@@ -17,6 +17,7 @@ def bedrock_agent_fixture(mocker):
         aws_region="us-west-2",
         bedrock_session_attributes={"first_name": "user_name"},
         bedrock_prompt_session_attributes={"timezone": "0"},
+        knowledge_base_configurations={}
     )
 
     return fixture