agent-squad 1.0.0__py3-none-any.whl

agent_squad/__init__.py

@@ -0,0 +1,3 @@
+from .shared import user_agent
+
+user_agent.inject_user_agent()

agent_squad/agents/__init__.py

@@ -0,0 +1,81 @@
+"""
+Code for Agents.
+"""
+from .agent import Agent, AgentOptions, AgentCallbacks, AgentProcessingResult, AgentResponse, AgentStreamResponse
+
+
+try:
+    from .lambda_agent import LambdaAgent, LambdaAgentOptions
+    from .bedrock_llm_agent import BedrockLLMAgent, BedrockLLMAgentOptions
+    from .lex_bot_agent import LexBotAgent, LexBotAgentOptions
+    from .amazon_bedrock_agent import AmazonBedrockAgent, AmazonBedrockAgentOptions
+    from .comprehend_filter_agent import ComprehendFilterAgent, ComprehendFilterAgentOptions
+    from .bedrock_translator_agent import BedrockTranslatorAgent, BedrockTranslatorAgentOptions
+    from .chain_agent import ChainAgent, ChainAgentOptions
+    from .bedrock_inline_agent import BedrockInlineAgent, BedrockInlineAgentOptions
+    from .bedrock_flows_agent import BedrockFlowsAgent, BedrockFlowsAgentOptions
+    _AWS_AVAILABLE = True
+except ImportError:
+    _AWS_AVAILABLE = False
+try:
+    from .anthropic_agent import AnthropicAgent, AnthropicAgentOptions
+    _ANTHROPIC_AVAILABLE = True
+except ImportError:
+    _ANTHROPIC_AVAILABLE = False
+
+
+try:
+    from .openai_agent import OpenAIAgent, OpenAIAgentOptions
+    _OPENAI_AVAILABLE = True
+except ImportError:
+    _OPENAI_AVAILABLE = False
+
+from .supervisor_agent import SupervisorAgent, SupervisorAgentOptions
+
+__all__ = [
+    'Agent',
+    'AgentOptions',
+    'AgentCallbacks',
+    'AgentProcessingResult',
+    'AgentResponse',
+    'AgentStreamResponse',
+    'SupervisorAgent',
+    'SupervisorAgentOptions'
+]
+
+
+if _AWS_AVAILABLE :
+    __all__.extend([
+        'LambdaAgent',
+        'LambdaAgentOptions',
+        'BedrockLLMAgent',
+        'BedrockLLMAgentOptions',
+        'LexBotAgent',
+        'LexBotAgentOptions',
+        'AmazonBedrockAgent',
+        'AmazonBedrockAgentOptions',
+        'ComprehendFilterAgent',
+        'ComprehendFilterAgentOptions',
+        'ChainAgent',
+        'ChainAgentOptions',
+        'BedrockTranslatorAgent',
+        'BedrockTranslatorAgentOptions',
+        'BedrockInlineAgent',
+        'BedrockInlineAgentOptions',
+        'BedrockFlowsAgent',
+        'BedrockFlowsAgentOptions'
+    ])
+
+
+if _ANTHROPIC_AVAILABLE:
+    __all__.extend([
+        'AnthropicAgent',
+        'AnthropicAgentOptions'
+    ])
+
+
+if _OPENAI_AVAILABLE:
+    __all__.extend([
+            'OpenAIAgent',
+            'OpenAIAgentOptions'
+        ])

agent_squad/agents/agent.py

