Structured Outputs with Amazon Bedrock in AG2#

%pip install ag2 boto3 pydantic --upgrade

import json
import os

from dotenv import load_dotenv
from pydantic import BaseModel

from autogen import ConversableAgent, LLMConfig

load_dotenv()

print("Libraries imported successfully!")

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

    explanation: str  # What operation or reasoning is being performed
    output: str  # The result of this step

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

    steps: list[Step]  # List of all steps taken
    final_answer: str  # The final answer

    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)
        )
        return f"{steps_output}\n\nFinal Answer: {self.final_answer}"

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

# Configure LLM with Bedrock and structured outputs
llm_config = LLMConfig(
    config_list={
        "api_type": "bedrock",
        "model": "qwen.qwen3-coder-480b-a35b-v1:0",
        "api_key": os.getenv("BEDROCK_API_KEY"),
        "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": MathReasoning,
        "aws_profile_name": os.getenv("AWS_PROFILE"),
    },
    cache_seed=42,  # Optional: for reproducible results
)
print("Bedrock LLM configuration created with structured outputs!")

# Create agent with structured output capability
math_agent = ConversableAgent(
    name="math_assistant",
    llm_config=llm_config,
    system_message="""You are a helpful math assistant that solves problems step by step.
    Always show your reasoning process clearly with explanations for each step.
    Return your response in the structured format requested.""",
    max_consecutive_auto_reply=1,
    human_input_mode="NEVER",
)

print(f"Agent '{math_agent.name}' created successfully!")

print("=== Example 1: Solve equation with structured output ===")

# Initiate chat with the agent
result1 = math_agent.run(
    message="Solve the equation: 2x + 5 = -25.",
    max_turns=5,
).process()

print("=== Example 2: Complex math problem ===")

result2 = math_agent.run(
    message="Designed to cause debate; some solve \\((2+2)\\) first (4), then \\(8\\div 2\\) (4), then \\(4\times 4\\) (16); others do \\(2\times 4\\) (8) first, then \\(8\\div 8\\) (1). use 10 steps to solve",
    max_turns=5,
).process()

# Define schema as a dictionary (JSON Schema format)
dict_schema = {
    "type": "object",
    "properties": {
        "problem": {"type": "string", "description": "The math problem being solved"},
        "solution_steps": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {"step": {"type": "string"}, "result": {"type": "string"}},
                "required": ["step", "result"],
            },
        },
        "answer": {"type": "string"},
    },
    "required": ["problem", "solution_steps", "answer"],
}

print("Dict schema defined:")
print(json.dumps(dict_schema, indent=2))

# Create a new LLM config with dict schema
llm_config_dict = LLMConfig(
    config_list={
        "api_type": "bedrock",
        "model": "eu.anthropic.claude-3-7-sonnet-20250219-v1:0",
        "api_key": os.getenv("BEDROCK_API_KEY"),
        "aws_region": os.getenv("AWS_REGION", "us-east-1"),
        "aws_access_key": os.getenv("AWS_ACCESS_KEY_ID"),
        "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
        "response_format": dict_schema,  # Using dict schema instead of Pydantic model
    },
)

# Create agent with dict schema
math_agent_dict = ConversableAgent(
    name="math_assistant_dict",
    llm_config=llm_config_dict,
    system_message="You are a helpful math assistant.",
    max_consecutive_auto_reply=1,
    human_input_mode="NEVER",
)

print("Agent created with dict schema!")

print("=== Example 3: Using dict schema ===")

result3 = math_agent_dict.run(
    message="Solve: x^2 - 5x + 6 = 0",
    max_turns=5,
).process()

# Inspect the tool that gets created from your schema
from autogen.oai.bedrock import BedrockClient

# Create a temporary client to see the tool creation
temp_client = BedrockClient(
    aws_region=os.getenv("AWS_REGION", "us-east-1"),
    aws_access_key=os.getenv("AWS_ACCESS_KEY_ID"),
    aws_secret_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
    response_format=MathReasoning,
)

# See how the schema is converted to a tool
tool = temp_client._create_structured_output_tool(MathReasoning)

print("Tool definition created from MathReasoning schema:")
print(json.dumps(tool, indent=2))

