From f77865e11a6992f3636ac459cff3abd0ca97df86 Mon Sep 17 00:00:00 2001
From: Carol Zheng <cazheng@google.com>
Date: Wed, 19 Feb 2025 00:29:37 -0800
Subject: [PATCH] Make sure swarm demo python is runnable and checked with
 lint' (#1029)

---
 website/snippets/python-examples/swarm.mdx    | 10 +++++++++-
 .../python-examples/swarmenhanced.mdx         | 20 +++++++++++--------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/website/snippets/python-examples/swarm.mdx b/website/snippets/python-examples/swarm.mdx
index 7b1a171677..df17163187 100644
--- a/website/snippets/python-examples/swarm.mdx
+++ b/website/snippets/python-examples/swarm.mdx
@@ -1,5 +1,13 @@
 ```python
-from autogen import AFTER_WORK, ON_CONDITION, AfterWorkOption, ConversableAgent, SwarmResult, initiate_swarm_chat
+from autogen import (
+    AFTER_WORK,
+    ON_CONDITION,
+    AfterWorkOption,
+    SwarmAgent,
+    SwarmResult,
+    initiate_swarm_chat,
+    register_hand_off,
+)
 
 llm_config = {"api_type": "openai", "model": "gpt-4o-mini", "cache_seed": None}
 
diff --git a/website/snippets/python-examples/swarmenhanced.mdx b/website/snippets/python-examples/swarmenhanced.mdx
index 2bd6f36121..4c876ccf42 100644
--- a/website/snippets/python-examples/swarmenhanced.mdx
+++ b/website/snippets/python-examples/swarmenhanced.mdx
@@ -11,6 +11,7 @@ from autogen import (
     SwarmResult,
     UserProxyAgent,
     initiate_swarm_chat,
+    register_hand_off,
 )
 
 # Put your key in the OPENAI_API_KEY environment variable
@@ -20,11 +21,9 @@ workflow_context = {
     # customer details
     "customer_name": None,
     "logged_in_username": None,
-
     # workflow status
     "logged_in": False,
     "requires_login": True,
-
     # order enquiry details
     "has_order_id": False,
     "order_id": None,
@@ -71,6 +70,7 @@ ORDER_DATABASE = {
     },
 }
 
+
 # ORDER FUNCTIONS
 def check_order_id(order_id: str, context_variables: dict) -> SwarmResult:
     """Check if the order ID is valid"""
@@ -127,6 +127,7 @@ def login_customer_by_username(username: str, context_variables: dict) -> SwarmR
             agent=authentication_agent,
         )
 
+
 # AGENTS
 
 # Human customer
@@ -228,7 +229,8 @@ nested_chat_two = {
 chat_queue = [nested_chat_one, nested_chat_two]
 
 # HANDOFFS
-order_triage_agent.register_hand_off(
+register_hand_off(
+    order_triage_agent,
     [
         ON_CONDITION(
             target=authentication_agent,
@@ -241,10 +243,11 @@ order_triage_agent.register_hand_off(
             available="logged_in",
         ),
         AFTER_WORK(AfterWorkOption.REVERT_TO_USER),
-    ]
+    ],
 )
 
-authentication_agent.register_hand_off(
+register_hand_off(
+    authentication_agent,
     [
         ON_CONDITION(
             target=order_triage_agent,
@@ -252,7 +255,7 @@ authentication_agent.register_hand_off(
             available="logged_in",
         ),
         AFTER_WORK(AfterWorkOption.REVERT_TO_USER),
-    ]
+    ],
 )
 
 
@@ -260,7 +263,8 @@ def has_order_in_context(agent: SwarmAgent, messages: List[Dict[str, Any]]) -> b
     return agent.get_context("has_order_id")
 
 
-order_mgmt_agent.register_hand_off(
+register_hand_off(
+    order_mgmt_agent,
     [
         ON_CONDITION(
             target={
@@ -276,7 +280,7 @@ order_mgmt_agent.register_hand_off(
         ),
         ON_CONDITION(target=order_triage_agent, condition="The customer has no more enquiries about this order."),
         AFTER_WORK(AfterWorkOption.REVERT_TO_USER),
-    ]
+    ],
 )
 
 chat_history = initiate_swarm_chat(