Amazon Bedrock V2 Client with AG2#

%pip install ag2 boto3 pydantic --upgrade

import os

from dotenv import load_dotenv
from pydantic import BaseModel

from autogen import ConversableAgent, LLMConfig

load_dotenv()

print("Libraries imported successfully!")

# Configure LLM to use Bedrock V2 client
llm_config_v2 = LLMConfig(
    config_list=[
        {
            "api_type": "bedrock_v2",  # <-- Key: use V2 client architecture
            "model": "qwen.qwen3-coder-480b-a35b-v1:0",
            "aws_region": os.getenv("AWS_REGION", "eu-north-1"),
            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
            "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
            "aws_profile_name": os.getenv("AWS_PROFILE"),
        }
    ],
)

# Compare with V1 client configuration
llm_config_v1 = LLMConfig(
    config_list=[
        {
            "api_type": "bedrock",  # <-- V1 client architecture
            "model": "qwen.qwen3-coder-480b-a35b-v1:0",
            "aws_region": os.getenv("AWS_REGION", "eu-north-1"),
            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
            "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
            "aws_profile_name": os.getenv("AWS_PROFILE"),
        }
    ],
    temperature=0.3,
)

print("Bedrock V2 and V1 configurations created!")

from autogen.llm_clients.bedrock_v2 import BedrockV2Client
from autogen.llm_clients.models import UnifiedResponse

# Create Bedrock V2 client directly
client = BedrockV2Client(
    aws_region=os.getenv("AWS_REGION", "eu-north-1"),
    aws_access_key=os.getenv("AWS_ACCESS_KEY"),
    aws_secret_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
)

# Make a request
response = client.create({
    "model": "qwen.qwen3-coder-480b-a35b-v1:0",
    "messages": [{"role": "user", "content": "Explain quantum computing in 2 sentences."}],
})

# Verify it's a UnifiedResponse
print(f"Response type: {type(response)}")
print(f"Is UnifiedResponse: {isinstance(response, UnifiedResponse)}")
print(f"\nProvider: {response.provider}")
print(f"Model: {response.model}")
print(f"\nText content: {response.text}")
print(f"\nUsage: {response.usage}")
print(f"Cost: ${response.cost:.6f}" if response.cost else "Cost: N/A")

# Access individual messages and content blocks
for i, message in enumerate(response.messages):
    print(f"\nMessage {i + 1}:")
    print(f"  Role: {message.role}")
    print(f"  Content blocks: {len(message.content)}")

    for j, block in enumerate(message.content):
        print(f"\n  Block {j + 1}:")
        print(f"    Type: {block.type}")
        print(f"    Class: {type(block).__name__}")

        # Access text content
        if hasattr(block, "text"):
            print(f"    Text: {block.text[:100]}..." if len(block.text) > 100 else f"    Text: {block.text}")

        # Access tool calls
        if hasattr(block, "name"):
            print(f"    Tool: {block.name}")
            print(
                f"    Arguments: {block.arguments[:100]}..."
                if len(block.arguments) > 100
                else f"    Arguments: {block.arguments}"
            )

# Use helper methods
print("\n" + "=" * 60)
print("Using helper methods:")
print(f"All text: {response.text}")
print(f"Tool calls: {len(response.get_content_by_type('tool_call'))}")
print(f"Text blocks: {len(response.get_content_by_type('text'))}")

# Define structured output model
class Step(BaseModel):
    """Represents a single step in solving a problem."""

    explanation: str
    output: str

class ProblemSolution(BaseModel):
    """Complete structured response for a problem solution."""

    problem: str
    steps: list[Step]
    final_answer: str
    confidence: float | None = None

    def format(self) -> str:
        """Format the structured output for human-readable display."""
        steps_output = "\n".join(
            f"Step {i + 1}: {step.explanation}\n  Output: {step.output}" for i, step in enumerate(self.steps)
        )
        confidence_str = f" (Confidence: {self.confidence})" if self.confidence else ""
        return f"Problem: {self.problem}\n\n{steps_output}\n\nFinal Answer: {self.final_answer}{confidence_str}"

