From b53ff16eb31bae27e345291bc01c73f02b9057c6 Mon Sep 17 00:00:00 2001
From: Kumaran Rajendhiran <kumaran@airt.ai>
Date: Mon, 30 Dec 2024 20:35:58 +0530
Subject: [PATCH] Add ExecuteCodeBlock

---
 autogen/agentchat/conversable_agent.py | 13 +++++--------
 autogen/messages.py                    | 24 ++++++++++++++++++++++++
 test/test_messages.py                  | 26 ++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/autogen/agentchat/conversable_agent.py b/autogen/agentchat/conversable_agent.py
index 0802d0c07..d69310a49 100644
--- a/autogen/agentchat/conversable_agent.py
+++ b/autogen/agentchat/conversable_agent.py
@@ -35,7 +35,7 @@
 from ..formatting_utils import colored
 from ..function_utils import get_function_schema, load_basemodels_if_needed, serialize_to_str
 from ..io.base import IOStream
-from ..messages import create_received_message_model, create_termination_and_human_reply
+from ..messages import create_execute_code_block, create_received_message_model, create_termination_and_human_reply
 from ..oai.client import ModelClient, OpenAIWrapper
 from ..runtime_logging import log_event, log_function_use, log_new_agent, logging_enabled
 from ..tools import Tool
@@ -2229,13 +2229,10 @@ def execute_code_blocks(self, code_blocks):
             lang, code = code_block
             if not lang:
                 lang = infer_lang(code)
-            iostream.print(
-                colored(
-                    f"\n>>>>>>>> EXECUTING CODE BLOCK {i} (inferred language is {lang})...",
-                    "red",
-                ),
-                flush=True,
-            )
+
+            execute_code_block = create_execute_code_block(code, lang, i, self)
+            execute_code_block.print(iostream.print)
+
             if lang in ["bash", "shell", "sh"]:
                 exitcode, logs, image = self.run_code(code, lang=lang, **self._code_execution_config)
             elif lang in PYTHON_VARIANTS:
diff --git a/autogen/messages.py b/autogen/messages.py
index 1e4092403..0a1b5556d 100644
--- a/autogen/messages.py
+++ b/autogen/messages.py
@@ -413,3 +413,27 @@ def create_termination_and_human_reply(
         sender_name=sender.name if sender else "No sender",
         recipient_name=recipient.name,
     )
+
+
+class ExecuteCodeBlock(BaseModel):
+    code: str
+    language: str
+    code_block_count: int
+    recipient_name: str
+
+    def print(self, f: Optional[Callable[..., Any]] = None) -> None:
+        f = f or print
+
+        f(
+            colored(
+                f"\n>>>>>>>> EXECUTING CODE BLOCK {self.code_block_count} (inferred language is {self.language})...",
+                "red",
+            ),
+            flush=True,
+        )
+
+
+def create_execute_code_block(code: str, language: str, code_block_count: int, recipient: Agent) -> ExecuteCodeBlock:
+    return ExecuteCodeBlock(
+        code=code, language=language, code_block_count=code_block_count, recipient_name=recipient.name
+    )
diff --git a/test/test_messages.py b/test/test_messages.py
index f244b4304..7c37a6f13 100644
--- a/test/test_messages.py
+++ b/test/test_messages.py
@@ -12,6 +12,7 @@
 from autogen.messages import (
     ClearAgentsHistory,
     ContentMessage,
+    ExecuteCodeBlock,
     FunctionCall,
     FunctionCallMessage,
     FunctionResponseMessage,
@@ -26,6 +27,7 @@
     ToolResponse,
     ToolResponseMessage,
     create_clear_agents_history,
+    create_execute_code_block,
     create_group_chat_resume,
     create_group_chat_run_chat,
     create_post_carryover_processing,
@@ -540,3 +542,27 @@ def test_termination_and_human_reply(sender: ConversableAgent, recipient: Conver
     # print(mock.call_args_list)
     expected_call_args_list = [call("\x1b[31m\n>>>>>>>> USING AUTO REPLY...\x1b[0m", flush=True)]
     assert mock.call_args_list == expected_call_args_list
+
+
+def test_execute_code_block(sender: ConversableAgent, recipient: ConversableAgent) -> None:
+    code = """print("hello world")"""
+    language = "python"
+    code_block_count = 0
+
+    actual = create_execute_code_block(code, language, code_block_count, recipient=recipient)
+
+    assert isinstance(actual, ExecuteCodeBlock)
+    assert actual.code == code
+    assert actual.language == language
+    assert actual.code_block_count == code_block_count
+
+    mock = MagicMock()
+    actual.print(f=mock)
+
+    # print(mock.call_args_list)
+
+    expected_call_args_list = [
+        call("\x1b[31m\n>>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\x1b[0m", flush=True)
+    ]
+
+    assert mock.call_args_list == expected_call_args_list