Swarm Concept Code

import autogen
import random

from autogen import (
 AfterWork,
 OnCondition,
 AfterWorkOption,
 AssistantAgent,
 SwarmResult,
 initiate_swarm_chat,
 register_hand_off
)

config_list = autogen.config_list_from_json("<path_to_your_config_file>")
llm_config = {"config_list": config_list}

# 1. A function that returns a value of "success" and updates the context variable "1" to True
def update_context_1(context_variables: dict) -> SwarmResult:
    context_variables["1"] = True
    return SwarmResult(value="success", context_variables=context_variables)

# 2. A function that returns an AssistantAgent object
def transfer_to_agent_2() -> AssistantAgent:
    """Transfer to agent 2"""
    return agent_2


# 3. A function that returns the value of "success", updates the context variable and transfers to agent 3
def update_context_2_and_transfer_to_3(context_variables: dict) -> SwarmResult:
    context_variables["2"] = True
    return SwarmResult(value="success", context_variables=context_variables, agent=agent_3)


# 4. A function that returns a normal value
def get_random_number() -> str:
    return random.randint(1, 100)


def update_context_3_with_random_number(context_variables: dict, random_number: int) -> SwarmResult:
 context_variables["3"] = random_number
    return SwarmResult(value="success", context_variables=context_variables)


agent_1 = AssistantAgent(
    name="Agent_1",
    system_message="You are Agent 1, first, call the function to update context 1, and transfer to Agent 2",
    llm_config=llm_config,
    functions=[update_context_1, transfer_to_agent_2],
)

agent_2 = AssistantAgent(
    name="Agent_2",
    system_message="You are Agent 2, call the function that updates context 2 and transfer to Agent 3",
    llm_config=llm_config,
    functions=[update_context_2_and_transfer_to_3],
)

agent_3 = AssistantAgent(
    name="Agent_3",
    system_message="You are Agent 3, tell a joke",
    llm_config=llm_config,
)

agent_4 = AssistantAgent(
    name="Agent_4",
    system_message="You are Agent 4, call the function to get a random number",
    llm_config=llm_config,
    functions=[get_random_number],
)

agent_5 = AssistantAgent(
    name="Agent_5",
    system_message="Update context 3 with the random number.",
    llm_config=llm_config,
    functions=[update_context_3_with_random_number],
)


# 5. This is equivalent to writing a transfer function
register_hand_off(agent=agent_3,hand_to=[OnCondition(agent_4, "Transfer to Agent 4")])
register_hand_off(agent=agent_4,hand_to=[AfterWork(agent_5)])

print("Agent 1 function schema:")
for func_schema in agent_1.llm_config["tools"]:
    print(func_schema)

print("Agent 3 function schema:")
for func_schema in agent_3.llm_config["tools"]:
    print(func_schema)

context_variables = {"1": False, "2": False, "3": False}
chat_result, context_variables, last_agent = initiate_swarm_chat(
    initial_agent=agent_1,
    agents=[agent_1, agent_2, agent_3, agent_4, agent_5],
    messages="start",
    context_variables=context_variables,
    after_work=AfterWork(AfterWorkOption.TERMINATE),  # this is the default value
)

from autogen import UserProxyAgent

user_agent = UserProxyAgent(name="User", code_execution_config=False)

agent_6 = AssistantAgent(
    name="Agent_6",
    system_message="You are Agent 6. Your job is to tell jokes.",
    llm_config=llm_config,
)

agent_7 = AssistantAgent(
    name="Agent_7",
    system_message="You are Agent 7, explain the joke.",
    llm_config=llm_config,
)

register_hand_off(
    agent=agent_6,
    hand_to=[
        OnCondition(
        agent_7, "Used to transfer to Agent 7. Don't call this function, unless the user explicitly tells you to."
        ),
        AfterWork(AfterWorkOption.REVERT_TO_USER),
    ]
)

chat_result, _, _ = initiate_swarm_chat(
    initial_agent=agent_6,
    agents=[agent_6, agent_7],
    user_agent=user_agent,
    messages="start",
)