Teams & Workflows

import asyncio
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.conditions import MaxMessageTermination

# Create agents with different roles
writer_agent = AssistantAgent(
    name="Writer",
    model_client=model_client,
    description="""You are an expert content writer specializing in technical documentation and blog posts.
    Your role is to create engaging, well-structured content that explains complex topics clearly.
    Focus on creating compelling introductions, logical flow, and actionable insights.
    Always consider your target audience and write in an accessible yet professional tone."""
)

editor_agent = AssistantAgent(
    name="Editor",
    model_client=model_client,
    description="""You are a professional editor and proofreader with expertise in technical content.
    Your role is to refine and improve written content for clarity, grammar, style, and structure.
    Check for consistency in terminology, improve readability, fix grammatical errors, and ensure
    the content flows logically. Provide specific suggestions for improvement while maintaining
    the author's voice and intent."""
)

reviewer_agent = AssistantAgent(
    name="Reviewer",
    model_client=model_client,
    description="""You are a content quality reviewer with expertise in technical accuracy and audience engagement.
    Your role is to evaluate content for technical correctness, completeness, and effectiveness.
    Assess whether the content meets its objectives, is appropriate for the target audience,
    and provides valuable insights. Identify gaps, suggest improvements, and ensure the content
    aligns with best practices and industry standards."""
)

# Create RoundRobin team
team = RoundRobinGroupChat(
    participants=[writer_agent, editor_agent, reviewer_agent],
    termination_condition=MaxMessageTermination(max_messages=9)
)

# Run the team
result = await team.run(
    task="Write and refine a blog post about AutoGen teams"
)

# Reflection pattern with RoundRobin
async def reflection_workflow():
    # Create specialized agents
    analyst = AssistantAgent(
        name="Analyst",
        model_client=model_client,
        description="""You are a systems analyst with expertise in algorithm design and problem-solving.
        Your role is to break down complex problems, analyze requirements, and propose initial solutions.
        Focus on understanding the core problem, identifying key constraints, and designing efficient
        approaches. Consider scalability, performance, and maintainability in your proposals.
        Present your analysis clearly with justification for your design choices."""
    )

    critic = AssistantAgent(
        name="Critic",
        model_client=model_client,
        description="""You are a senior technical critic with expertise in identifying weaknesses and edge cases.
        Your role is to rigorously evaluate proposed solutions and identify potential issues.
        Look for scalability problems, edge cases, security vulnerabilities, performance bottlenecks,
        and maintainability concerns. Provide constructive criticism with specific examples and
        suggest areas that need improvement. Be thorough but fair in your evaluation."""
    )

    refiner = AssistantAgent(
        name="Refiner",
        model_client=model_client,
        description="""You are a solution architect specializing in optimization and refinement.
        Your role is to take feedback from analysis and criticism to improve and polish solutions.
        Address identified issues, optimize performance, enhance robustness, and ensure best practices.
        Integrate feedback constructively while maintaining the solution's core effectiveness.
        Focus on creating production-ready, well-documented solutions."""
    )

    # Each round: Analyze → Critique → Refine
    reflection_team = RoundRobinGroupChat(
        participants=[analyst, critic, refiner],
        termination_condition=MaxMessageTermination(max_messages=12)
    )

    result = await reflection_team.run(
        task="Design an efficient algorithm for large-scale data processing"
    )
    return result

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import SelectorGroupChat
from autogen_agentchat.conditions import TextMentionTermination

# Create specialized research agents
web_searcher = AssistantAgent(
    name="WebSearcher",
    model_client=model_client,
    description="""You are a research specialist with expertise in web search and information gathering.
    Your role is to find current, accurate, and relevant information from reliable sources on the internet.
    Focus on identifying authoritative sources, fact-checking information, and gathering comprehensive data.
    Organize your findings clearly and provide source citations. Prioritize recent information and
    reputable publications. Signal when you have sufficient data for analysis.""",
    tools=[web_search_tool]
)

