Enter fullscreen mode

Building Swarm-based agents with AG2

Authors: Yiran Wu, Mark Sze

AG2 now provides an implementation of the swarm orchestration from OpenAI's Swarm framework, with some additional features!

Background: the swarm orchestration is a multi-agent collaboration where agents execute tasks and are responsible for handing them off to other agents.

Here are the key features of the swarm orchestration:

Hand-offs: Agents can transfer control to another agent, enabling smooth and direct transitions within workflows.
Context variables: Agents can dynamically update a shared memory through function calls, maintaining context and adaptability throughout the process.

Besides these core features, AG2 provides:

Simple Interface for Tool Call/Handoff Registration: When creating a swarm agent, you can pass in a list of functions to be registered directly. We also provide a simple method to register handoffs.
Transition Beyond Tool Calls: We enable the ability to automatically transfer to a nominated swarm agent when an agent has completed their task. We will extend this to allow other transitions in the future (e.g., use a function to determine the next agent ).
Built-in human-in-the-loop: Adding a user agent (UserProxyAgent) to your swarm will allow swarm agents to transition back to the user. Provides a means to clarify and confirm with the user without breaking out of the swarm.

This feature builds on GroupChat, offering a simpler interface to use swarm orchestration. For comparison, see two implementations of the same example: one using swarm orchestration and another naive implementation with GroupChat (Legacy).

Handoffs

Before we dive into a swarm example, an important concept in swarm orchestration is when and how an agent hands off to another agent.

Providing additional flexibility, we introduce the capability to define an after-work handoff. Think of it as the agent's next action after completing their task. It can be to hand off to another agent, revert to the user, stay with the agent for another iteration, or terminate the conversation.

The following are the prioritized handoffs for each iteration of the swarm.

Agent-level: Calls a tool that returns a swarm agent: A swarm agent's tool call returns the next agent to hand off to.
Agent-level: Calls a pre-defined conditional handoff: A swarm agent has an ON_CONDITION handoff that is chosen by the LLM (behaves like a tool call).
Agent-level: After work hand off: When no tool calls are made it can use an, optional, AFTER_WORK handoff that is a preset option or a nominated swarm agent.
Swarm-level: After work handoff: If the agent does not have an AFTER_WORK handoff, the swarm's AFTER_WORK handoff will be used.

In the following code sample a SwarmAgent named responder has:

Two conditional handoffs registered (ON_CONDITION), specifying the agent to hand off to and the condition to trigger the handoff.
An after-work handoff (AFTER_WORK) nominated using one of the preset options (TERMINATE, REVERT_TO_USER, STAY). This could also be a swarm agent.

responder.register_hand_off(
    hand_to=[
 ON_CONDITION(weather, "If you need weather data, hand off to the Weather_Agent"),
 ON_CONDITION(travel_advisor, "If you have weather data but need formatted recommendations, hand off to the Travel_Advisor_Agent"),
 AFTER_WORK(AfterWorkOption.REVERT_TO_USER),
 ]
)

You can specify the swarm-level after work handoff when initiating the swarm (here we nominate to terminate):

history, context, last_agent = initiate_swarm_chat(
    initial_agent=responder,
    agents=my_list_of_swarm_agents,
    max_rounds=30,
    messages=messages,
    after_work=AFTER_WORK(AfterWorkOption.TERMINATE)
)

Creating a swarm

Define the functions that can be used by your SwarmAgents.
Create your SwarmAgents (which derives from ConversableAgent).
For each swarm agent, specify the handoffs (transitions to another agent) and what to do when they have finished their work (termed After Work).
Optionally, create your context dictionary.
Call initiate_swarm_chat.

Example

This example of managing refunds demonstrates the context handling, swarm and agent-level conditional and after work hand offs, and the human-in-the-loop feature.

from autogen import initiate_swarm_chat, SwarmAgent, SwarmResult, ON_CONDITION, AFTER_WORK, AfterWorkOption
from autogen import UserProxyAgent
import os

# All our swarm agents will use GPT-4o-mini from OpenAI
llm_config = {...}

