Merge pull request #277 from awslabs/custom-user-agent

Custom user agent

Author: brnaba-aws

python/src/multi_agent_orchestrator/__init__.py

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

python/src/multi_agent_orchestrator/agents/amazon_bedrock_agent.py

@@ -14,6 +14,7 @@
 from multi_agent_orchestrator.agents import Agent, AgentOptions, AgentStreamResponse
 from multi_agent_orchestrator.types import ConversationMessage, ParticipantRole
 from multi_agent_orchestrator.utils import Logger
+from multi_agent_orchestrator.shared import user_agent
 
 
 @dataclass
@@ -69,7 +70,10 @@ def __init__(self, options: AmazonBedrockAgentOptions):
         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'))
+                                       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

python/src/multi_agent_orchestrator/agents/bedrock_flows_agent.py

@@ -1,17 +1,18 @@
 from typing import List, Dict, Any, Optional, Callable
 from dataclasses import dataclass
 import os
-import json
 import boto3
 from multi_agent_orchestrator.utils import (Logger, conversation_to_dict)
 from multi_agent_orchestrator.agents import (Agent, AgentOptions)
 from multi_agent_orchestrator.types import (ConversationMessage, ParticipantRole)
+from multi_agent_orchestrator.shared import user_agent
 
 # BedrockFlowsAgentOptions Dataclass
 @dataclass
 class BedrockFlowsAgentOptions(AgentOptions):
     flowIdentifier: str = None
     flowAliasIdentifier: str = None
+    region: Optional[str] = None
     bedrock_agent_client: Optional[Any] = None
     enableTrace: Optional[bool] = False
     flow_input_encoder: Optional[Callable] = None
@@ -28,13 +29,10 @@ def __init__(self, options: BedrockFlowsAgentOptions):
         if options.bedrock_agent_client:
             self.bedrock_agent_client = options.bedrock_agent_client
         else:
-            if options.region:
-                self.bedrock_agent_client = boto3.client(
-                    'bedrock-agent-runtime',
-                    region_name=options.region or os.environ.get('AWS_REGION')
-                )
-            else:
-                self.bedrock_agent_client = boto3.client('bedrock-agent-runtime')
+            self.client = boto3.client('bedrock-agent-runtime',
+                                       region_name=options.region or os.environ.get('AWS_REGION'))
+
+        user_agent.register_feature_to_client(self.bedrock_agent_client, feature="bedrock-flows-agent")
 
         self.enableTrace = options.enableTrace
         self.flowAliasIdentifier = options.flowAliasIdentifier

python/src/multi_agent_orchestrator/agents/bedrock_llm_agent.py

@@ -12,7 +12,7 @@
                        AgentProviderType)
 from multi_agent_orchestrator.utils import conversation_to_dict, Logger, AgentTools, AgentTool
 from multi_agent_orchestrator.retrievers import Retriever
-
+from multi_agent_orchestrator.shared import user_agent
 
 @dataclass
 class BedrockLLMAgentOptions(AgentOptions):
@@ -36,11 +36,13 @@ def __init__(self, options: BedrockLLMAgentOptions):
             if options.region:
                 self.client = boto3.client(
                     'bedrock-runtime',
-                    region_name=options.region or os.environ.get('AWS_REGION')
-                )
+                    region_name=options.region)
             else:
                 self.client = boto3.client('bedrock-runtime')
 
+        user_agent.register_feature_to_client(self.client, feature="bedrock-llm-agent")
+
+
         self.model_id: str = options.model_id or BEDROCK_MODEL_ID_CLAUDE_3_HAIKU
         self.streaming: bool = options.streaming
         self.inference_config: dict[str, Any]

python/src/multi_agent_orchestrator/agents/lambda_agent.py

@@ -5,6 +5,7 @@
 from multi_agent_orchestrator.agents import Agent, AgentOptions
 from multi_agent_orchestrator.types import ConversationMessage, ParticipantRole
 from multi_agent_orchestrator.utils import conversation_to_dict
+from multi_agent_orchestrator.shared import user_agent
 
 @dataclass
 class LambdaAgentOptions(AgentOptions):
@@ -25,7 +26,11 @@ class LambdaAgent(Agent):
     def __init__(self, options: LambdaAgentOptions):
         super().__init__(options)
         self.options = options
+
         self.lambda_client = boto3.client('lambda', region_name=self.options.function_region)
+
+        user_agent.register_feature_to_client(self.lambda_client, feature="lambda-agent")
+
         if self.options.input_payload_encoder is None:
             self.encoder = self.__default_input_payload_encoder
         else:

python/src/multi_agent_orchestrator/agents/lex_bot_agent.py