data_analyst = AssistantAgent(
    name="DataAnalyst",
    model_client=model_client,
    description="""You are a data analyst with expertise in statistical analysis and pattern recognition.
    Your role is to process, analyze, and interpret data to extract meaningful insights and trends.
    Focus on identifying patterns, calculating relevant statistics, creating data visualizations,
    and drawing evidence-based conclusions. Present your analysis clearly with supporting data
    and explain the significance of your findings. Highlight key trends and actionable insights."""
)

report_writer = AssistantAgent(
    name="ReportWriter",
    model_client=model_client,
    description="""You are a professional report writer specializing in research documentation and executive summaries.
    Your role is to synthesize research findings and analysis into comprehensive, well-structured reports.
    Focus on creating clear executive summaries, organizing information logically, and presenting
    insights in an actionable format. Ensure your reports are suitable for decision-makers and
    include recommendations based on the research and analysis provided."""
)

# Create SelectorGroupChat team
research_team = SelectorGroupChat(
    participants=[web_searcher, data_analyst, report_writer],
    model_client=model_client,
    termination_condition=TextMentionTermination("FINAL_REPORT")
)

# The selector will choose appropriate agents based on task needs
result = await research_team.run(
    task="Research current trends in renewable energy adoption and create a comprehensive report"
)

# Custom speaker selection based on conversation history
def custom_speaker_selection(messages):
    last_message = messages[-1]

    # If last message mentioned data, choose analyst
    if any(keyword in last_message.content.lower()
           for keyword in ["data", "statistics", "numbers", "analysis"]):
        return "DataAnalyst"

    # If last message was a question, choose searcher
    if "?" in last_message.content:
        return "WebSearcher"

    # If analysis is complete, choose writer
    if "analysis complete" in last_message.content.lower():
        return "ReportWriter"

    # Default to searcher for initial research
    return "WebSearcher"

# Use custom selection function
custom_team = SelectorGroupChat(
    participants=[web_searcher, data_analyst, report_writer],
    model_client=model_client,
    speaker_selection_method=custom_speaker_selection,
    termination_condition=TextMentionTermination("FINAL_REPORT")
)

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import Swarm
from autogen_agentchat.messages import HandoffMessage

# Customer support swarm agents
frontline_agent = AssistantAgent(
    name="FrontlineSupport",
    model_client=model_client,
    description="""You are a frontline customer support representative with expertise in initial problem assessment.
    Your role is to greet customers warmly, understand their issues, and provide immediate help when possible.
    For technical problems (errors, bugs, performance issues), use HandoffMessage to transfer to TechnicalSupport.
    For billing, payment, or account issues, use HandoffMessage to transfer to BillingSupport.
    For general questions, provide helpful information and escalate only when necessary.
    Always maintain a friendly, professional tone and ensure customers feel heard."""
)

technical_agent = AssistantAgent(
    name="TechnicalSupport",
    model_client=model_client,
    description="""You are a technical support specialist with expertise in troubleshooting and problem resolution.
    Your role is to diagnose technical issues, provide step-by-step solutions, and resolve technical problems.
    For complex issues requiring advanced expertise (system architecture, deep debugging, critical bugs),
    use HandoffMessage to transfer to SpecialistSupport with a detailed summary of the problem and
    troubleshooting steps already attempted. Focus on clear technical communication and systematic problem-solving."""
)

billing_agent = AssistantAgent(
    name="BillingSupport",
    model_client=model_client,
    description="""You are a billing and account specialist with expertise in financial transactions and account management.
    Your role is to handle billing inquiries, payment issues, subscription changes, and account-related questions.
    Provide clear explanations of charges, help with payment methods, process refunds when appropriate,
    and resolve account access issues. Ensure customers understand their billing and maintain accurate records.
    Escalate complex billing disputes to management when necessary."""
)

