Added code example for using Bedrock Prompt Routing in BedrockClassif…

…ier and BedrockLLMAgent

README.md

@@ -95,6 +95,7 @@ Get hands-on experience with the Multi-Agent Orchestrator through our diverse se
   - [`fast-api-streaming`](https://github.com/awslabs/multi-agent-orchestrator/tree/main/examples/fast-api-streaming): FastAPI implementation with streaming support
   - [`text-2-structured-output`](https://github.com/awslabs/multi-agent-orchestrator/tree/main/examples/text-2-structured-output): Natural Language to Structured Data
   - [`bedrock-inline-agents`](https://github.com/awslabs/multi-agent-orchestrator/tree/main/examples/bedrock-inline-agents): Bedrock Inline Agents sample
+  - [`bedrock-prompt-routing`](https://github.com/awslabs/multi-agent-orchestrator/tree/main/examples/bedrock-prompt-routing): Bedrock Prompt Routing sample code for `BedrockClassifier`and `BedrockLLMAgent`
 
 
 All examples are available in both Python and TypeScript implementations. Check out our [documentation](https://awslabs.github.io/multi-agent-orchestrator/) for comprehensive guides on setting up and using the Multi-Agent Orchestrator!

examples/bedrock-prompt-routing/main.py

@@ -0,0 +1,107 @@
+
+import uuid
+import asyncio
+import os
+from typing import Optional, Any
+import json
+import sys
+from multi_agent_orchestrator.orchestrator import MultiAgentOrchestrator, OrchestratorConfig
+from multi_agent_orchestrator.agents import (BedrockLLMAgent,
+                        BedrockLLMAgentOptions,
+                        AgentResponse,
+                        AgentCallbacks)
+from multi_agent_orchestrator.types import ConversationMessage, ParticipantRole
+from multi_agent_orchestrator.classifiers import BedrockClassifier, BedrockClassifierOptions
+
+class LLMAgentCallbacks(AgentCallbacks):
+    def on_llm_new_token(self, token: str) -> None:
+        # handle response streaming here
+        print(token, end='', flush=True)
+
+
+async def handle_request(_orchestrator: MultiAgentOrchestrator, _user_input:str, _user_id:str, _session_id:str):
+    response:AgentResponse = await _orchestrator.route_request(_user_input, _user_id, _session_id)
+
+    # Print metadata
+    print("\nMetadata:")
+    print(f"Selected Agent: {response.metadata.agent_name}")
+    if isinstance(response, AgentResponse) and response.streaming is False:
+        # Handle regular response
+        if isinstance(response.output, str):
+            print(response.output)
+        elif isinstance(response.output, ConversationMessage):
+                print(response.output.content[0].get('text'))
+
+def custom_input_payload_encoder(input_text: str,
+                                 chat_history: list[Any],
+                                 user_id: str,
+                                 session_id: str,
+                                 additional_params: Optional[dict[str, str]] = None) -> str:
+    return json.dumps({
+        'hello':'world'
+    })
+
+def custom_output_payload_decoder(response: dict[str, Any]) -> Any:
+    decoded_response = json.loads(
+        json.loads(
+            response['Payload'].read().decode('utf-8')
+        )['body'])['response']
+    return ConversationMessage(
+            role=ParticipantRole.ASSISTANT.value,
+            content=[{'text': decoded_response}]
+        )
+
+if __name__ == "__main__":
+
+    # Initialize the orchestrator with some options
+    orchestrator = MultiAgentOrchestrator(options=OrchestratorConfig(
+        LOG_AGENT_CHAT=True,
+        LOG_CLASSIFIER_CHAT=True,
+        LOG_CLASSIFIER_RAW_OUTPUT=True,
+        LOG_CLASSIFIER_OUTPUT=True,
+        LOG_EXECUTION_TIMES=True,
+        MAX_RETRIES=3,
+        USE_DEFAULT_AGENT_IF_NONE_IDENTIFIED=True,
+        MAX_MESSAGE_PAIRS_PER_AGENT=10,
+    ),
+    classifier=BedrockClassifier(BedrockClassifierOptions(
+        model_id=f"arn:aws:bedrock:us-east-1:{os.getenv('AWS_ACCOUNT_ID')}:default-prompt-router/anthropic.claude:1"))
+    )
+
+    # Add some agents
+    tech_agent = BedrockLLMAgent(BedrockLLMAgentOptions(
+        name="Tech Agent",
+        streaming=True,
+        description="Specializes in technology areas including software development, hardware, AI, \
+            cybersecurity, blockchain, cloud computing, emerging tech innovations, and pricing/costs \
+            related to technology products and services.",
+        model_id="anthropic.claude-3-sonnet-20240229-v1:0",
+        callbacks=LLMAgentCallbacks()
+    ))
+    orchestrator.add_agent(tech_agent)
+
+    # Add some agents
+    health_agent = BedrockLLMAgent(BedrockLLMAgentOptions(
+        name="Health Agent",
+        streaming=False,
+        model_id=f"arn:aws:bedrock:us-east-1:{os.getenv('AWS_ACCOUNT_ID')}:default-prompt-router/anthropic.claude:1",
+        description="Specialized agent for giving health advice.",
+        callbacks=LLMAgentCallbacks()
+    ))
+    orchestrator.add_agent(health_agent)
+
+    USER_ID = "user123"
+    SESSION_ID = str(uuid.uuid4())
+
+    print("Welcome to the interactive Multi-Agent system. Type 'quit' to exit.")
+
+    while True:
+        # Get user input
+        user_input = input("\nYou: ").strip()
+
+        if user_input.lower() == 'quit':
+            print("Exiting the program. Goodbye!")
+            sys.exit()
+
+        # Run the async function
+        asyncio.run(handle_request(orchestrator, user_input, USER_ID, SESSION_ID))

examples/bedrock-prompt-routing/readme.md

@@ -0,0 +1,31 @@
+# Bedrock Prompt Routing Example
+
+This guide demonstrates how to implement and utilize [Amazon Bedrock Prompt Routing](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-routing.html) functionality with your `BedrockClassifier` or `BedrockLLMAgent`. Prompt routing helps optimize model selection based on your input patterns, improving both performance and cost-effectiveness.
+
+## Prerequisites
+Before running this example, ensure you have:
+
+- An active AWS account with Bedrock access
+- Python installed on your system (version 3.11 or higher recommended)
+- The AWS SDK for Python (Boto3) installed
+- Appropriate IAM permissions configured for Bedrock access
+
+## Installation
+
+First, install the required dependencies by running:
+
+```bash
+pip install boto3 multi-agent-orchestrator
+```
+
+export your AWS_ACCOUNT_ID variable by running:
+```bash
+export AWS_ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text)
+```
+
+## Running the Example
+
+```bash
+python main.py
+```
+