@@ -0,0 +1,181 @@
+from typing import Union, AsyncIterable, Optional, Any, TypeAlias
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from agent_squad.types import ConversationMessage
+from agent_squad.utils import Logger
+
+# Type aliases for complex types
+AgentParamsType: TypeAlias = dict[str, Any]
+AgentOutputType: TypeAlias = Union[str, "AgentStreamResponse", Any]  # Forward reference
+
+@dataclass
+class AgentProcessingResult:
+    """
+    Contains metadata about the result of an agent's processing.
+
+    Attributes:
+        user_input: The original input from the user
+        agent_id: Unique identifier for the agent
+        agent_name: Display name of the agent
+        user_id: Identifier for the user
+        session_id: Identifier for the current session
+        additional_params: Optional additional parameters for the agent
+    """
+    user_input: str
+    agent_id: str
+    agent_name: str
+    user_id: str
+    session_id: str
+    additional_params: AgentParamsType = field(default_factory=dict)
+
+
+@dataclass
+class AgentStreamResponse:
+    """
+    Represents a streaming response from an agent.
+
+    Attributes:
+        text: The current text in the stream
+        final_message: The complete message when streaming is complete
+    """
+    text: str = ""
+    final_message: Optional[ConversationMessage] = None
+
+
+@dataclass
+class AgentResponse:
+    """
+    Complete response from an agent, including metadata and output.
+
+    Attributes:
+        metadata: Processing metadata
+        output: The actual output from the agent
+        streaming: Whether this response is streaming
+    """
+    metadata: AgentProcessingResult
+    output: AgentOutputType
+    streaming: bool
+
+
+class AgentCallbacks:
+    """
+    Defines callbacks that can be triggered during agent processing.
+    Provides default implementations that can be overridden by subclasses.
+    """
+    def on_llm_new_token(self, token: str) -> None:
+        """
+        Called when a new token is generated by the LLM.
+
+        Args:
+            token: The new token generated
+        """
+        pass  # Default implementation does nothing
+
+
+@dataclass
+class AgentOptions:
+    """
+    Configuration options for an agent.
+
+    Attributes:
+        name: The display name of the agent
+        description: A description of the agent's purpose and capabilities
+        save_chat: Whether to save the chat history
+        callbacks: Optional callbacks for agent events
+        LOG_AGENT_DEBUG_TRACE: Whether to enable debug tracing for this agent
+    """
+    name: str
+    description: str
+    save_chat: bool = True
+    callbacks: Optional[AgentCallbacks] = None
+    # Optional: Flag to enable/disable agent debug trace logging
+    # If true, the agent will log additional debug information
+    LOG_AGENT_DEBUG_TRACE: Optional[bool] = False
+
+class Agent(ABC):
+    """
+    Abstract base class for all agents in the system.
+
+    Implements common functionality and defines the required interface
+    for concrete agent implementations.
+    """
+
+    def __init__(self, options: AgentOptions):
+        """
+        Initialize a new agent with the given options.
+
+        Args:
+            options: Configuration options for this agent
+        """
+        self.name = options.name
+        self.id = self.generate_key_from_name(options.name)
+        self.description = options.description
+        self.save_chat = options.save_chat
+        self.callbacks = options.callbacks if options.callbacks is not None else AgentCallbacks()
+        self.log_debug_trace = options.LOG_AGENT_DEBUG_TRACE
+
+    def is_streaming_enabled(self) -> bool:
+        """
+        Whether this agent supports streaming responses.
+
+        Returns:
+            True if streaming is enabled, False otherwise
+        """
+        return False
+
+    @staticmethod
+    def generate_key_from_name(name: str) -> str:
+        """
+        Generate a standardized key from an agent name.
+
+        Args:
+            name: The display name to convert
+
+        Returns:
+            A lowercase, hyphenated key with special characters removed
+        """
+        # Remove special characters and replace spaces with hyphens
+        key = re.sub(r"[^a-zA-Z0-9\s-]", "", name)
+        key = re.sub(r"\s+", "-", key)
+        return key.lower()
+
+    @abstractmethod
+    async def process_request(
+        self,
+        input_text: str,
+        user_id: str,
+        session_id: str,
+        chat_history: list[ConversationMessage],
+        additional_params: Optional[AgentParamsType] = None,
+    ) -> Union[ConversationMessage, AsyncIterable[AgentOutputType]]:
+        """
+        Process a user request and generate a response.
+
+        Args:
+            input_text: The user's input text
+            user_id: Identifier for the user
+            session_id: Identifier for the current session
+            chat_history: List of previous messages in the conversation
+            additional_params: Optional additional parameters
+
+        Returns:
+            Either a complete message or an async iterable for streaming responses
+        """
+        pass
+
+    def log_debug(self, class_name: str, message: str, data: Any = None) -> None:
+        """
+        Log a debug message if debug tracing is enabled.
+
+        Args:
+            class_name: Name of the class logging the message
+            message: The message to log
+            data: Optional data to include in the log
+        """
+        if self.log_debug_trace:
+            prefix = f"> {class_name} \n> {self.name} \n>"
+            if data:
+                Logger.info(f"{prefix} {message} \n> {data}")
+            else:
+                Logger.info(f"{prefix} {message} \n>")

agent_squad/agents/amazon_bedrock_agent.py