print("Pydantic models defined:")
print(f"- Step: {Step.model_json_schema()}")
print(f"- ProblemSolution: {ProblemSolution.model_json_schema()}")

# Configure Bedrock V2 with structured outputs

llm_config_v2_structured = LLMConfig(
    config_list={
        "api_type": "bedrock_v2",
        "model": "qwen.qwen3-coder-480b-a35b-v1:0",
        "aws_region": os.getenv("AWS_REGION", "eu-north-1"),
        "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
        "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
        "total_max_attempts": 8,
        "mode": "adaptive",  # Retries with client-side throttling
        "response_format": ProblemSolution,
    },
)

# Create agent with structured outputs
math_agent = ConversableAgent(
    name="math_assistant",
    llm_config=llm_config_v2_structured,
    system_message="You are a helpful math assistant that solves problems step by step. Always show your reasoning process clearly.",
    max_consecutive_auto_reply=1,
    human_input_mode="NEVER",
)

print("Agent created with Bedrock V2 and structured outputs!")

# Test the agent with a math problem
print("=== Solving Math Problem with Structured Output ===")

result = math_agent.run(
    message="Solve the equation: 3x + 7 = 22. Show all steps.",
    max_turns=1,
).process()

print("\nResponse received!")

# Create agents with different client versions
agent_v2 = ConversableAgent(
    name="agent_v2",
    llm_config=llm_config_v2,
    system_message="You are a helpful assistant using V2 client architecture.",
    max_consecutive_auto_reply=1,
    human_input_mode="NEVER",
)

agent_v1 = ConversableAgent(
    name="agent_v1",
    llm_config=llm_config_v1,
    system_message="You are a helpful assistant using V1 client architecture.",
    max_consecutive_auto_reply=1,
    human_input_mode="NEVER",
)

print("Agents created with V1 and V2 clients!")

# Test both agents with the same question
question = "What are the three main benefits of renewable energy?"

print("=== V2 Client Response ===")
result_v2 = agent_v2.run(message=question, max_turns=1).process()

print("\n=== V1 Client Response ===")
result_v1 = agent_v1.run(message=question, max_turns=1).process()

print("\nBoth clients work seamlessly with the same interface!")

from autogen import GroupChat, GroupChatManager

# Planner agent - uses V2 client
planner = ConversableAgent(
    name="planner_agent",
    llm_config=llm_config_v2,
    system_message="Create detailed project plans. Break down tasks into clear steps.",
    description="Creates project plans",
)

# Reviewer agent - uses V1 client (demonstrates compatibility)
reviewer = ConversableAgent(
    name="reviewer_agent",
    llm_config=llm_config_v1,
    system_message="Review plans and provide constructive feedback. Keep reviews concise.",
    description="Reviews plans",
)

# Coordinator agent - uses V2 client
coordinator = ConversableAgent(
    name="coordinator_agent",
    llm_config=llm_config_v2,
    system_message="Coordinate between planner and reviewer. Say DONE! when the plan is finalized.",
    description="Coordinates the planning process",
)

# Setup group chat
groupchat = GroupChat(
    agents=[coordinator, planner, reviewer],
    speaker_selection_method="auto",
    messages=[],
)

# Create manager with V2 client
manager = GroupChatManager(
    name="group_manager",
    groupchat=groupchat,
    llm_config=llm_config_v2,
    is_termination_msg=lambda x: "DONE!" in (x.get("content", "") or "").upper(),
)

print("Group chat created with mixed V1/V2 Bedrock clients!")

# Start the conversation
print("=== Starting Group Chat ===")
chat_result = coordinator.initiate_chat(
    recipient=manager,
    message="Let's create a plan for organizing a tech conference.",
)

print("\n=== Chat History ===")
for msg in chat_result.chat_history:
    print(
        f"\n[{msg.get('role', 'unknown')}]: {msg.get('content', '')[:200]}..."
        if len(msg.get("content", "")) > 200
        else f"\n[{msg.get('role', 'unknown')}]: {msg.get('content', '')}"
    )

