Base Agent

The BedrockAgent class serves as the foundation for all agents in the Bedrock Swarm framework. It provides core functionality for model interaction, tool management, and memory handling.

Class Documentation

Base class for Bedrock-powered agents.

Each agent has: 1. A specific role/domain of expertise 2. A set of tools it can use 3. A memory system for context 4. A Bedrock model for processing

PARAMETER	DESCRIPTION
model_id	Bedrock model ID TYPE: str
name	Name of the agent TYPE: str
role	Description of agent's role/expertise TYPE: str
tools	Optional list of tools available to the agent TYPE: Optional[List[BaseTool]] DEFAULT: None
memory	Optional memory system (defaults to SimpleMemory) TYPE: Optional[BaseMemory] DEFAULT: None
system_prompt	Optional system prompt TYPE: Optional[str] DEFAULT: None

RAISES	DESCRIPTION
InvalidModelError	If model ID is not supported

Source code in src/bedrock_swarm/agents/base.py

def __init__(
    self,
    model_id: str,
    name: str,
    role: str,
    tools: Optional[List[BaseTool]] = None,
    memory: Optional[BaseMemory] = None,
    system_prompt: Optional[str] = None,
) -> None:
    """Initialize the agent.

    Args:
        model_id: Bedrock model ID
        name: Name of the agent
        role: Description of agent's role/expertise
        tools: Optional list of tools available to the agent
        memory: Optional memory system (defaults to SimpleMemory)
        system_prompt: Optional system prompt

    Raises:
        InvalidModelError: If model ID is not supported
    """
    self.model_id = model_id
    self.name = name
    self.role = role
    self.tools = {tool.name: tool for tool in tools} if tools else {}
    self.memory = memory or SimpleMemory()
    self.system_prompt = system_prompt

    logger.debug(f"Initializing agent {name} with role: {role}")
    logger.debug(f"Available tools: {list(self.tools.keys())}")

    # Initialize AWS session
    self.session = boto3.Session(
        region_name=AWSConfig.region,
        profile_name=AWSConfig.profile,
    )

    # Initialize model
    self.model = self._initialize_model()

Attributes

last_token_count: int property

Get the token count from the last request.

RETURNS	DESCRIPTION
int	Number of tokens used in last request TYPE: int

Functions

_initialize_model() -> BedrockModel

Initialize the Bedrock model.

Source code in src/bedrock_swarm/agents/base.py

def _initialize_model(self) -> BedrockModel:
    """Initialize the Bedrock model."""
    try:
        return ModelFactory.create_model(self.model_id)
    except ValueError as e:
        raise InvalidModelError(str(e))

_build_prompt(message: str) -> str

Build the complete prompt including context and conversation history.

This method builds a comprehensive prompt that includes: 1. System prompt and role context 2. Available tools and their schemas 3. Recent conversation history 4. Response format instructions 5. Current message

Source code in src/bedrock_swarm/agents/base.py

