logo
Menu

AI | Amazon Bedrock + Agent + KB

I found that I still needed some help after reading the link below, but I couldn't find any end-to-end code samples to test with.

Published Feb 17, 2024
Here’s my code sample with added comments explaining each part:
import logging
import boto3
from botocore.exceptions import ClientError
# Set up logging to help with debugging
logger = logging.getLogger(__name__)
class BedrockAgentRuntimeWrapper:
"""
A class to encapsulate interactions with the Amazon Bedrock Agent Runtime.
It simplifies the process of sending prompts and receiving responses.
"""

def __init__(self, agent_runtime_client):
"""
Initialize the wrapper with a Bedrock Agent Runtime client.
:param agent_runtime_client: A Boto3 client for the Bedrock Agent Runtime service.
"""
self.agent_runtime_client = agent_runtime_client
self.agent_id = '' # The unique identifier for your Bedrock agent
def invoke_agent(self, agent_alias_id, session_id, prompt):
"""
Send a prompt to the Bedrock agent and return its response.
:param agent_alias_id: The alias ID of the agent to be invoked.
:param session_id: A unique identifier for the session, used to maintain context in conversations.
:param prompt: The input text to send to the agent.
:return: The agent's response to the input text.
"""
try:
# Make the call to Bedrock to invoke the agent
response = self.agent_runtime_client.invoke_agent(
agentId=self.agent_id,
agentAliasId=agent_alias_id,
sessionId=session_id,
inputText=prompt,
)
# Process and concatenate the response from the agent
completion = ""
for event in response.get("completion"):
chunk = event["chunk"]
completion += chunk["bytes"].decode()
except ClientError as e:
# Log an error if something goes wrong
logger.error(f"Couldn't invoke agent. {e}")
raise
return completion
# Example usage of the wrapper class
if __name__ == "__main__":
# Set up the Bedrock Agent Runtime client
agent_runtime_client = boto3.client('bedrock-agent-runtime', region_name="us-east-1")
# Create an instance of the wrapper class
bedrock_wrapper = BedrockAgentRuntimeWrapper(agent_runtime_client)
# Define parameters for invoking the agent
agent_alias_id = '' # Replace with the actual agent alias ID
session_id = 'your-session-id' # Replace with a unique session ID
prompt_text = "Hello, Claude!" # The prompt to send to the agent
# Invoke the agent and get the response
response = bedrock_wrapper.invoke_agent(agent_alias_id, session_id, prompt_text)
# Print the response from the agent
print(response)
Each comment should clearly explain the purpose of the code sections, and parameters, making it easier to understand each part of the Python script.
Remember to replace the placeholders with the actual values you intend to use in your application.
  1. self.agent_id
2. Agent_alias_id
3. Session_id
4. Prompt_text