# We configure our starting context dictionary
context_variables = {
    "passport_number": "",
    "customer_verified": False,
    "refund_approved": False,
    "payment_processed": False
}

# Functions that our swarm agents will be assigned
# They can return a SwarmResult, a SwarmAgent, or a string
# SwarmResult allows you to update context_variables and/or hand off to another agent
def verify_customer_identity(passport_number: str, context_variables: dict) -> str:
    context_variables["passport_number"] = passport_number
    context_variables["customer_verified"] = True
    return SwarmResult(values="Customer identity verified", context_variables=context_variables)

def approve_refund_and_transfer(context_variables: dict) -> str:
    context_variables["refund_approved"] = True
    return SwarmResult(values="Refund approved", context_variables=context_variables, agent=payment_processor)

def process_refund_payment(context_variables: dict) -> str:
    context_variables["payment_processed"] = True
    return SwarmResult(values="Payment processed successfully", context_variables=context_variables)

# Swarm Agents, similar to ConversableAgent, but with functions and hand offs (specified later)
customer_service = SwarmAgent(
    name="CustomerServiceRep",
    system_message="""You are a customer service representative.
      First verify the customer's identity by asking for the customer's passport number,
      then calling the verify_customer_identity function,
      finally, transfer the case to the refund specialist.""",
    llm_config=llm_config,
    functions=[verify_customer_identity],
)

refund_specialist = SwarmAgent(
    name="RefundSpecialist",
    system_message="""You are a refund specialist.
    Review the case and approve the refund, then transfer to the payment processor.""",
    llm_config=llm_config,
    functions=[approve_refund_and_transfer],
)

payment_processor = SwarmAgent(
    name="PaymentProcessor",
    system_message="""You are a payment processor.
    Process the refund payment and provide a confirmation message to the customer.""",
    llm_config=llm_config,
    functions=[process_refund_payment],
)

satisfaction_surveyor = SwarmAgent(
    name="SatisfactionSurveyor",
    system_message="""You are a customer satisfaction specialist.
    Ask the customer to rate their experience with the refund process.""",
    llm_config=llm_config,
)

# Conditional and After work hand offs

customer_service.register_hand_off(
    hand_to=[
        ON_CONDITION(refund_specialist, "After customer verification, transfer to refund specialist"),
        AFTER_WORK(AfterWorkOption.REVERT_TO_USER)
    ]
)

payment_processor.register_hand_off(
    hand_to=[
        AFTER_WORK(satisfaction_surveyor),
    ]
)

# Our human, you, allowing swarm agents to revert back for more information
user = UserProxyAgent(name="User", code_execution_config=False)

# Initiate the swarm
# Returns the ChatResult, final context, and last speaker

chat_result, context_variables, last_speaker = initiate_swarm_chat(
    initial_agent=customer_service, # Starting agent
    agents=[customer_service, refund_specialist, payment_processor, satisfaction_surveyor],
    user_agent=user, # Human user
    messages="Customer requesting refund for order #12345",
    context_variables=context_variables, # Context
    after_work=AFTER_WORK(AfterWorkOption.TERMINATE) # Swarm-level after work hand off
)


print(f"Context Variables:\n{json.dumps(context_variables, indent=2)}")

And here's the output, showing the flow of agents.

User (to chat_manager):

Customer requesting refund for order #12345

--------------------------------------------------------------------------------

Next speaker: CustomerServiceRep

CustomerServiceRep (to chat_manager):

To assist you with the refund for order #12345,
I need to verify your identity.
Could you please provide your passport number?

--------------------------------------------------------------------------------

Next speaker: User

Replying as User. Provide feedback to chat_manager.
Press enter to skip and use auto-reply,
or type 'exit' to end the conversation: AUS923828C

User (to chat_manager):

AUS923828C

--------------------------------------------------------------------------------

Next speaker: CustomerServiceRep

CustomerServiceRep (to chat_manager):

***** Suggested tool call (call_wfx2VoCmuCKDFKV5xcmB6kSc): verify_customer_identity *****
Arguments:
{"passport_number":"AUS923828C"}
*****************************************************************************************