def _build_prompt(self, message: str) -> str:
    """Build the complete prompt including context and conversation history.

    This method builds a comprehensive prompt that includes:
    1. System prompt and role context
    2. Available tools and their schemas
    3. Recent conversation history
    4. Response format instructions
    5. Current message
    """
    prompt = []

    # Add system prompt if provided
    if self.system_prompt:
        prompt.append(f"System: {self.system_prompt}")

    # Add role context
    prompt.append(f"You are a specialized agent with expertise in: {self.role}")

    # Add available tools
    if self.tools:
        prompt.append("\n<tools>")
        for tool in self.tools.values():
            prompt.append(f"- {tool.name}: {tool.description}")
            schema = tool.get_schema()
            prompt.append(f"  Schema: {json.dumps(schema, indent=2)}")
        prompt.append("</tools>")

    # Add conversation history from memory
    recent_messages = self.memory.get_messages()[-5:]  # Get last 5 messages
    if recent_messages:
        prompt.append("\n<conversation_history>")
        for msg in recent_messages:
            # Include metadata about tool usage if available
            tool_info = ""
            if msg.metadata and msg.metadata.get("type") == "tool_result":
                tool_info = (
                    f" [Tool Result: {msg.metadata.get('tool_call_id', 'unknown')}]"
                )
            prompt.append(f"{msg.role}{tool_info}: {msg.content}")
        prompt.append("</conversation_history>")

    # Add response format instructions using XML
    prompt.extend(
        [
            "\n<response_format>",
            "You must respond in one of two formats:",
            "\n1. To use a tool:",
            "Respond with ONLY a JSON object in this exact format (no explanation text outside the JSON):",
            '{"type": "tool_call", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "tool_name", "arguments": {"arg1": "value1"}}}]}',
            "\n2. For normal responses:",
            "Respond with ONLY a JSON object in this exact format (no explanation text outside the JSON):",
            '{"type": "message", "content": "your natural language response here"}',
            "\n<rules>",
            "- Always use proper JSON with double quotes",
            "- Never include explanations or text outside the JSON object",
            "- Tool arguments must be a valid JSON object, not a string",
            "- Respond with exactly one complete JSON object",
            "- Maintain conversation context from history",
            "- Reference previous tool results when relevant",
            "</rules>",
            "</response_format>",
            f"\n<input>{message}</input>",
        ]
    )

    final_prompt = "\n".join(prompt)
    logger.debug(f"Built prompt for agent {self.name}:\n{final_prompt}")
    return final_prompt

generate(message: str) -> AgentResponse

Generate a response to a message.

This method: 1. Records the incoming message in memory 2. Generates a response using the model 3. Records the response and any tool calls in memory 4. Returns the processed response

PARAMETER	DESCRIPTION
message	Message to respond to TYPE: str

RETURNS	DESCRIPTION
AgentResponse	Response containing either tool calls or direct message

Source code in src/bedrock_swarm/agents/base.py

def generate(self, message: str) -> AgentResponse:
    """Generate a response to a message.

    This method:
    1. Records the incoming message in memory
    2. Generates a response using the model
    3. Records the response and any tool calls in memory
    4. Returns the processed response

    Args:
        message: Message to respond to

    Returns:
        Response containing either tool calls or direct message
    """
    logger.debug(f"Agent {self.name} generating response for message: {message}")

    # Record incoming message in memory
    self.memory.add_message(
        Message(
            role="user",
            content=message,
            timestamp=datetime.now(),
            metadata={"type": "user_message", "agent": self.name},
        )
    )

    # Get bedrock client
    client = self.session.client(
        "bedrock-runtime",
        endpoint_url=AWSConfig.endpoint_url,
    )

    # Build prompt and get response
    prompt = self._build_prompt(message)
    response = self.model.invoke(client=client, message=prompt)
    logger.debug(f"Raw model response: {response}")

    # Record the response in memory with appropriate metadata
    if response.get("type") == "tool_call":
        # Record tool call intent
        self.memory.add_message(
            Message(
                role="assistant",
                content=json.dumps(response["tool_calls"]),
                timestamp=datetime.now(),
                metadata={
                    "type": "tool_call_intent",
                    "agent": self.name,
                    "tool_calls": response["tool_calls"],
                },
            )
        )
    else:
        # Record normal message response
        self.memory.add_message(
            Message(
                role="assistant",
                content=response.get("content", ""),
                timestamp=datetime.now(),
                metadata={"type": "assistant_response", "agent": self.name},
            )
        )

    # Return the processed response
    return response

_format_prompt(message: str, history: List[Message]) -> str

Format the prompt with message history.

PARAMETER	DESCRIPTION
message	Current message TYPE: str
history	Message history TYPE: List[Message]

RETURNS	DESCRIPTION
str	Formatted prompt string

Source code in src/bedrock_swarm/agents/base.py

