From d0f27e657746c05a5ce181796d6763845d9efc9c Mon Sep 17 00:00:00 2001
From: Kumaran Rajendhiran <kumaran@airt.ai>
Date: Mon, 30 Dec 2024 18:55:48 +0530
Subject: [PATCH] Add GroupChatRunChat

---
 autogen/agentchat/groupchat.py | 10 ++++++++--
 autogen/messages.py            | 15 +++++++++++++++
 test/test_messages.py          | 24 ++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/autogen/agentchat/groupchat.py b/autogen/agentchat/groupchat.py
index a552418cd..08bcb529d 100644
--- a/autogen/agentchat/groupchat.py
+++ b/autogen/agentchat/groupchat.py
@@ -18,7 +18,12 @@
 from ..formatting_utils import colored
 from ..graph_utils import check_graph_validity, invert_disallowed_to_allowed
 from ..io.base import IOStream
-from ..messages import create_clear_agents_history, create_group_chat_resume, create_speaker_attempt
+from ..messages import (
+    create_clear_agents_history,
+    create_group_chat_resume,
+    create_group_chat_run_chat,
+    create_speaker_attempt,
+)
 from ..oai.client import ModelClient
 from ..runtime_logging import log_new_agent, logging_enabled
 from .agent import Agent
@@ -1149,7 +1154,8 @@ def run_chat(
                 speaker = groupchat.select_speaker(speaker, self)
                 if not silent:
                     iostream = IOStream.get_default()
-                    iostream.print(colored(f"\nNext speaker: {speaker.name}\n", "green"), flush=True)
+                    group_chat_run_chat = create_group_chat_run_chat(speaker, silent)
+                    group_chat_run_chat.print(iostream.print)
                 # let the speaker speak
                 reply = speaker.generate_reply(sender=self)
             except KeyboardInterrupt:
diff --git a/autogen/messages.py b/autogen/messages.py
index dd8e95abc..828f42670 100644
--- a/autogen/messages.py
+++ b/autogen/messages.py
@@ -368,3 +368,18 @@ def create_group_chat_resume(
     last_speaker_name: str, messages: list[dict[str, Any]], silent: Optional[bool] = False
 ) -> GroupChatResume:
     return GroupChatResume(last_speaker_name=last_speaker_name, messages=messages, verbose=not silent)
+
+
+class GroupChatRunChat(BaseModel):
+    speaker_name: str
+    verbose: Optional[bool] = False
+
+    def print(self, f: Optional[Callable[..., Any]] = None) -> None:
+        f = f or print
+
+        if self.verbose:
+            f(colored(f"\nNext speaker: {self.speaker_name}\n", "green"), flush=True)
+
+
+def create_group_chat_run_chat(speaker: Agent, silent: Optional[bool] = False) -> GroupChatRunChat:
+    return GroupChatRunChat(speaker_name=speaker.name, verbose=not silent)
diff --git a/test/test_messages.py b/test/test_messages.py
index 13224c862..53d1b12a7 100644
--- a/test/test_messages.py
+++ b/test/test_messages.py
@@ -16,6 +16,7 @@
     FunctionCallMessage,
     FunctionResponseMessage,
     GroupChatResume,
+    GroupChatRunChat,
     MessageRole,
     PostCarryoverProcessing,
     SpeakerAttempt,
@@ -25,6 +26,7 @@
     ToolResponseMessage,
     create_clear_agents_history,
     create_group_chat_resume,
+    create_group_chat_run_chat,
     create_post_carryover_processing,
     create_received_message_model,
     create_speaker_attempt,
@@ -487,3 +489,25 @@ def test_group_chat_resume() -> None:
     ]
 
     assert mock.call_args_list == expected_call_args_list
+
+
+def test_group_chat_run_chat() -> None:
+    speaker = ConversableAgent(
+        "assistant uno", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER"
+    )
+    silent = False
+
+    actual = create_group_chat_run_chat(speaker, silent)
+
+    assert isinstance(actual, GroupChatRunChat)
+    assert actual.speaker_name == "assistant uno"
+    assert actual.verbose is True
+
+    mock = MagicMock()
+    actual.print(f=mock)
+
+    # print(mock.call_args_list)
+
+    expected_call_args_list = [call("\x1b[32m\nNext speaker: assistant uno\n\x1b[0m", flush=True)]
+
+    assert mock.call_args_list == expected_call_args_list