--------------------------------------------------------------------------------

Next speaker: Tool_Execution


>>>>>>>> EXECUTING FUNCTION verify_customer_identity...
Tool_Execution (to chat_manager):

***** Response from calling tool (call_wfx2VoCmuCKDFKV5xcmB6kSc) *****
Customer identity verified
**********************************************************************

--------------------------------------------------------------------------------

Next speaker: CustomerServiceRep

CustomerServiceRep (to chat_manager):

***** Suggested tool call (call_Jz1viRLeJuOltPRcKfYZ8bgH): transfer_to_RefundSpecialist *****
Arguments:
{}
*********************************************************************************************

--------------------------------------------------------------------------------

Next speaker: Tool_Execution


>>>>>>>> EXECUTING FUNCTION transfer_to_RefundSpecialist...
Tool_Execution (to chat_manager):

***** Response from calling tool (call_Jz1viRLeJuOltPRcKfYZ8bgH) *****
SwarmAgent --> RefundSpecialist
**********************************************************************

--------------------------------------------------------------------------------

Next speaker: RefundSpecialist

RefundSpecialist (to chat_manager):

***** Suggested tool call (call_c4uhy8Mi3Ihe3tVRFVG9M8Uw): approve_refund_and_transfer *****
Arguments:
{}
********************************************************************************************

--------------------------------------------------------------------------------

Next speaker: Tool_Execution


>>>>>>>> EXECUTING FUNCTION approve_refund_and_transfer...
Tool_Execution (to chat_manager):

***** Response from calling tool (call_c4uhy8Mi3Ihe3tVRFVG9M8Uw) *****
Refund approved
**********************************************************************

--------------------------------------------------------------------------------

Next speaker: PaymentProcessor

PaymentProcessor (to chat_manager):

***** Suggested tool call (call_qrs5ysx89rQqMFtVuqMUb9O0): process_refund_payment *****
Arguments:
{}
***************************************************************************************

--------------------------------------------------------------------------------

Next speaker: Tool_Execution


>>>>>>>> EXECUTING FUNCTION process_refund_payment...
Tool_Execution (to chat_manager):

***** Response from calling tool (call_qrs5ysx89rQqMFtVuqMUb9O0) *****
Payment processed successfully
**********************************************************************

--------------------------------------------------------------------------------

Next speaker: PaymentProcessor

PaymentProcessor (to chat_manager):

Your refund for order #12345 has been processed successfully.
You should see the amount credited back to your account shortly.
If you have any further questions or need assistance, feel free to reach out.
Thank you for your patience!

--------------------------------------------------------------------------------

Next speaker: SatisfactionSurveyor

SatisfactionSurveyor (to chat_manager):

Thank you for your patience during the refund process!
We hope you are satisfied with your experience.
Could you please rate your satisfaction with the refund process from 1 to 5,
where 1 is very dissatisfied and 5 is very satisfied?
Your feedback is valuable to us!

--------------------------------------------------------------------------------
Context Variables:
{
  "customer_verified": true,
  "refund_approved": true,
  "payment_processed": true,
  "passport_number": "AUS923828C"
}

Notes

Behind-the-scenes, swarm agents are supported by a tool execution agent, that executes tools on their behalf. Hence, the appearance of Tool Execution in the output.
Currently only swarm agents can be added to a swarm. This is to maintain their ability to manage context variables, auto-execute functions, and support hand offs. Eventually, we may allow ConversableAgent to have the same capability and make "SwarmAgent" a simpler subclass with certain defaults changed (like AssistantAgent and UserProxyAgent).
Would you like to enhance the swarm feature or have found a bug? Please let us know by creating an issue on the AG2 GitHub.

For Further Reading

Finding this useful?

The AG2 team is working hard to create content like this, not to mention building a powerful, open-source, end-to-end platform for multi-agent automation.

The easiest way to show your support is just to star AG2 repo, but also take a look at it for contributions or simply to give it a try.

Also, let us know if you have any interesting use cases for Swarm-based agents? Or maybe you would like to see more features or improvements? Do join our Discord server for discussion.