specialist_agent = AssistantAgent(
    name="SpecialistSupport",
    model_client=model_client,
    description="""You are a senior technical specialist with deep expertise in complex system issues and escalations.
    Your role is to handle the most challenging technical problems that require advanced knowledge.
    Focus on root cause analysis, system-level debugging, performance optimization, and critical issue resolution.
    Provide detailed technical solutions and work directly with engineering teams when needed.
    Document complex solutions for future reference and mentor other support team members."""
)

# Create swarm with handoff logic
support_swarm = Swarm(
    agents=[frontline_agent, technical_agent, billing_agent, specialist_agent],
    starting_agent=frontline_agent
)

# Agents use HandoffMessage to transfer control
async def handle_customer_inquiry(inquiry):
    result = await support_swarm.run(
        task=f"Customer inquiry: {inquiry}"
    )
    return result

# Financial research swarm example
data_collector = AssistantAgent(
    name="DataCollector",
    model_client=model_client,
    description="""You are a financial data specialist with expertise in market research and data gathering.
    Your role is to collect comprehensive financial data, market information, news, and relevant metrics.
    Use available tools to gather stock prices, financial statements, market trends, news sentiment,
    and economic indicators. Organize data systematically and ensure accuracy and completeness.
    When data collection is complete, use HandoffMessage to transfer to FinancialAnalyst with
    a summary of collected data and sources.""",
    tools=[stock_data_tool, news_search_tool]
)

financial_analyst = AssistantAgent(
    name="FinancialAnalyst",
    model_client=model_client,
    description="""You are a senior financial analyst with expertise in investment analysis and valuation.
    Your role is to analyze financial data, calculate key metrics, identify trends, and develop investment insights.
    Focus on fundamental analysis, technical indicators, financial ratios, and market comparisons.
    Provide clear analysis with supporting calculations and reasoning. When analysis is complete,
    use HandoffMessage to transfer to RiskAssessment with your findings and preliminary investment thesis."""
)

risk_assessor = AssistantAgent(
    name="RiskAssessment",
    model_client=model_client,
    description="""You are a risk management specialist with expertise in investment risk evaluation and compliance.
    Your role is to assess potential risks, evaluate risk-return profiles, and ensure regulatory compliance.
    Analyze market risks, credit risks, liquidity risks, and operational risks. Consider macroeconomic factors,
    volatility, correlation, and scenario analysis. When risk assessment is complete,
    use HandoffMessage to transfer to ReportWriter with risk evaluation and recommended risk parameters."""
)

report_writer = AssistantAgent(
    name="ReportWriter",
    model_client=model_client,
    description="""You are an investment report specialist with expertise in creating comprehensive financial reports.
    Your role is to synthesize data analysis and risk assessment into professional investment reports.
    Create clear executive summaries, detailed analysis sections, and actionable investment recommendations.
    Ensure reports meet institutional standards and provide clear rationale for recommendations.
    Include risk disclosures, assumptions, and supporting evidence for all conclusions."""
)

# Create financial research swarm
finance_swarm = Swarm(
    agents=[data_collector, financial_analyst, risk_assessor, report_writer],
    starting_agent=data_collector
)

# Example handoff in agent logic:
# return HandoffMessage(
#     target="FinancialAnalyst",
#     content="Data collection complete. Here's the financial data for analysis..."
# )

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import GraphFlow
from autogen_agentchat.conditions import TextMentionTermination
import networkx as nx

# Create workflow agents
planner = AssistantAgent(
    name="Planner",
    model_client=model_client,
    description="""You are a project planning specialist with expertise in software development lifecycle management.
    Your role is to analyze requirements, break down complex projects into manageable tasks, and create detailed
    implementation plans. Consider technical constraints, dependencies, timelines, and resource requirements.
    Provide clear specifications, acceptance criteria, and technical guidance for the development team.
    Ensure plans are realistic, well-structured, and include proper testing and review phases."""
)

