diff --git a/.doc_gen/metadata/bedrock-agent_metadata.yaml b/.doc_gen/metadata/bedrock-agent_metadata.yaml index ce7ea8e8345..cf3486fe0de 100644 --- a/.doc_gen/metadata/bedrock-agent_metadata.yaml +++ b/.doc_gen/metadata/bedrock-agent_metadata.yaml @@ -111,6 +111,38 @@ bedrock-agent_ListAgents: services: bedrock-agent: {ListAgents} +bedrock-agent_ListAgentActionGroups: + title: List the action groups for a &BR; agent using an &AWS; SDK + title_abbrev: List the action groups for an agent + synopsis: list the action groups for a &BR; agent. + languages: + Python: + versions: + - sdk_version: 3 + github: python/example_code/bedrock-agent + excerpts: + - description: List the action groups for an agent. + snippet_tags: + - python.example_code.bedrock-agent.ListAgentActionGroups + services: + bedrock-agent: {ListAgentActionGroups} + +bedrock-agent_ListAgentKnowledgeBases: + title: List the knowledge bases associated with a &BR; agent using an &AWS; SDK + title_abbrev: List the knowledge bases associated with an agent + synopsis: list the knowledge bases associated with a &BR; agent. + languages: + Python: + versions: + - sdk_version: 3 + github: python/example_code/bedrock-agent + excerpts: + - description: List the knowledge bases associated with an agent. + snippet_tags: + - python.example_code.bedrock-agent.ListAgentKnowledgeBases + services: + bedrock-agent: {ListAgentKnowledgeBases} + bedrock-agent_PrepareAgent: title: Prepare an &BR; agent using an &AWS; SDK title_abbrev: Prepare an agent @@ -128,7 +160,7 @@ bedrock-agent_PrepareAgent: bedrock-agent: {PrepareAgent} bedrock-agent_GettingStartedWithBedrockAgents: title: An end-to-end example showing how to create and invoke &BR; agents using an &AWS; SDK - title_abbrev: Create and invoke agents + title_abbrev: Create and invoke an agent synopsis_list: - Create an execution role for the agent. - Create the agent and deploy a DRAFT version. @@ -144,7 +176,7 @@ bedrock-agent_GettingStartedWithBedrockAgents: - sdk_version: 3 github: python/example_code/bedrock-agent excerpts: - - description: + - description: Create and invoke an agent. snippet_tags: - python.example_code.bedrock-agent.Scenario_GettingStartedBedrockAgents services: @@ -156,7 +188,9 @@ bedrock-agent_GettingStartedWithBedrockAgents: DeleteAgent, DeleteAgentAlias, GetAgent, + ListAgentActionGroups, ListAgents, + ListAgentKnowledgeBases, PrepareAgent, } bedrock-agent-runtime: diff --git a/python/example_code/bedrock-agent-runtime/README.md b/python/example_code/bedrock-agent-runtime/README.md index 561b50f1e7f..ec12e79298a 100644 --- a/python/example_code/bedrock-agent-runtime/README.md +++ b/python/example_code/bedrock-agent-runtime/README.md @@ -32,9 +32,6 @@ python -m pip install -r requirements.txt ``` - -> ⚠ You must request access to a model before you can use it. If you try to use the model (with the API or console) before you have requested access to it, you will receive an error message. For more information, see [Model access](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html). - ### Single actions @@ -45,19 +42,6 @@ Code excerpts that show you how to call individual service functions. - -#### Create and invoke agents - -This example shows you how to do the following: - -- Create an execution role for the agent. -- Create the agent and deploy a DRAFT version. -- Create a Lambda function that implements the agent's capabilities. -- Create an action group that connects the agent to the Lambda function. -- Deploy the fully configured agent. -- Invoke the agent with user-provided prompts. -- Delete all created resources. - ## Run the examples diff --git a/python/example_code/bedrock-agent/README.md b/python/example_code/bedrock-agent/README.md index 904447a1613..7f5581bdb95 100644 --- a/python/example_code/bedrock-agent/README.md +++ b/python/example_code/bedrock-agent/README.md @@ -44,15 +44,17 @@ Code excerpts that show you how to call individual service functions. - [Delete an agent](bedrock_agent_wrapper.py#L119) (`DeleteAgent`) - [Delete an agent alias](bedrock_agent_wrapper.py#L140) (`DeleteAgentAlias`) - [Get information about an agent](bedrock_agent_wrapper.py#L162) (`GetAgent`) +- [List the action groups for an agent](bedrock_agent_wrapper.py#L209) (`ListAgentActionGroups`) - [List the agents](bedrock_agent_wrapper.py#L186) (`ListAgents`) -- [Prepare an agent](bedrock_agent_wrapper.py#L209) (`PrepareAgent`) +- [List the knowledge bases associated with an agent](bedrock_agent_wrapper.py#L238) (`ListAgentKnowledgeBases`) +- [Prepare an agent](bedrock_agent_wrapper.py#L267) (`PrepareAgent`) ### Scenarios Code examples that show you how to accomplish a specific task by calling multiple functions within the same service. -- [Create and invoke agents](scenario_get_started_with_agents.py) +- [Create and invoke an agent](scenario_get_started_with_agents.py) @@ -68,7 +70,7 @@ functions within the same service. -#### Create and invoke agents +#### Create and invoke an agent This example shows you how to do the following: diff --git a/python/example_code/bedrock-agent/bedrock_agent_wrapper.py b/python/example_code/bedrock-agent/bedrock_agent_wrapper.py index d60ba77ddff..e4687d5cc10 100644 --- a/python/example_code/bedrock-agent/bedrock_agent_wrapper.py +++ b/python/example_code/bedrock-agent/bedrock_agent_wrapper.py @@ -206,6 +206,64 @@ def list_agents(self): # snippet-end:[python.example_code.bedrock-agent.ListAgents] + # snippet-start:[python.example_code.bedrock-agent.ListAgentActionGroups] + def list_agent_action_groups(self, agent_id, agent_version): + """ + List the action groups for a version of an Amazon Bedrock Agent. + + :param agent_id: The unique identifier of the agent. + :param agent_version: The version of the agent. + :return: The list of action group summaries for the version of the agent. + """ + + try: + action_groups = [] + + paginator = self.client.get_paginator("list_agent_action_groups") + for page in paginator.paginate( + agentId=agent_id, + agentVersion=agent_version, + PaginationConfig={"PageSize": 10}, + ): + action_groups.extend(page["actionGroupSummaries"]) + + except ClientError as e: + logger.error(f"Couldn't list action groups. {e}") + raise + else: + return action_groups + + # snippet-end:[python.example_code.bedrock-agent.ListAgentActionGroups] + + # snippet-start:[python.example_code.bedrock-agent.ListAgentKnowledgeBases] + def list_agent_knowledge_bases(self, agent_id, agent_version): + """ + List the knowledge bases associated with a version of an Amazon Bedrock Agent. + + :param agent_id: The unique identifier of the agent. + :param agent_version: The version of the agent. + :return: The list of knowledge base summaries for the version of the agent. + """ + + try: + knowledge_bases = [] + + paginator = self.client.get_paginator("list_agent_knowledge_bases") + for page in paginator.paginate( + agentId=agent_id, + agentVersion=agent_version, + PaginationConfig={"PageSize": 10}, + ): + knowledge_bases.extend(page["agentKnowledgeBaseSummaries"]) + + except ClientError as e: + logger.error(f"Couldn't list knowledge bases. {e}") + raise + else: + return knowledge_bases + + # snippet-end:[python.example_code.bedrock-agent.ListAgentKnowledgeBases] + # snippet-start:[python.example_code.bedrock-agent.PrepareAgent] def prepare_agent(self, agent_id): """ diff --git a/python/example_code/bedrock-agent/scenario_get_started_with_agents.py b/python/example_code/bedrock-agent/scenario_get_started_with_agents.py index 1eb62e94816..0adf4b4da67 100644 --- a/python/example_code/bedrock-agent/scenario_get_started_with_agents.py +++ b/python/example_code/bedrock-agent/scenario_get_started_with_agents.py @@ -26,6 +26,7 @@ import re import string import sys +import uuid import yaml import zipfile @@ -39,12 +40,16 @@ import demo_tools.question as q from demo_tools.retries import wait +logger = logging.getLogger(__name__) + +# snippet-start:[python.example_code.bedrock-agent.Scenario_GettingStartedBedrockAgents] REGION = "us-east-1" ROLE_POLICY_NAME = "agent_permissions" -logger = logging.getLogger(__name__) class BedrockAgentScenarioWrapper: + """Runs a scenario that shows how to get started using Agents for Amazon Bedrock.""" + def __init__( self, bedrock_agent_client, runtime_client, lambda_client, iam_resource, postfix ): @@ -58,11 +63,10 @@ def __init__( self.agent = None self.agent_alias = None self.agent_role = None - self.agent_version = None + self.prepared_agent_details = None self.lambda_role = None self.lambda_function = None - # snippet-start:[python.example_code.bedrock-agent.Scenario_GettingStartedBedrockAgents] def run_scenario(self): print("=" * 88) print("Welcome to the Amazon Bedrock Agents demo.") @@ -71,36 +75,42 @@ def run_scenario(self): # Query input from user print("Let's start with creating an agent:") print("-" * 40) - name, model_id = self._request_name_and_model_from_user() + name, foundation_model = self._request_name_and_model_from_user() print("-" * 40) # Create an execution role for the agent - self.agent_role = self._create_agent_role(model_id) + self.agent_role = self._create_agent_role(foundation_model) # Create the agent - self.agent = self._create_agent(name, model_id) + self.agent = self._create_agent(name, foundation_model) # Prepare a DRAFT version of the agent - self.agent_version = self._prepare_agent() + self.prepared_agent_details = self._prepare_agent() - # Create the lambda function + # Create the agent's Lambda function self.lambda_function = self._create_lambda_function() # Configure permissions for the agent to invoke the Lambda function - self._allow_agent_to_invoke_function(self.lambda_function) - self._let_function_accept_invocations_from_agent(self.agent) + self._allow_agent_to_invoke_function() + self._let_function_accept_invocations_from_agent() - # Create an action group that connects the agent to the Lambda function + # Create an action group to connects the agent with the Lambda function self._create_agent_action_group() - # Prepare the DRAFT version of the agent, including the action group - self.agent_version = self._prepare_agent() + # If the agent has been modified or any components have been added, prepare the agent again + components = [self._get_agent()] + components += self._get_agent_action_groups() + components += self._get_agent_knowledge_bases() + + latest_update = max(component["updatedAt"] for component in components) + if latest_update > self.prepared_agent_details["preparedAt"]: + self.prepared_agent_details = self._prepare_agent() # Create an agent alias self.agent_alias = self._create_agent_alias() # Test the agent - self._chat_with_agent() + self._chat_with_agent(self.agent_alias) print("=" * 88) print("Thanks for running the demo!\n") @@ -116,8 +126,6 @@ def run_scenario(self): print("=" * 88) print("Thanks again for running the demo!") - # snippet-end:[python.example_code.bedrock-agent.Scenario_GettingStartedBedrockAgents] - def _request_name_and_model_from_user(self): existing_agent_names = [ agent["agentName"] for agent in self.bedrock_wrapper.list_agents() @@ -203,10 +211,10 @@ def _prepare_agent(self): print("Preparing the agent...") agent_id = self.agent["agentId"] - version = self.bedrock_wrapper.prepare_agent(agent_id)["agentVersion"] + prepared_agent_details = self.bedrock_wrapper.prepare_agent(agent_id) self._wait_for_agent_status(agent_id, "PREPARED") - return version + return prepared_agent_details def _create_lambda_function(self): print("Creating the Lambda function...") @@ -270,7 +278,7 @@ def _create_lambda_role(self): return role - def _allow_agent_to_invoke_function(self, lambda_function): + def _allow_agent_to_invoke_function(self): policy = self.iam_resource.RolePolicy( self.agent_role.role_name, ROLE_POLICY_NAME ) @@ -279,19 +287,19 @@ def _allow_agent_to_invoke_function(self, lambda_function): { "Effect": "Allow", "Action": "lambda:InvokeFunction", - "Resource": lambda_function["FunctionArn"], + "Resource": self.lambda_function["FunctionArn"], } ) self.agent_role.Policy(ROLE_POLICY_NAME).put(PolicyDocument=json.dumps(doc)) - def _let_function_accept_invocations_from_agent(self, agent): + def _let_function_accept_invocations_from_agent(self): try: self.lambda_client.add_permission( FunctionName=self.lambda_function["FunctionName"], + SourceArn=self.agent["agentArn"], StatementId="BedrockAccess", Action="lambda:InvokeFunction", Principal="bedrock.amazonaws.com", - SourceArn=agent["agentArn"], ) except ClientError as e: logger.error( @@ -308,7 +316,7 @@ def _create_agent_action_group(self): name="current_date_and_time", description="Gets the current date and time.", agent_id=self.agent["agentId"], - agent_version=self.agent_version, + agent_version=self.prepared_agent_details["agentVersion"], function_arn=self.lambda_function["FunctionArn"], api_schema=json.dumps(yaml.safe_load(file)), ) @@ -316,6 +324,19 @@ def _create_agent_action_group(self): logger.error(f"Couldn't create agent action group. Here's why: {e}") raise + def _get_agent(self): + return self.bedrock_wrapper.get_agent(self.agent["agentId"]) + + def _get_agent_action_groups(self): + return self.bedrock_wrapper.list_agent_action_groups( + self.agent["agentId"], self.prepared_agent_details["agentVersion"] + ) + + def _get_agent_knowledge_bases(self): + return self.bedrock_wrapper.list_agent_knowledge_bases( + self.agent["agentId"], self.prepared_agent_details["agentVersion"] + ) + def _create_agent_alias(self): print("Creating an agent alias...") @@ -332,26 +353,29 @@ def _wait_for_agent_status(self, agent_id, status): while self.bedrock_wrapper.get_agent(agent_id)["agentStatus"] != status: wait(2) - def _chat_with_agent(self): + def _chat_with_agent(self, agent_alias): print("-" * 88) print("The agent is ready to chat.") print("Try asking for the date or time. Type 'exit' to quit.") + # Create a unique session ID for the conversation + session_id = uuid.uuid4().hex + while True: prompt = q.ask("Prompt: ", q.non_empty) if prompt == "exit": break - response = asyncio.run(self._invoke_agent(prompt)) + response = asyncio.run(self._invoke_agent(agent_alias, prompt, session_id)) print(f"Agent: {response}") - async def _invoke_agent(self, prompt): + async def _invoke_agent(self, agent_alias, prompt, session_id): response = self.bedrock_agent_runtime_client.invoke_agent( agentId=self.agent["agentId"], - agentAliasId=self.agent_alias["agentAliasId"], - sessionId="Session", + agentAliasId=agent_alias["agentAliasId"], + sessionId=session_id, inputText=prompt, ) @@ -369,12 +393,10 @@ def _delete_resources(self): if self.agent_alias: agent_alias_id = self.agent_alias["agentAliasId"] - agent_alias_name = self.agent_alias["agentAliasName"] - print(f"Deleting agent alias '{agent_alias_name}'...") + print("Deleting agent alias...") self.bedrock_wrapper.delete_agent_alias(agent_id, agent_alias_id) - agent_name = self.agent["agentName"] - print(f"Deleting agent '{agent_name}'...") + print("Deleting agent...") agent_status = self.bedrock_wrapper.delete_agent(agent_id)["agentStatus"] while agent_status == "DELETING": wait(5) @@ -457,3 +479,5 @@ def _create_deployment_package(function_name): scenario.run_scenario() except Exception as e: logging.exception(f"Something went wrong with the demo. Here's what: {e}") + +# snippet-end:[python.example_code.bedrock-agent.Scenario_GettingStartedBedrockAgents] diff --git a/python/example_code/bedrock-agent/test/test_bedrock_agent_wrapper.py b/python/example_code/bedrock-agent/test/test_bedrock_agent_wrapper.py index 8b3dcde5287..75c9b8e4d40 100644 --- a/python/example_code/bedrock-agent/test/test_bedrock_agent_wrapper.py +++ b/python/example_code/bedrock-agent/test/test_bedrock_agent_wrapper.py @@ -222,7 +222,7 @@ def test_get_agent(stubber, wrapper, error_code): @pytest.mark.parametrize("error_code", [None, "ClientError"]) def test_list_agents(stubber, wrapper, error_code): - expected_params = {} + expected_params = {"maxResults": 10} response = { "agentSummaries": [ { @@ -247,6 +247,54 @@ def test_list_agents(stubber, wrapper, error_code): assert exc_info.value.response["Error"]["Code"] == error_code +@pytest.mark.parametrize("error_code", [None, "ClientError"]) +def test_list_agent_action_groups(stubber, wrapper, error_code): + expected_params = { + "agentId": Fake.AGENT_ID, + "agentVersion": Fake.VERSION, + "maxResults": 10, + } + response = {"actionGroupSummaries": []} + + stubber.stub_list_agent_action_groups( + expected_params, response, error_code=error_code + ) + + if error_code is None: + got_action_groups = wrapper.list_agent_action_groups( + Fake.AGENT_ID, Fake.VERSION + ) + assert got_action_groups is not None + else: + with pytest.raises(ClientError) as exc_info: + wrapper.list_agent_action_groups(Fake.AGENT_ID, Fake.VERSION) + assert exc_info.value.response["Error"]["Code"] == error_code + + +@pytest.mark.parametrize("error_code", [None, "ClientError"]) +def test_list_agent_knowledge_bases(stubber, wrapper, error_code): + expected_params = { + "agentId": Fake.AGENT_ID, + "agentVersion": Fake.VERSION, + "maxResults": 10, + } + response = {"agentKnowledgeBaseSummaries": []} + + stubber.stub_list_agent_knowledge_bases( + expected_params, response, error_code=error_code + ) + + if error_code is None: + got_action_groups = wrapper.list_agent_knowledge_bases( + Fake.AGENT_ID, Fake.VERSION + ) + assert got_action_groups is not None + else: + with pytest.raises(ClientError) as exc_info: + wrapper.list_agent_knowledge_bases(Fake.AGENT_ID, Fake.VERSION) + assert exc_info.value.response["Error"]["Code"] == error_code + + @pytest.mark.parametrize("error_code", [None, "ClientError"]) def test_prepare_agent(stubber, wrapper, error_code): agent_id = Fake.AGENT_ID diff --git a/python/test_tools/bedrock_agent_stubber.py b/python/test_tools/bedrock_agent_stubber.py index bbd7b48e69b..78f17cba5dc 100644 --- a/python/test_tools/bedrock_agent_stubber.py +++ b/python/test_tools/bedrock_agent_stubber.py @@ -64,6 +64,21 @@ def stub_list_agents(self, expected_params, response, error_code=None): "list_agents", expected_params, response, error_code=error_code ) + def stub_list_agent_action_groups(self, expected_params, response, error_code=None): + self._stub_bifurcator( + "list_agent_action_groups", expected_params, response, error_code=error_code + ) + + def stub_list_agent_knowledge_bases( + self, expected_params, response, error_code=None + ): + self._stub_bifurcator( + "list_agent_knowledge_bases", + expected_params, + response, + error_code=error_code, + ) + def stub_prepare_agent(self, expected_params, response, error_code=None): self._stub_bifurcator( "prepare_agent", expected_params, response, error_code=error_code