Feat/multiagent structured output

[afarntrog]

Description

This pull request introduces support for structured outputs in multi-agent orchestration by allowing a Pydantic model to be specified for structured output from nodes. This enhancement is implemented across the core multi-agent abstractions (MultiAgentBase), as well as the Graph and Swarm orchestrators. The changes ensure that a structured_output_model can be passed through all relevant execution paths, and that agents or nodes can override or inherit this model as needed.

Structured Output Model Support

• Added an optional structured_output_model parameter (of type Type[BaseModel]) to the invoke_async, stream_async, and __call__ methods in MultiAgentBase, Graph, and Swarm classes, allowing users to specify a Pydantic model for structured node outputs.

• Updated all internal orchestration and node execution methods to propagate the structured_output_model parameter, ensuring that it is available at every level of the orchestration stack.

Agent and Node Behavior

• Ensured that when executing agent nodes, the agent's own default structured output model takes precedence, but falls back to the graph-level model if not specified, supporting flexible and hierarchical model assignment.

These changes make it easier to enforce structured outputs from multi-agent workflows, improving type safety and downstream integration.

Related Issues

#538
#1309

Documentation PR

Will add documentation once we merge (if needed)

Type of Change

New feature

Testing

How have you tested the change? Verify that the changes do not break functionality or introduce warnings in consuming repositories: agents-docs, agents-tools, agents-cli

Tested locally and added integration tests proving these changes

• I ran hatch run prepare

Checklist

• I have read the CONTRIBUTING document
• I have added any necessary tests that prove my fix is effective or my feature works
• I have updated the documentation accordingly
• I have added an appropriate example to the documentation to outline the feature, or no new docs are needed
• My changes generate no new warnings
• Any dependent changes have been merged and published

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Example

Here's a complete example you can run to see how we can allow for a graph agent with conditional edge routing:

from pydantic import BaseModel, Field
from typing import Literal, List, Optional

from strands import Agent
from strands.multiagent.graph import GraphBuilder, GraphState


# ============================================================================
# STRUCTURED OUTPUT MODELS
# ============================================================================

# Model for routing decisions (used at graph level)
class ClassificationResult(BaseModel):
    """Classifier's structured output for routing."""
    category: Literal["technical", "billing", "general"]
    confidence: float = Field(ge=0.0, le=1.0)
    reasoning: str


# Model for technical support responses
class TechnicalSupportResponse(BaseModel):
    """Technical agent's structured response."""
    issue_type: Literal["bug", "configuration", "performance", "compatibility", "other"]
    severity: Literal["low", "medium", "high", "critical"]
    troubleshooting_steps: List[str]
    estimated_resolution_time: str
    escalation_needed: bool


# Model for billing support responses  
class BillingSupportResponse(BaseModel):
    """Billing agent's structured response."""
    issue_type: Literal["overcharge", "refund", "subscription", "payment_method", "invoice", "other"]
    amount_involved: Optional[float] = None
    resolution_action: str
    refund_eligible: bool
    follow_up_required: bool


# Model for general support responses
class GeneralSupportResponse(BaseModel):
    """General support agent's structured response."""
    question_type: Literal["hours", "location", "contact", "policy", "other"]
    answer: str
    additional_resources: List[str] = []
    satisfaction_followup: bool


# ============================================================================
# AGENTS WITH THEIR OWN OUTPUT MODELS
# ============================================================================

# Classifier agent - uses ClassificationResult for routing
classifier_agent = Agent(
    name="classifier",
    system_prompt="""You are a customer support classifier. 
Analyze the incoming query and classify it into one of these categories:
- technical: Issues with software, apps, crashes, bugs, performance
- billing: Payment issues, charges, refunds, subscriptions, invoices
- general: Business hours, contact info, policies, general questions

Provide your classification with confidence level and reasoning.""",
)

# Technical agent - has its OWN structured output model
technical_agent = Agent(
    name="technical_support",
    system_prompt="""You are a technical support specialist.
Analyze the technical issue and provide a structured diagnosis including:
- Type of issue (bug, configuration, performance, compatibility, other)
- Severity level
- Step-by-step troubleshooting instructions
- Estimated time to resolve
- Whether escalation to engineering is needed""",
    structured_output_model=TechnicalSupportResponse,
)