developer = AssistantAgent(
    name="Developer",
    model_client=model_client,
    description="""You are a software developer with expertise in implementing robust, maintainable code solutions.
    Your role is to write high-quality code that meets specifications and follows best practices.
    Focus on clean code principles, proper error handling, security considerations, and performance optimization.
    Document your code clearly and ensure it's testable and maintainable. Communicate any issues or
    clarifications needed during implementation."""
)

tester = AssistantAgent(
    name="Tester",
    model_client=model_client,
    description="""You are a quality assurance specialist with expertise in comprehensive software testing.
    Your role is to design and execute test cases, identify bugs, and verify system functionality.
    Focus on functional testing, edge cases, error handling, performance testing, and user experience.
    Report issues clearly with reproduction steps, expected vs actual behavior, and severity levels.
    Ensure all requirements are properly tested before approval."""
)

reviewer = AssistantAgent(
    name="Reviewer",
    model_client=model_client,
    description="""You are a senior code reviewer with expertise in code quality, architecture, and best practices.
    Your role is to review code for maintainability, security, performance, and adherence to standards.
    Focus on code structure, design patterns, potential vulnerabilities, and long-term maintainability.
    Provide constructive feedback with specific suggestions for improvement. Ensure code meets
    team standards and industry best practices before approval."""
)

# Build workflow graph
workflow_graph = nx.DiGraph()

# Add nodes (agents)
workflow_graph.add_node("Planner", agent=planner)
workflow_graph.add_node("Developer", agent=developer)
workflow_graph.add_node("Tester", agent=tester)
workflow_graph.add_node("Reviewer", agent=reviewer)

# Add edges (workflow connections)
workflow_graph.add_edge("Planner", "Developer")
workflow_graph.add_edge("Developer", "Tester")
workflow_graph.add_edge("Developer", "Reviewer")  # Parallel review
workflow_graph.add_edge("Tester", "Developer")    # Bug fixes
workflow_graph.add_edge("Reviewer", "Developer")  # Code improvements

# Advanced GraphFlow with conditional routing
def route_based_on_content(source, messages):
    """Custom routing function based on message content"""
    last_message = messages[-1]

    if source == "Planner":
        return ["Developer"]  # Always go to developer after planning

    elif source == "Developer":
        # Route to both tester and reviewer in parallel
        return ["Tester", "Reviewer"]

    elif source == "Tester":
        if "PASS" in last_message.content:
            return []  # End workflow - tests passed
        else:
            return ["Developer"]  # Back to developer for fixes

    elif source == "Reviewer":
        if "APPROVED" in last_message.content:
            return []  # End workflow - review approved
        else:
            return ["Developer"]  # Back to developer for improvements

    return []

# Create GraphFlow with custom routing
development_workflow = GraphFlow(
    graph=workflow_graph,
    starting_node="Planner",
    routing_func=route_based_on_content,
    termination_condition=TextMentionTermination("WORKFLOW_COMPLETE")
)

# Run the workflow
result = await development_workflow.run(
    task="Develop a user authentication system with comprehensive testing"
)

# Message filtering for focused communication
def filter_messages_for_agent(agent_name, messages):
    """Filter messages relevant to specific agents"""

    if agent_name == "Tester":
        # Only send implementation and bug-related messages
        return [msg for msg in messages
                if any(keyword in msg.content.lower()
                       for keyword in ["implement", "code", "test", "bug", "error"])]

    elif agent_name == "Reviewer":
        # Only send code quality and review-related messages
        return [msg for msg in messages
                if any(keyword in msg.content.lower()
                       for keyword in ["code", "review", "quality", "standard", "best practice"])]

    # Default: send all messages
    return messages

# Apply message filtering in GraphFlow
filtered_workflow = GraphFlow(
    graph=workflow_graph,
    starting_node="Planner",
    message_filter_func=filter_messages_for_agent,
    termination_condition=TextMentionTermination("WORKFLOW_COMPLETE")
)