From 06d33f19f285edd5cf8fcb323e8fca41c25afae5 Mon Sep 17 00:00:00 2001 From: Hikmat Dhamee Date: Thu, 20 Feb 2025 21:50:16 +0545 Subject: [PATCH 1/2] Currently, The BedrockAgentTarget doesn't seem to support KB filtering capability if I need to apply filters to KB for tests. The agent only uses associated KBs but doesn't apply any filters to vector stores. The proposed changes adds support for KB filtering capability to BedrockAgentTarget. I was trying to use this tool but realized that I need add this support as my use case needs KB filters. --- src/agenteval/targets/bedrock_agent/target.py | 3 +++ tests/src/agenteval/targets/bedrock_agent/test_target.py | 1 + 2 files changed, 4 insertions(+) diff --git a/src/agenteval/targets/bedrock_agent/target.py b/src/agenteval/targets/bedrock_agent/target.py index ebe3bc3..568ac79 100644 --- a/src/agenteval/targets/bedrock_agent/target.py +++ b/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: diff --git a/tests/src/agenteval/targets/bedrock_agent/test_target.py b/tests/src/agenteval/targets/bedrock_agent/test_target.py index a823a02..5df1c62 100644 --- a/tests/src/agenteval/targets/bedrock_agent/test_target.py +++ b/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 From ef0370cd98b14f62320a8c05051be73a6d488e94 Mon Sep 17 00:00:00 2001 From: Hikmat Dhamee Date: Thu, 20 Feb 2025 21:52:14 +0545 Subject: [PATCH 2/2] Currently, The BedrockAgentTarget doesn't seem to support KB filtering capability if I need to apply filters to KB for tests. The agent only uses associated KBs but doesn't apply any filters to vector stores. The proposed changes adds support for KB filtering capability to BedrockAgentTarget. I was trying to use this tool but realized that I need add this support as my use case needs KB filters. --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a19a96e..b20476e 100644 --- a/CHANGELOG.md +++ b/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.