From 4c1bed8bd351b86a4b9a1ae140d834dfdb4a82e1 Mon Sep 17 00:00:00 2001
From: Zac <2215540+zboyles@users.noreply.github.com>
Date: Sat, 19 Oct 2024 23:01:57 -0400
Subject: [PATCH] Add shutdown method to clean up resources for chat agents

---
 .../src/autogen_agentchat/agents/_base_chat_agent.py          | 4 ++++
 .../autogen_agentchat/teams/_group_chat/_base_group_chat.py   | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py
index b5274596874..03b939e6005 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py
@@ -81,6 +81,10 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token:
         """Handle incoming messages and return a response message."""
         ...
 
+    async def shutdown(self):
+        """Shutdown the agent and clean up resources."""
+        pass
+
 
 class BaseToolUseChatAgent(BaseChatAgent):
     """Base class for a chat agent that can use tools.
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
index 633fc0f8524..babd37bce72 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
@@ -169,5 +169,9 @@ async def collect_group_chat_messages(
         # Wait for the runtime to stop.
         await runtime.stop_when_idle()
 
+        # Call shutdown on all participants
+        for participant in self._participants:
+            await participant.shutdown()
+
         # Return the result.
         return TeamRunResult(messages=group_chat_messages)