# Define structured output models for orchestration
class TaskDetails(BaseModel):
    """Details about the task being processed."""

    task_type: str
    description: str
    priority: str | None = None
    requirements: list[str] = []

class RoutingDecision(BaseModel):
    """Structured routing decision from the orchestrator."""

    request_analysis: str
    task_details: TaskDetails
    selected_agent: str
    routing_reason: str
    expected_outcome: str
    next_steps: list[str] = []

    def format(self) -> str:
        """Format the structured output for human-readable display."""
        output = "🎯 Routing Decision\n"
        output += f"{'=' * 60}\n\n"
        output += f"Request Analysis:\n{self.request_analysis}\n\n"
        output += f"Task Type: {self.task_details.task_type}\n"
        output += f"Description: {self.task_details.description}\n\n"
        output += "Routing Decision:\n"
        output += f"  → Selected Agent: {self.selected_agent}\n"
        output += f"  → Reason: {self.routing_reason}\n"
        output += f"  → Expected Outcome: {self.expected_outcome}\n"
        if self.next_steps:
            output += "\nNext Steps:\n"
            for i, step in enumerate(self.next_steps, 1):
                output += f"  {i}. {step}\n"
        return output

print("Structured output models defined for orchestration!")

# Configure orchestrator with structured outputs
orchestrator_llm_config = LLMConfig(
    config_list=[
        {
            "api_type": "bedrock_v2",  # V2 client with structured outputs
            "model": "aqwen.qwen3-coder-480b-a35b-v1:0",
            "aws_region": os.getenv("AWS_REGION", "eu-north-1"),
            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
            "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
            "response_format": RoutingDecision,  # Structured output for routing
        }
    ],
    temperature=0.3,
)

# Regular config for other agents
regular_llm_config = LLMConfig(
    config_list=[
        {
            "api_type": "bedrock_v2",  # V2 client without structured outputs
            "model": "aqwen.qwen3-coder-480b-a35b-v1:0",
            "aws_region": os.getenv("AWS_REGION", "eu-north-1"),
            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
            "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
        }
    ],
    temperature=0.3,
)

print("LLM configurations created!")

from autogen import UserProxyAgent
from autogen.agentchat import initiate_group_chat
from autogen.agentchat.group.patterns.auto import AutoPattern

# Create orchestrator agent with structured outputs
orchestrator = ConversableAgent(
    name="pipeline_orchestrator",
    system_message="""🎯 You are the Pipeline Orchestrator. Your role is to:
    • Analyze user requests and determine the workflow path
    • Route tasks to appropriate specialized agents
    • Monitor pipeline progress and coordinate handoffs
    • Report final results to the user

    You MUST provide structured routing decisions that include:
    - Analysis of the user's request
    - Task type and details
    - Selected agent and reasoning
    - Expected outcome and next steps

    Workflow Decision Logic:
    ┌─────────────────────────────────────────────────────────┐
    │ 1. Research task? → Route to researcher                  │
    │ 2. Writing task? → Route to writer                       │
    │ 3. Analysis task? → Route to analyst                     │
    │ 4. Planning task? → Route to planner                     │
    └─────────────────────────────────────────────────────────┘
    Always provide clear routing decisions with reasoning.""",
    llm_config=orchestrator_llm_config,
)

# Create specialized agents
researcher = ConversableAgent(
    name="researcher",
    system_message="""🔍 You are a Researcher. Your role is to:
    • Conduct thorough research on given topics
    • Gather relevant information from multiple sources
    • Provide well-sourced findings
    • Summarize key points clearly
    Always cite your sources and provide comprehensive research.""",
    llm_config=regular_llm_config,
)

writer = ConversableAgent(
    name="writer",
    system_message="""✍️ You are a Writer. Your role is to:
    • Create well-structured written content
    • Follow style guidelines and best practices
    • Ensure clarity and coherence
    • Adapt tone and style to the audience
    Write engaging, clear, and professional content.""",
    llm_config=regular_llm_config,
)