@@ -1,12 +1,12 @@
-from typing import List, Dict, Optional
+import os
+from typing import Any, Optional
 from dataclasses import dataclass
 import boto3
 from botocore.exceptions import BotoCoreError, ClientError
 from multi_agent_orchestrator.agents import Agent, AgentOptions
 from multi_agent_orchestrator.types import ConversationMessage, ParticipantRole
 from multi_agent_orchestrator.utils import Logger
-import os
-from typing import Any
+from multi_agent_orchestrator.shared import user_agent
 
 @dataclass
 class LexBotAgentOptions(AgentOptions):
@@ -30,6 +30,9 @@ def __init__(self, options: LexBotAgentOptions):
         else:
             self.lex_client = boto3.client('lexv2-runtime', region_name=self.region)
 
+        user_agent.register_feature_to_client(self.lex_client, feature="lex-agent")
+
+
         self.bot_id = options.bot_id
         self.bot_alias_id = options.bot_alias_id
         self.locale_id = options.locale_id
@@ -38,8 +41,8 @@ def __init__(self, options: LexBotAgentOptions):
             raise ValueError("bot_id, bot_alias_id, and locale_id are required for LexBotAgent")
 
     async def process_request(self, input_text: str, user_id: str, session_id: str,
-                        chat_history: List[ConversationMessage],
-                        additional_params: Optional[Dict[str, str]] = None) -> ConversationMessage:
+                        chat_history: list[ConversationMessage],
+                        additional_params: Optional[dict[str, str]] = None) -> ConversationMessage:
         try:
             params = {
                 'botId': self.bot_id,

python/src/multi_agent_orchestrator/classifiers/bedrock_classifier.py

@@ -6,7 +6,7 @@
 from multi_agent_orchestrator.utils import Logger
 from multi_agent_orchestrator.types import ConversationMessage, ParticipantRole, BEDROCK_MODEL_ID_CLAUDE_3_5_SONNET
 from multi_agent_orchestrator.classifiers import Classifier, ClassifierResult
-
+from multi_agent_orchestrator.shared import user_agent
 
 class BedrockClassifierOptions:
     def __init__(
@@ -29,7 +29,10 @@ def __init__(self, options: BedrockClassifierOptions):
         if options.client:
             self.client = options.client
         else:
-            self.client = boto3.client('bedrock-runtime', region_name=self.region)
+            self.client = boto3.client('bedrock-runtime',region_name=self.region)
+
+        user_agent.register_feature_to_client(self.client, feature="bedrock-classifier")
+
         self.model_id = options.model_id or BEDROCK_MODEL_ID_CLAUDE_3_5_SONNET
         self.system_prompt: str
         self.inference_config = {

python/src/multi_agent_orchestrator/shared/__init__.py

@@ -0,0 +1,208 @@
+import logging
+import os
+
+from .version import VERSION
+
+mao_version = VERSION
+inject_header = True
+
+try:
+    import botocore
+except ImportError:
+    # if botocore failed to import, user might be using custom runtime and we can't inject header
+    inject_header = False
+
+logger = logging.getLogger(__name__)
+
+EXEC_ENV = os.environ.get("AWS_EXECUTION_ENV", "NA")
+TARGET_SDK_EVENT = "request-created"
+FEATURE_PREFIX = "MAOPY"
+DEFAULT_FEATURE = "no-op"
+HEADER_NO_OP = f"{FEATURE_PREFIX}/{DEFAULT_FEATURE}/{mao_version} MAOPYEnv/{EXEC_ENV}"
+
+
+def _initializer_botocore_session(session):
+    """
+    This function is used to add an extra header for the User-Agent in the Botocore session,
+    as described in the pull request: https://github.com/boto/botocore/pull/2682
+
+    Parameters
+    ----------
+    session : botocore.session.Session
+        The Botocore session to which the user-agent function will be registered.
+
+    Raises
+    ------
+    Exception
+        If there is an issue while adding the extra header for the User-Agent.
+
+    """
+    try:
+        session.register(TARGET_SDK_EVENT, _create_feature_function(DEFAULT_FEATURE))
+    except Exception:
+        logger.debug("Can't add extra header User-Agent")
+
+
+def _create_feature_function(feature):
+    """
+    Create and return the `add_mao_feature` function.
+
+    The `add_mao_feature` function is designed to be registered in boto3's event system.
+    When registered, it appends the given feature string to the User-Agent header of AWS SDK requests.
+
+    Parameters
+    ----------
+    feature : str
+        The feature string to be appended to the User-Agent header.
+
+    Returns
+    -------
+    add_mao_feature : Callable
+        The `add_mao_feature` function that modifies the User-Agent header.
+
+
+    """
+
+    def add_mao_feature(request, **kwargs):
+        try:
+            headers = request.headers
+            header_user_agent = (
+                f"{headers['User-Agent']} {FEATURE_PREFIX}/{feature}/{mao_version} MAOEnv/{EXEC_ENV}"
+            )
+
+            # This function is exclusive to client and resources objects created in MAO
+            # and must remove the no-op header, if present
+            if HEADER_NO_OP in headers["User-Agent"] and feature != DEFAULT_FEATURE:
+                # Remove HEADER_NO_OP + space
+                header_user_agent = header_user_agent.replace(f"{HEADER_NO_OP} ", "")
+
+            headers["User-Agent"] = f"{header_user_agent}"
+        except Exception:
+            logger.debug("Can't find User-Agent header")
+
+    return add_mao_feature
+
+
+# Add feature user-agent to given sdk boto3.session
+def register_feature_to_session(session, feature):
+    """
+    Register the given feature string to the event system of the provided boto3 session
+    and append the feature to the User-Agent header of the request
+
+    Parameters
+    ----------
+    session : boto3.session.Session
+        The boto3 session to which the feature will be registered.
+    feature : str
+        The feature string to be added to the User-Agent header, e.g., "00000001"  (Bedrock) in MAO.
+
+    Raises
+    ------
+    AttributeError
+        If the provided session does not have an event system.
+
+    """
+    try:
+        session.events.register(TARGET_SDK_EVENT, _create_feature_function(feature))
+    except AttributeError as e:
+        logger.debug(f"session passed in doesn't have a event system:{e}")
+
+
+# Add feature user-agent to given sdk botocore.session.Session
+def register_feature_to_botocore_session(botocore_session, feature):
+    """
+    Register the given feature string to the event system of the provided botocore session
+
+    Please notice this function is for patching botocore session and is different from
+    previous one which is for patching boto3 session
+
+    Parameters
+    ----------
+    botocore_session : botocore.session.Session
+        The botocore session to which the feature will be registered.
+    feature : str
+        The feature value to be added to the User-Agent header, e.g., "00000001" (Bedrock runtime) in MAO.
+
+    Raises
+    ------
+    AttributeError
+        If the provided session does not have an event system.
+
+    Examples
+    --------
+    **register led-bot user-agent to botocore session**
+
+        >>> from multi_agent_orchestrator.shared.user_agent import (
+        >>>    register_feature_to_botocore_session
+        >>> )
+        >>>
+        >>> session = botocore.session.Session()
+        >>> register_feature_to_botocore_session(botocore_session=session, feature="data-masking")
+        >>> key_provider = StrictAwsKmsMasterKeyProvider(key_ids=self.keys, botocore_session=session)
+
+    """
+    try:
+        botocore_session.register(TARGET_SDK_EVENT, _create_feature_function(feature))
+    except AttributeError as e:
+        logger.debug(f"botocore session passed in doesn't have a event system:{e}")
+
+
+# Add feature user-agent to given sdk boto3.client
+def register_feature_to_client(client, feature):
+    """
+    Register the given feature string to the event system of the provided boto3 client
+    and append the feature to the User-Agent header of the request
+
+    Parameters
+    ----------
+    client : boto3.session.Session.client
+        The boto3 client to which the feature will be registered.
+    feature : str
+        The feature value to be added to the User-Agent header, e.g., "00000001" (Bedrock runtime) in MAO.
+
+    Raises
+    ------
+    AttributeError
+        If the provided client does not have an event system.
+
+    """
+    try:
+        client.meta.events.register(TARGET_SDK_EVENT, _create_feature_function(feature))
+    except AttributeError as e:
+        logger.debug(f"session passed in doesn't have a event system:{e}")
+
+
+# Add feature user-agent to given sdk boto3.resource
+def register_feature_to_resource(resource, feature):
+    """
+    Register the given feature string to the event system of the provided boto3 resource
+    and append the feature to the User-Agent header of the request
+
+    Parameters
+    ----------
+    resource : boto3.session.Session.resource
+        The boto3 resource to which the feature will be registered.
+    feature : str
+        The feature value to be added to the User-Agent header, e.g., "00000001" (Bedrock runtime) in MAO.
+
+    Raises
+    ------
+    AttributeError
+        If the provided resource does not have an event system.
+
+    """
+    try:
+        resource.meta.client.meta.events.register(TARGET_SDK_EVENT, _create_feature_function(feature))
+    except AttributeError as e:
+        logger.debug(f"resource passed in doesn't have a event system:{e}")
+
+
+def inject_user_agent():
+    if inject_header:
+        # Some older botocore versions doesn't support register_initializer. In those cases, we disable the feature.
+        if not hasattr(botocore, "register_initializer"):
+            return
+
+        # Customize botocore session to inject Boto3 header
+        # See: https://github.com/boto/botocore/pull/2682
+        botocore.register_initializer(_initializer_botocore_session)

python/src/multi_agent_orchestrator/shared/user_agent.py

@@ -0,0 +1,3 @@
+"""Exposes version constant."""
+
+VERSION = "0.1.10"