from autogen import initiate_swarm_chat, ConversableAgent, SwarmResult, OnCondition, AFTER_WORK, AfterWorkOption
from autogen import UserProxyAgent
import os
# All our swarm agents will use GPT-4o-mini from OpenAI
llm_config = LLMConfig.from_json(path="OAI_CONFIG_LIST")
# 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 ConversableAgent, 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)
with llm_config:
customer_service = ConversableAgent(
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.""",
functions=[verify_customer_identity],
)
refund_specialist = ConversableAgent(
name="RefundSpecialist",
system_message="""You are a refund specialist.
Review the case and approve the refund, then transfer to the payment processor.""",
functions=[approve_refund_and_transfer],
)
payment_processor = ConversableAgent(
name="PaymentProcessor",
system_message="""You are a payment processor.
Process the refund payment and provide a confirmation message to the customer.""",
functions=[process_refund_payment],
)
satisfaction_surveyor = ConversableAgent(
name="SatisfactionSurveyor",
system_message="""You are a customer satisfaction specialist.
Ask the customer to rate their experience with the refund process.""",
)
# Conditional and After work hand offs
register_hand_off(
agent=customer_service,
hand_to=[
OnCondition(refund_specialist, "After customer verification, transfer to refund specialist"),
AFTER_WORK(AfterWorkOption.REVERT_TO_USER)
]
)
register_hand_off(
agent=payment_processor,
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)}")