Update codebase to replace traditional agents by inline agents

[johncordeiro]

Migrating to AWS Bedrock Inline Agents

This document provides an overview of the migration from traditional AWS Bedrock Agents to the new Inline Agents approach in Nexus AI.

Overview

We've migrated our agent system from traditional AWS Bedrock Agents (which required pre-creating and deploying agents) to the more flexible Inline Agents implementation. This change brings several key benefits:

1. Faster iteration - No more waiting for agent preparation and deployment
2. Runtime configuration - Configure agent capabilities at invocation time
3. Simplified architecture - Remove complex agent lifecycle management
4. Improved flexibility - Easily switch between models or configurations
5. Multi-agent collaboration - Better support for agent teams and specialization

Implementation Details

Core Changes

• Removed traditional agent creation, preparation, and version management
• Added support for configuring and invoking agents inline at runtime
• Preserved Lambda function integration for custom actions/skills
• Maintained knowledge base and session management
• Streamlined IAM permissions and security model

Key Components

1. BedrockFileDatabase Class

The core client for interacting with AWS Bedrock services has been updated with:

• New invoke_inline_agent() method for direct agent invocation
• Streaming support via invoke_inline_agent_stream()
• Simplified Lambda function integration

2. AgentUsecase Class

The agent use case layer has been simplified:

• Removed traditional agent lifecycle management
• Added new invoke_inline_agent() method
• Added data transfer objects for inline agent configuration
• Preserved skill/Lambda management capabilities

3. Lambda Function Integration

Lambda functions are now directly attached to inline agents:

• Simplified creation and permission management
• No need for complex action group management
• Direct invocation via action groups at runtime

Usage Examples

Basic Agent Invocation

from nexus.usecases.agents.agents import AgentUsecase, InlineAgentDTO

# Configure the agent
agent_config = InlineAgentDTO(
    session_id="user-123-session-456",
    instruction="You are a helpful assistant for our company...",
    input_text="Hello, I need help with my order",
    foundation_model="anthropic.claude-3-7-sonnet-20250219-v1:0",
    enable_trace=True
)

# Invoke the agent
agent_usecase = AgentUsecase()
response = agent_usecase.invoke_inline_agent(agent_config)

Using Knowledge Base

# Configure knowledge base integration
knowledge_base_config = {
    "knowledgeBaseId": "kb-12345",
    "retrievalConfiguration": {
        "vectorSearchConfiguration": {
            "filter": {
                "equals": {
                    "key": "contentBaseUuid",
                    "value": "cb-67890"
                }
            }
        }
    }
}

# Add to agent configuration
agent_config = InlineAgentDTO(
    session_id="user-123-session-456",
    instruction="You are a helpful assistant...",
    input_text="What documents do I need for onboarding?",
    foundation_model="anthropic.claude-3-7-sonnet-20250219-v1:0",
    knowledge_bases=[knowledge_base_config]
)

Streaming Responses

for chunk in agent_usecase.invoke_inline_agent_stream(
    session_id="user-123-session-456",
    input_text="Hello",
    instruction="You are a helpful assistant...",
    content_base=content_base,
    message=message
):
    if chunk['type'] == 'chunk':
        # Process content chunk
        content = chunk['content']
        # Send to user...
    elif chunk['type'] == 'trace':
        # Process trace data for debugging
        trace = chunk['content']['trace']

Using Lambda Functions

# Define action groups with Lambda functions
action_groups = [
    {
        'name': 'OrderLookup',
        'actionGroupExecutor': {
            'lambda': 'arn:aws:lambda:region:account:function:order-lookup'
        },
        'functionSchema': {
            'functions': [
                {
                    'name': 'lookupOrder',
                    'parameters': {
                        'order_id': {'type': 'string'}
                    }
                }
            ]
        }
    }
]

# Add to agent configuration
agent_config.action_groups = action_groups

Best Practices

1. Session Management

   • Generate consistent session IDs for conversations
   • Use UUIDs or another unique identifier format
   • Consider session expiration based on use case

2. Knowledge Base Integration

   • Use specific filters to ensure relevant knowledge retrieval
   • Consider limiting result count for performance
   • Test different retrieval configurations

3. Lambda Function Design

   • Create focused, single-purpose functions
   • Use structured inputs and outputs
   • Implement proper error handling
   • Keep dependencies minimal

4. Security

   • Use the provided IAM role with minimal permissions
   • Securely manage credentials
   • Review traces and logs regularly

IAM Permissions

The inline agent approach requires the following IAM permissions:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": "bedrock:InvokeInlineAgent",
            "Resource": "*"
        },
        {
            "Effect": "Allow",
            "Action": "bedrock:InvokeModel",
            "Resource": "arn:aws:bedrock:region::foundation-model/modelId"
        },
        {
            "Effect": "Allow",
            "Action": [
                "lambda:InvokeFunction"
            ],
            "Resource": "arn:aws:lambda:region:account:function:*"
        }
    ]
}

Troubleshooting

Common issues and solutions:

1. Session State Issues

   • Ensure session state is properly formatted
   • Check knowledge base configurations
   • Verify credentials format

2. Lambda Function Errors

   • Check function permissions
   • Verify Lambda handler path
   • Review function timeout settings

3. Response Formatting Problems

   • Check model configuration
   • Verify instruction format
   • Review any prompt overrides

[gemini-code-assist]

Hello @johncordeiro, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