def _format_prompt(self, message: str, history: List[Message]) -> str:
    """Format the prompt with message history.

    Args:
        message: Current message
        history: Message history

    Returns:
        Formatted prompt string
    """
    # Start with system prompt if provided
    prompt = f"{self.system_prompt}\n\n" if self.system_prompt else ""

    # Add message history
    for msg in history:
        prompt += f"{msg.role}: {msg.content}\n"

    # Add current message
    prompt += f"human: {message}\nassistant:"

    return prompt

Core Features

The base agent provides:

1. Model Integration
2. Model initialization
3. Request formatting
4. Response processing

5. Token management

6. Tool Management

7. Tool registration
8. Tool validation
9. Tool execution

10. Error handling

11. Memory System

12. Context management
13. Message history
14. State persistence
15. Memory cleanup

Usage Examples

from bedrock_swarm.agents import BedrockAgent
from bedrock_swarm.tools import CalculatorTool

# Create agent
agent = BedrockAgent(
    name="calculator",
    model_id="us.anthropic.claude-3-5-sonnet-20241022-v2:0",
    tools=[CalculatorTool()],
    system_prompt="You are a calculation specialist."
)

# Process message
response = agent.generate("Calculate 15 * 7")
print(response)

# Check memory
history = agent.memory.get_messages()
print(history)

# Get token usage
print(agent.last_token_count)

Tool Integration

Agents can use multiple tools:

from bedrock_swarm.tools import CalculatorTool, TimeTool

# Create agent with multiple tools
agent = BedrockAgent(
    name="utility_agent",
    model_id="us.anthropic.claude-3-5-sonnet-20241022-v2:0",
    tools=[
        CalculatorTool(),
        TimeTool()
    ]
)

# Tools are available in prompts
response = agent.generate(
    "What will be 2 * 5 in 3 hours?"
)

Memory Management

The agent includes memory management:

# Access memory system
memory = agent.memory

# Store message
memory.add_message(
    role="user",
    content="Hello!",
    metadata={"timestamp": "2024-02-29"}
)

# Get recent messages
recent = memory.get_recent(n=5)

# Clear memory
memory.clear()

Error Handling

The agent handles various errors:

1. Model Errors

   try:
       response = agent.generate("Invalid request")
   except ModelInvokeError as e:
       print(f"Model error: {e}")
   

2. Tool Errors

   try:
       response = agent.execute_tool("invalid_tool")
   except ToolError as e:
       print(f"Tool error: {e}")
   

3. Memory Errors

   try:
       messages = agent.memory.get_messages()
   except MemoryError as e:
       print(f"Memory error: {e}")
   

Implementation Details

The agent implementation includes:

1. Request Processing
2. Message formatting
3. Tool integration
4. Memory context

5. Response handling

6. Tool Management

7. Tool registration
8. Schema validation
9. Execution handling

10. Result processing

11. Memory Integration

12. Message storage
13. Context retrieval
14. State management
15. Cleanup handling

Agent Configuration

Agents can be configured with:

agent_config = {
    "name": "specialized_agent",
    "model_id": "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
    "system_prompt": "You are a specialized agent...",
    "tools": [CalculatorTool()],
    "memory": CustomMemory(),
    "model_config": {
        "temperature": 0.7,
        "max_tokens": 1000
    }
}

agent = BedrockAgent(**agent_config)

Best Practices

1. Tool Organization

   # Group related tools
   math_tools = [
       CalculatorTool(),
       StatisticsTool()
   ]
   
   # Create specialized agent
   math_agent = BedrockAgent(
       name="math_expert",
       tools=math_tools
   )
   

2. Memory Usage

   # Use memory for context
   agent.memory.add_context(
       key="user_preferences",
       value={"format": "metric"}
   )
   
   # Clear old context
   agent.memory.clear_context("old_key")
   

3. Error Recovery

   try:
       response = agent.generate(message)
   except Exception as e:
       # Log error
       logger.error(f"Agent error: {e}")
       # Attempt recovery
       agent.reset_state()
       # Retry with fallback
       response = agent.generate_fallback(message)