@@ -0,0 +1,172 @@
+"""
+Amazon Bedrock Agent Integration Module
+
+This module provides a robust implementation for interacting with Amazon Bedrock agents,
+offering a flexible and extensible way to process conversational interactions using
+AWS Bedrock's agent runtime capabilities.
+"""
+
+from typing import Any, Optional
+from dataclasses import dataclass
+import os
+import boto3
+from botocore.exceptions import BotoCoreError, ClientError
+from agent_squad.agents import Agent, AgentOptions, AgentStreamResponse
+from agent_squad.types import ConversationMessage, ParticipantRole
+from agent_squad.utils import Logger
+from agent_squad.shared import user_agent
+
+
+@dataclass
+class AmazonBedrockAgentOptions(AgentOptions):
+    """
+    Configuration options for Amazon Bedrock Agent initialization.
+
+    Provides flexible configuration for Bedrock agent runtime:
+    - agent_id: Unique identifier for the Bedrock agent
+    - agent_alias_id: Specific alias for the agent
+    - client: Optional custom boto3 client (allows dependency injection)
+    - streaming: Flag to enable streaming response mode (on final response)
+    - enableTrace: Flag to enable detailed event tracing
+    """
+    region: Optional[str] = None
+    agent_id: str = None
+    agent_alias_id: str = None
+    client: Any | None = None
+    streaming: bool | None = False
+    enableTrace: bool | None = False
+
+
+class AmazonBedrockAgent(Agent):
+    """
+    Specialized agent for interacting with Amazon Bedrock's intelligent agent runtime.
+
+    This class extends the base Agent class to provide:
+    - Direct integration with AWS Bedrock agent runtime
+    - Configurable response handling (streaming/non-streaming)
+    - Comprehensive error management
+    - Flexible session and conversation state management
+    """
+
+    def __init__(self, options: AmazonBedrockAgentOptions):
+        """
+        Initialize the Bedrock agent with comprehensive configuration.
+
+        Handles client creation, either using a provided client or
+        automatically creating one based on AWS configuration.
+
+        :param options: Detailed configuration for agent initialization
+        """
+        super().__init__(options)
+
+        # Store core agent identifiers
+        self.agent_id = options.agent_id
+        self.agent_alias_id = options.agent_alias_id
+
+        # Set up Bedrock runtime client
+        if options.client:
+            # Use provided client (useful for testing or custom configurations)
+            self.client = options.client
+        else:
+            # Create default client using AWS region from options or environment
+            self.client = boto3.client('bedrock-agent-runtime',
+                                       region_name=options.region or os.environ.get('AWS_REGION'))
+
+        user_agent.register_feature_to_client(self.client, feature="bedrock-agent")
+
+
+        # Configure response handling modes
+        self.streaming = options.streaming
+        self.enableTrace = options.enableTrace
+
+    def is_streaming_enabled(self) -> bool:
+        """
+        Check if streaming mode is active for response processing.
+
+        :return: Boolean indicating streaming status
+        """
+        return self.streaming is True
+
+    async def process_request(
+        self,
+        input_text: str,
+        user_id: str,
+        session_id: str,
+        chat_history: list[ConversationMessage],
+        additional_params: dict[str, str] | None = None
+    ) -> ConversationMessage:
+        """
+        Process a user request through the Bedrock agent runtime.
+
+        Handles the entire interaction lifecycle:
+        - Manages session state
+        - Invokes agent with configured parameters
+        - Processes streaming or non-streaming responses
+        - Handles potential errors
+
+        :param input_text: User's input message
+        :param user_id: Identifier for the user
+        :param session_id: Unique conversation session identifier
+        :param chat_history: Previous conversation messages
+        :param additional_params: Optional supplementary parameters
+        :return: Agent's response as a conversation message
+        """
+        # Initialize session state, defaulting to empty if not provided
+        session_state = {}
+        if (additional_params and 'sessionState' in additional_params):
+            session_state = additional_params['sessionState']
+
+        try:
+            # Configure streaming behavior
+            streamingConfigurations = {
+                'streamFinalResponse': self.streaming
+            }
+
+            # Invoke Bedrock agent with comprehensive configuration
+            response = self.client.invoke_agent(
+                agentId=self.agent_id,
+                agentAliasId=self.agent_alias_id,
+                sessionId=session_id,
+                inputText=input_text,
+                enableTrace=self.enableTrace,
+                sessionState=session_state,
+                streamingConfigurations=streamingConfigurations if self.streaming else {}
+            )
+
+            completion = ""
+
+            if self.streaming:
+                async def generate_chunks():
+                    nonlocal completion
+                    for event in response['completion']:
+                        if 'chunk' in event:
+                            chunk = event['chunk']
+                            decoded_response = chunk['bytes'].decode('utf-8')
+                            self.callbacks.on_llm_new_token(decoded_response)
+                            completion += decoded_response
+                            yield AgentStreamResponse(text=decoded_response)
+                        elif 'trace' in event and self.enableTrace:
+                            Logger.info(f"Received event: {event}")
+                    yield AgentStreamResponse(
+                        final_message=ConversationMessage(
+                            role=ParticipantRole.ASSISTANT.value,
+                            content=[{'text':completion}]))
+                return generate_chunks()
+            else:
+                for event in response['completion']:
+                    if 'chunk' in event:
+                        chunk = event['chunk']
+                        decoded_response = chunk['bytes'].decode('utf-8')
+                        self.callbacks.on_llm_new_token(decoded_response)
+                        completion += decoded_response
+                    elif 'trace' in event and self.enableTrace:
+                        Logger.info(f"Received event: {event}")
+
+                return ConversationMessage(
+                    role=ParticipantRole.ASSISTANT.value,
+                    content=[{"text": completion}]
+                )
+        except (BotoCoreError, ClientError) as error:
+            # Comprehensive error logging and propagation
+            Logger.error(f"Error processing request: {str(error)}")
+            raise error