# Billing agent - has its OWN structured output model
billing_agent = Agent(
    name="billing_support",
    system_prompt="""You are a billing specialist.
Analyze the billing issue and provide a structured response including:
- Type of billing issue
- Amount involved (if applicable)
- Resolution action to take
- Whether customer is eligible for a refund
- Whether follow-up is required""",
    structured_output_model=BillingSupportResponse,
)

# General agent - has its OWN structured output model
general_agent = Agent(
    name="general_support",
    system_prompt="""You are a general support agent.
Answer general inquiries and provide:
- Type of question being asked
- Clear, helpful answer
- Links to additional resources if relevant
- Whether to follow up on customer satisfaction""",
    structured_output_model=GeneralSupportResponse,
)


# ============================================================================
# CONDITIONAL EDGE FUNCTIONS
# ============================================================================

def route_to_technical(state: GraphState) -> bool:
    """Route to technical agent if classification is 'technical'."""
    classifier_result = state.results.get("classifier")
    if classifier_result and classifier_result.result:
        structured = classifier_result.result.structured_output
        if structured:
            return structured.category == "technical"
    return False


def route_to_billing(state: GraphState) -> bool:
    """Route to billing agent if classification is 'billing'."""
    classifier_result = state.results.get("classifier")
    if classifier_result and classifier_result.result:
        structured = classifier_result.result.structured_output
        if structured:
            return structured.category == "billing"
    return False


def route_to_general(state: GraphState) -> bool:
    """Route to general agent if classification is 'general'."""
    classifier_result = state.results.get("classifier")
    if classifier_result and classifier_result.result:
        structured = classifier_result.result.structured_output
        if structured:
            return structured.category == "general"
    return False


# ============================================================================
# GRAPH CONSTRUCTION
# ============================================================================

def create_support_graph():
    """Create the support routing graph."""
    builder = GraphBuilder()
    
    # Add all nodes
    builder.add_node(classifier_agent, "classifier")
    builder.add_node(technical_agent, "technical")
    builder.add_node(billing_agent, "billing")
    builder.add_node(general_agent, "general")
    
    # Conditional edges from classifier to specialists
    builder.add_edge("classifier", "technical", condition=route_to_technical)
    builder.add_edge("classifier", "billing", condition=route_to_billing)
    builder.add_edge("classifier", "general", condition=route_to_general)
    
    builder.set_entry_point("classifier")
    builder.set_max_node_executions(5)
    
    return builder.build()


# ============================================================================
# EXECUTION EXAMPLE
# ============================================================================

def run_example_graph_and_agent_structured_output():
    """Run example showing both graph-level and agent-level structured output."""
    graph = create_support_graph()
    
    queries = [
        "My application keeps crashing when I try to upload large files over 100MB",
        # "I was charged $49.99 twice for my monthly subscription last week",
        # "What are your support hours on weekends?",
    ]
    
    for query in queries:
        print(f"\n{'='*70}")
        print(f"QUERY: {query}")
        print('='*70)
        
        # Execute graph with ClassificationResult for routing
        result = graph(query, structured_output_model=ClassificationResult)
        
        print(f"\nStatus: {result.status}")
        print(f"Execution path: {' -> '.join(n.node_id for n in result.execution_order)}")
        
        # Show classifier's structured output (used for routing)
        classifier_result = result.results.get("classifier")
        if classifier_result and classifier_result.result:
            classification = classifier_result.result.structured_output
            if classification:
                print(f"\n📋 CLASSIFICATION:")
                print(f"   Category: {classification.category}")
                print(f"   Confidence: {classification.confidence:.0%}")
                print(f"   Reasoning: {classification.reasoning}")
        
        # Show the specialist agent's structured output
        for node_id in ["technical", "billing", "general"]:
            node_result = result.results.get(node_id)
            if node_result and node_result.result:
                structured = node_result.result.structured_output
                if structured:
                    print(f"\n🎯 {node_id.upper()} AGENT RESPONSE:")
                    print(node_result.result.structured_output.model_dump_json())

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!