diff --git a/autogen/agentchat/contrib/swarm_agent.py b/autogen/agentchat/contrib/swarm_agent.py
index 2b0fb9eb78..ca52c017c6 100644
--- a/autogen/agentchat/contrib/swarm_agent.py
+++ b/autogen/agentchat/contrib/swarm_agent.py
@@ -88,9 +88,9 @@ def initiate_swarm_chat(
                 - REVERT_TO_USER : Revert to the user agent if a user agent is provided. If not provided, terminate the conversation.
                 - STAY : Stay with the last speaker.
 
-            Callable: A custom function that takes the current agent, messages, and groupchat as arguments and returns an AfterWorkOption or a SwarmAgent.
+            Callable: A custom function that takes the current agent, messages, and groupchat as arguments and returns an AfterWorkOption or a SwarmAgent (by reference or string name).
                 ```python
-                def custom_afterwork_func(last_speaker: SwarmAgent, messages: List[Dict[str, Any]], groupchat: GroupChat) -> Union[AfterWorkOption, SwarmAgent]:
+                def custom_afterwork_func(last_speaker: SwarmAgent, messages: List[Dict[str, Any]], groupchat: GroupChat) -> Union[AfterWorkOption, SwarmAgent, str]:
                 ```
     Returns:
         ChatResult:     Conversations chat history.
diff --git a/website/docs/topics/swarm.ipynb b/website/docs/topics/swarm.ipynb
index 1724eea982..d848ed2c05 100644
--- a/website/docs/topics/swarm.ipynb
+++ b/website/docs/topics/swarm.ipynb
@@ -137,28 +137,57 @@
     "See the documentation on [registering a nested chat](https://ag2ai.github.io/ag2/docs/reference/agentchat/conversable_agent#register_nested_chats) for further information on the parameters `reply_func_from_nested_chats`, `use_async`, and `config`.\n",
     "\n",
     "Once a nested chat is complete, the resulting output from the last chat in the nested chats will be returned as the agent that triggered the nested chat's response.\n",
-    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
     "### AFTER_WORK\n",
     "\n",
-    "When the last active agent's response doesn't suggest a tool call or handoff, the chat will terminate by default. However, you can register an `AFTER_WORK` handoff to define a fallback agent if you don't want the chat to end at this agent. At the swarm chat level, you also pass in an `AFTER_WORK` handoff to define the fallback mechanism for the entire chat.\n",
-    "If this parameter is set for the agent and the chat, we will prioritize the agent's setting. There should only be one `AFTER_WORK`. If multiple `AFTER_WORK` handoffs are passed, only the last one will be used.\n",
+    "When the active agent's response doesn't suggest a tool call or handoff, the chat will terminate by default. However, you can register an `AFTER_WORK` handoff to control what to do next. You can register these `AFTER_WORK` handoffs at the agent level and also the swarm level (through the `after_work` parameter on `initiate_swarm_chat`). The agent level takes precedence over the swarm level.\n",
+    "\n",
+    "The AFTER_WORK takes a single parameter and this can be an agent, an agent's name, an `AfterWorkOption`, or a callable function.\n",
     "\n",
-    "Besides fallback to an agent, we provide 3 `AfterWorkOption`:     \n",
+    "The `AfterWorkOption` options are:\n",
     "- `TERMINATE`: Terminate the chat    \n",
     "- `STAY`: Stay at the current agent    \n",
-    "- `REVERT_TO_USER`: Revert to the user agent. Only if a user agent is passed in when initializing. (See below for more details)    \n",
+    "- `REVERT_TO_USER`: Revert to the user agent. Only if a user agent is passed in when initializing. (See below for more details)\n",
     "\n",
-    "```python\n",
-    "agent_1 = SwarmAgent(...)\n",
+    "The callable function signature is:\n",
+    "`def my_after_work_func(last_speaker: SwarmAgent, messages: List[Dict[str, Any]], groupchat: GroupChat) -> Union[AfterWorkOption, SwarmAgent, str]:`\n",
     "\n",
-    "# Register the handoff\n",
+    "Note: there should only be one `AFTER_WORK`, if your requirement is more complex, use the callable function parameter.\n",
+    "\n",
+    "Here are examples of registering AFTER_WORKS\n",
+    "\n",
+    "```python\n",
+    "# Register the handoff to an agent\n",
     "agent_1.handoff(hand_to=[\n",
     " ON_CONDITION(...), \n",
     " ON_CONDITION(...),\n",
-    " AFTER_WORK(agent_4) # Fallback to agent_4 if no handoff is suggested\n",
+    " AFTER_WORK(agent_4) # Fallback to agent_4 if no ON_CONDITION handoff is suggested\n",
     "])\n",
     "\n",
+    "# Register the handoff to an AfterWorkOption\n",
     "agent_2.handoff(hand_to=[AFTER_WORK(AfterWorkOption.TERMINATE)]) # Terminate the chat if no handoff is suggested\n",
+    "\n",
+    "def my_after_work_func(last_speaker: SwarmAgent, messages: List[Dict[str, Any]], groupchat: GroupChat) -> Union[AfterWorkOption, SwarmAgent, str]:\n",
+    "    if last_speaker.get_context(\"agent_1_done\"):\n",
+    "        return agent_2\n",
+    "    else:\n",
+    "        return AfterWorkOption.TERMINATE\n",
+    "\n",
+    "# Register the handoff to a function that will return an agent or AfterWorkOption\n",
+    "agent_3.handoff(hand_to=[AFTER_WORK(my_after_work_func)])\n",
+    "\n",
+    "# Register the swarm level AFTER_WORK that becomes the default for agents that don't have one specified\n",
+    "chat_history, context_variables, last_active_agent = initiate_swarm_chat(\n",
+    "    ...\n",
+    "    after_work=AfterWorkOption.TERMINATE # Or an agent or Callable\n",
+    ")\n",
+    "\n",
     "```\n",
     "\n",
     "### Initialize SwarmChat with `initiate_swarm_chat`\n",