class DetailedMathReasoning(BaseModel):
    """Enhanced math reasoning with custom formatting."""

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

    def format(self) -> str:
        """Custom formatted output with problem statement."""
        output = f"Problem: {self.problem}\n\n"
        output += "Solution Steps:\n"
        for i, step in enumerate(self.steps, 1):
            output += f"  {i}. {step.explanation}\n"
            output += f"     → {step.output}\n"
        output += f"\nFinal Answer: {self.final_answer}"
        if self.verification:
            output += f"\n\nVerification: {self.verification}"
        return output

    def to_markdown(self) -> str:
        """Export as Markdown format."""
        md = f"## Problem\n\n{self.problem}\n\n"
        md += "## Solution\n\n"
        for i, step in enumerate(self.steps, 1):
            md += f"### Step {i}\n\n"
            md += f"**Explanation**: {step.explanation}\n\n"
            md += f"**Result**: `{step.output}`\n\n"
        md += f"## Final Answer\n\n`{self.final_answer}`"
        return md

# Create agent with enhanced model
enhanced_llm_config = LLMConfig(
    config_list={
        "api_type": "bedrock",
        "model": "eu.anthropic.claude-3-7-sonnet-20250219-v1:0",
        "aws_region": os.getenv("AWS_REGION", "us-east-1"),
        "aws_access_key": os.getenv("AWS_ACCESS_KEY_ID"),
        "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
        "response_format": DetailedMathReasoning,
    },
)

enhanced_agent = ConversableAgent(
    name="enhanced_math_assistant",
    llm_config=enhanced_llm_config,
    system_message="You are a detailed math assistant. Always verify your answers.",
    max_consecutive_auto_reply=1,
)

print("Enhanced agent created with custom formatting!")

# Test the enhanced agent
result = enhanced_agent.run(
    recipient=enhanced_agent,
    message="Solve: 4x - 8 = 12. Show your work and verify the answer.",
    max_turns=10,
).process()

from autogen.agentchat.group.patterns import AutoPattern

# Create a simple reviewer agent (without structured output) to work with the enhanced agent
reviewer_llm_config = LLMConfig(
    config_list={
        "api_type": "bedrock",
        "model": "eu.anthropic.claude-3-7-sonnet-20250219-v1:0",
        "aws_region": os.getenv("AWS_REGION", "us-east-1"),
        "aws_access_key": os.getenv("AWS_ACCESS_KEY_ID"),
        "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
        # No response_format - this agent will use regular text responses
    },
)

reviewer_agent = ConversableAgent(
    name="math_reviewer",
    llm_config=reviewer_llm_config,
    system_message="You are a math reviewer. Review the solutions provided by the math assistant and provide feedback. Keep your reviews brief and constructive.",
    max_consecutive_auto_reply=1,
)

# Create AutoPattern for groupchat
pattern = AutoPattern(
    initial_agent=enhanced_agent,
    agents=[enhanced_agent, reviewer_agent],
    group_manager_args={
        "llm_config": enhanced_llm_config,  # Use same config for group manager
    },
)

print("AutoPattern created with structured output agent!")

import os

from dotenv import load_dotenv
from pydantic import BaseModel

from autogen import ConversableAgent, LLMConfig, UserProxyAgent
from autogen.agentchat import initiate_group_chat
from autogen.agentchat.group.patterns.auto import AutoPattern

load_dotenv()

# Define structured output models
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] = []

class WorkflowStatus(BaseModel):
    """Status of the current workflow execution."""

    current_stage: str
    completed_stages: list[str] = []
    pending_stages: list[str] = []
    issues: list[str] = []
    progress_percentage: int | None = None

class PipelineOrchestrationResponse(BaseModel):
    """Complete structured response from the pipeline orchestrator."""

    routing_decision: RoutingDecision
    workflow_status: WorkflowStatus | None = None

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

# Regular LLM config for other agents
llm_config = LLMConfig(
    config_list={
        "api_type": "bedrock",
        "model": "eu.anthropic.claude-3-7-sonnet-20250219-v1:0",
        "api_key": os.getenv("BEDROCK_API_KEY"),
        "aws_region": os.getenv("AWS_REGION"),
        "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
        "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
        "aws_profile_name": os.getenv("AWS_PROFILE"),
    },
    cache_seed=42,
)

