Make sure swarm demo python is runnable and checked with lint' (#1029)

ag2ai · Feb 19, 2025 · f77865e · f77865e
1 parent b758689
commit f77865e
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 9 deletions.
diff --git 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
@@ -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,26 +243,28 @@ 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,
             condition="The customer is logged in, continue with the order triage.",
             available="logged_in",
         ),
         AFTER_WORK(AfterWorkOption.REVERT_TO_USER),
-    ]
+    ],
 )
 
 
 def has_order_in_context(agent: SwarmAgent, messages: List[Dict[str, Any]]) -> bool:
     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(