analyst = ConversableAgent(
    name="analyst",
    system_message="""📊 You are an Analyst. Your role is to:
    • Analyze data and information systematically
    • Identify patterns and trends
    • Provide insights and recommendations
    • Support conclusions with evidence
    Provide thorough, data-driven analysis.""",
    llm_config=regular_llm_config,
)

planner_agent = ConversableAgent(
    name="planner",
    system_message="""📋 You are a Planner. Your role is to:
    • Create detailed plans and roadmaps
    • Break down complex tasks into steps
    • Identify dependencies and timelines
    • Anticipate potential issues
    Create comprehensive, actionable plans.""",
    llm_config=regular_llm_config,
)

# Create user proxy
user = UserProxyAgent(
    name="user",
    human_input_mode="TERMINATE",
    code_execution_config={"work_dir": "coding", "use_docker": False},
)

print("All agents created!")

# Create AutoPattern for groupchat
pattern = AutoPattern(
    initial_agent=orchestrator,
    agents=[
        orchestrator,
        researcher,
        writer,
        analyst,
        planner_agent,
    ],
    user_agent=user,
    group_manager_args={"llm_config": orchestrator_llm_config},
)

print("AutoPattern created with Bedrock V2 agents!")

# Run the group chat
print("=== Starting Group Chat with Bedrock V2 and Structured Outputs ===")

result, context, last_agent = initiate_group_chat(
    pattern=pattern,
    messages="I need help creating a marketing strategy for a new coffee shop. Research the market, analyze competitors, and create a comprehensive plan.",
    max_rounds=8,
)

print("\n=== Final Result ===")
print(result)

# Example: Custom workflow that processes V2 responses
def process_v2_response(response: UnifiedResponse):
    """Process a UnifiedResponse and extract all relevant information."""
    print(f"\n{'=' * 60}")
    print(f"Processing Response from {response.provider.upper()}")
    print(f"{'=' * 60}")

    print(f"\n📝 Model: {response.model}")
    print(f"🆔 ID: {response.id}")
    print(f"✅ Status: {response.status}")
    print(f"🏁 Finish Reason: {response.finish_reason}")

    print(f"\n💬 Messages ({len(response.messages)}):")
    for i, msg in enumerate(response.messages, 1):
        print(f"\n  Message {i}:")
        print(f"    Role: {msg.role}")
        print(f"    Content Blocks: {len(msg.content)}")

        # Count content types
        text_blocks = [b for b in msg.content if b.type == "text"]
        tool_blocks = [b for b in msg.content if b.type == "tool_call"]
        image_blocks = [b for b in msg.content if b.type == "image"]

        if text_blocks:
            print(f"    📄 Text blocks: {len(text_blocks)}")
            print(
                f"       Text: {text_blocks[0].text[:100]}..."
                if len(text_blocks[0].text) > 100
                else f"       Text: {text_blocks[0].text}"
            )

        if tool_blocks:
            print(f"    🔧 Tool calls: {len(tool_blocks)}")
            for tool in tool_blocks:
                print(f"       - {tool.name}")

        if image_blocks:
            print(f"    🖼️  Images: {len(image_blocks)}")

    print("\n📊 Usage:")
    print(f"    Prompt tokens: {response.usage.get('prompt_tokens', 0)}")
    print(f"    Completion tokens: {response.usage.get('completion_tokens', 0)}")
    print(f"    Total tokens: {response.usage.get('total_tokens', 0)}")

    if response.cost:
        print(f"    💰 Cost: ${response.cost:.6f}")

    print("\n🔍 Provider Metadata:")
    for key, value in response.provider_metadata.items():
        print(f"    {key}: {value}")

# Test the function
test_response = client.create({
    "model": "aqwen.qwen3-coder-480b-a35b-v1:0",
    "messages": [{"role": "user", "content": "List 3 benefits of cloud computing."}],
})

process_v2_response(test_response)