Description

This PR enhances the memory tool with comprehensive boto3 session injection capabilities, enabling flexible AWS authentication patterns for accessing Amazon Bedrock Knowledge Bases. The changes add two new public functions (set_memory_session and get_memory_session) and extend the existing memory tool to support custom session management.

Key Changes:

• Added set_memory_session() function to store boto3 sessions with optional keys for multi-session management
• Added get_memory_session() function to retrieve stored sessions by key
• Enhanced memory tool with session_key parameter for selecting which stored session to use
• Implemented global session store (_SESSION_STORE) for managing multiple AWS sessions simultaneously
• Updated MemoryServiceClient to accept and use custom boto3 sessions
• Added comprehensive session validation and error handling throughout the tool
• Maintained full backward compatibility with existing memory tool usage patterns

New Capabilities:

• Multi-environment authentication: Support for dev/staging/prod environments with separate credentials
• Cross-account AWS access: Enable access to Knowledge Bases in different AWS accounts via STS assume role
• Regional session management: Support for region-specific sessions for data compliance requirements
• Session lifecycle management: Proper session storage, retrieval, and cleanup functionality
• Custom AWS profile support: Easy integration with existing AWS credential configurations
• Federated login sessions: Improved support for enterprise authentication patterns

Related Issues

#380

Documentation and Tools PR

• feat! REQUIRED: Docs update for: Added option to pass existing boto3 … docs#143
• [feature] memory tool: use existing boto3 session tools#138

Type of Change

New feature

Testing

How have you tested the change?

Code Quality and Validation:

• Created comprehensive validation scripts to test all code examples (27 examples validated)
• Verified function signatures match between implementation and documentation
• Tested session storage and retrieval functionality with various session keys
• Validated backward compatibility - existing memory tool usage continues to work unchanged
• Tested error handling for invalid sessions, missing keys, and authentication failures

Integration Testing:

• Verified session injection works with different AWS authentication patterns
• Tested multi-session management scenarios (production, staging, development)
• Validated cross-account access patterns using STS assume role
• Confirmed regional session management functionality
• Tested session cleanup and lifecycle management

Repository Compatibility:

[x] I ran hatch run prepare - All formatting and linting checks pass
[x] Verified changes don't break functionality in consuming repositories
[x] Confirmed no new warnings introduced
[x] Validated all existing tests continue to pass

Checklist

[x] I have read the CONTRIBUTING document
[x] I have added any necessary tests that prove my fix is effective or my feature works
[x] I have updated the documentation accordingly
[x] I have added an appropriate example to the documentation to outline the feature, or no new docs are needed
[x] My changes generate no new warnings
[x] Any dependent changes have been merged and published

Additional Quality Assurance Completed:

[x] All 27 code examples in documentation are syntactically correct and functional
[x] Function signatures validated against actual implementation
[x] Session management follows AWS security best practices
[x] Comprehensive error handling provides clear user feedback
[x] Full backward compatibility maintained for existing users