This pull request migrates the agent system from traditional AWS Bedrock Agents to the more flexible Inline Agents implementation. This change removes the need for pre-creating and deploying agents, allowing for faster iteration, runtime configuration, a simplified architecture, improved flexibility, and better support for multi-agent collaboration. The core changes involve removing traditional agent lifecycle management and adding support for configuring and invoking agents inline at runtime, while preserving Lambda function integration for custom actions/skills, maintaining knowledge base and session management, and streamlining IAM permissions and the security model.

Highlights

• Inline Agent Invocation: Introduces invoke_inline_agent() and invoke_inline_agent_stream() methods in the BedrockFileDatabase class for direct agent invocation and streaming support.
• Lambda Function Integration: Simplifies Lambda function integration by directly attaching them to inline agents, removing the need for complex action group management.
• Session and Knowledge Base Management: Maintains knowledge base and session management capabilities, including the ability to configure knowledge base integration and manage session state.
• IAM Permissions: IAM permissions are updated to allow bedrock:InvokeInlineAgent and lambda:InvokeFunction actions.

Changelog

Click here to see the changelog

• nexus/task_managers/file_database/bedrock.py
  • Removes the update_agent method.
  • Adds invoke_inline_agent method to invoke an inline agent with specified configurations.
  • Adds invoke_inline_agent_stream method for streaming responses from inline agents, including knowledge base integration and credentials.
  • Removes methods related to associating sub-agents, attaching knowledge bases, creating agents, and managing agent versions.
  • Updates create_lambda_function to create a default alias for the function and simplifies the permission management.
  • Removes methods related to disassociating sub-agents, converting agents to supervisors, deleting agents, and invoking supervisors with streaming.
  • Removes prepare_agent method.
  • Updates _allow_agent_lambda to allow Bedrock to invoke Lambda functions with inline agents.
  • Removes wait_agent_status_update method.
  • Removes update_agent_action_group and _create_lambda_iam_role methods.
• nexus/task_managers/tasks_bedrock.py
  • Removes agent_external_id and agent_version parameters from run_create_lambda_function and run_update_lambda_function tasks.
  • Updates docstrings for run_create_lambda_function and run_update_lambda_function to reflect changes in parameters.
• nexus/usecases/agents/agents.py
  • Removes BedrockSubAgent import.
  • Replaces UpdateAgentDTO with InlineAgentDTO for inline agent configuration.
  • Removes methods related to assigning agents, creating agents, creating external agents, creating external agent aliases, updating agents to supervisors, creating supervisors, creating agent versions, and deleting agent versions.
  • Adds invoke_inline_agent method to invoke an inline agent with the provided configuration.
  • Adds invoke_inline_agent_stream method to invoke an inline agent with streaming response.
  • Removes agent_external_id and agent_version parameters from create_skill and update_skill methods.
  • Removes methods related to invoking supervisors with streaming, preparing agents, unassigning agents, updating agents, and creating external supervisors.
  • Removes methods related to updating supervisor collaborators and updating multi-agent configurations.
  • Adds validate_inline_agent_dto method to validate the InlineAgentDTO fields.
  • Adds create_dict_to_inline_dto method to convert a dictionary to an InlineAgentDTO.
  • Removes methods related to updating dictionaries to DTOs, creating dictionaries to DTOs, handling agent DTOs, getting update fields, processing skill updates, and waiting for agent status updates.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

---

From agents old to inline new,
Bedrock's power shines anew.
No more waiting, swift and free,
Agents run inline, for you and me.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The code changes introduce inline agents, which is a significant improvement over traditional agents. The implementation is generally well-structured, but there are some areas that could be improved in terms of error handling, security, and configurability.

Summary of Findings

• Missing Error Handling: The code lacks proper error handling for Lambda function creation and updates. This could lead to unhandled exceptions and potential instability. It's important to implement try-except blocks to catch potential errors and handle them gracefully.
• Insecure Lambda Permissions: The Lambda permission policy uses a wildcard resource ARN, which is overly permissive and could pose a security risk. It's recommended to restrict the resource ARN to the specific Lambda functions that the Bedrock agent needs to invoke.
• Hardcoded Values: Several hardcoded values, such as the agent foundation model and session TTL, should be configurable via settings or environment variables to allow for easier customization and deployment.

Merge Readiness

The code changes introduce inline agents, which is a significant improvement over traditional agents. However, there are some areas that need to be addressed before merging, particularly error handling and security. I would recommend addressing these before merging. I am unable to directly approve the pull request, and that users should have others review and approve this code before merging.

[gemini-code-assist]

This Lambda permission policy uses a wildcard resource ARN, which is overly permissive and could pose a security risk. It's recommended to restrict the resource ARN to the specific Lambda functions that the Bedrock agent needs to invoke.

            FunctionName=lambda_function_name,
            StatementId=f"allow_bedrock_inline_{uuid.uuid4()}",
            Action="lambda:InvokeFunction",
            Principal="bedrock.amazonaws.com",
            # Restrict the resource ARN to the specific Lambda function
            SourceArn=f"arn:aws:lambda:{self.region_name}:{self.account_id}:function:{lambda_function_name}"

[gemini-code-assist]

Consider adding more specific error handling here. Catching all exceptions might mask important issues. At a minimum, log the exception details.

Suggested change

	except Exception as e:
	print(f"Error fetching credentials: {str(e)}")
	except Exception as e:
	print(f"Error fetching credentials: {str(e)}")
	# Consider logging the error with more details
	# logging.exception("Error fetching credentials")