orchestrator_llm_config = LLMConfig(
    config_list={
        "api_type": "bedrock",
        "model": "eu.anthropic.claude-3-7-sonnet-20250219-v1:0",
        "api_key": os.getenv("BEDROCK_API_KEY"),
        "aws_region": os.getenv("AWS_REGION"),
        "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
        "aws_secret_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
        "aws_profile_name": os.getenv("AWS_PROFILE"),
        "response_format": PipelineOrchestrationResponse,
    },
    cache_seed=42,
)

print("Bedrock LLM configuration created 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. New project request? → Route to project_creator      │
    │ 2. Code development? → Route to code_developer          │
    │ 3. Code analysis needed? → Route to code_quality_analyzer│
    │ 4. Configuration updates? → Route to config_manager     │
    │ 5. Build/deploy request? → Route to build_deploy_agent  │
    │ 6. Validation needed? → Route to deployment_validator   │
    └─────────────────────────────────────────────────────────┘
    Always provide clear routing decisions with reasoning.""",
    llm_config=orchestrator_llm_config,  # Use config with structured output
)

project_creator = ConversableAgent(
    name="project_creator",
    system_message="""🏗️ You are the Project Creator. Your role is to:
    Use the APPLY_PATCH tool to:
    ✓ Create project structure and directories
    ✓ Generate initial files (README.md, .gitignore, LICENSE)
    ✓ Set up configuration files (package.json, requirements.txt, pom.xml)
    ✓ Create initial source code templates
    ✓ Set up test directory structure
    Project Creation Checklist:
    ┌──────────────────────────────────────────────┐
    │ □ Project directory structure created?       │
    │ □ Configuration files initialized?           │
    │ □ Initial source code files created?         │
    │ □ Test directory structure set up?           │
    │ □ Documentation files added?                 │
    └──────────────────────────────────────────────┘
    After project creation:
    → Route to code_developer for implementation
    Create clean, well-structured project foundations.""",
    llm_config=llm_config,
)

code_developer = ConversableAgent(
    name="code_developer",
    system_message="""💻 You are the Code Developer. Your role is to:
    Use the APPLY_PATCH tool to:
    ✓ Write application code and implement features
    ✓ Create modules, classes, and functions
    ✓ Add business logic and algorithms
    ✓ Implement API endpoints and routes
    ✓ Write unit tests for new code
    Development Guidelines:
    ┌──────────────────────────────────────────────┐
    │ 1. Follow best practices and coding standards│
    │ 2. Write clean, maintainable code            │
    │ 3. Include proper error handling             │
    │ 4. Add comprehensive unit tests              │
    │ 5. Document complex logic                    │
    └──────────────────────────────────────────────┘

    After development:
    → Route to code_quality_analyzer for testing and validation
    Write production-ready code with tests!""",
    llm_config=llm_config,
)

code_analyzer = ConversableAgent(
    name="code_quality_analyzer",
    system_message="""🔍 You are the Code Quality Analyzer. Your role is to:
    Use the SHELL tool to:
    ✓ Run test suites: pytest, unittest, npm test, mvn test
    ✓ Run linters: pylint, flake8, black --check, eslint, prettier
    ✓ Check git status: git status, git diff, git log
    ✓ Analyze code coverage: coverage report, pytest --cov
    ✓ Check build status: npm run build, mvn compile
    Analysis Checklist:
    ┌──────────────────────────────────────────────┐
    │ □ All tests passing?                         │
    │ □ No linting errors?                         │
    │ □ Code coverage acceptable?                  │
    │ □ No merge conflicts?                        │
    │ □ Build successful?                          │
    └──────────────────────────────────────────────┘
    Routing Rules:
    • Issues found → Route to code_fixer
    • All checks pass → Route to config_manager for CI/CD setup
    Provide detailed analysis reports with specific issues identified.""",
    llm_config=llm_config,
)

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

pattern = AutoPattern(
    initial_agent=orchestrator,
    agents=[
        orchestrator,
        project_creator,
        code_developer,
        code_analyzer,
    ],
    user_agent=user,
    group_manager_args={"llm_config": orchestrator_llm_config},
)

result, context, last_agent = initiate_group_chat(
    pattern=pattern,
    messages="I want to create a new project. for a coffee machine business",
    max_rounds=5,
)
print(result)