Repository: microsoft/autogen
Pull request: Add code generation support to `CodeExecutorAgent`
## Why are these changes needed?
- To add support for code generation, execution and reflection to `CodeExecutorAgent`.

## Related issue number
Closes #5824 

## Checks
- [x] I've included any doc changes needed for <https://microsoft.github.io/autogen/>. See <https://github.com/microsoft/autogen/blob/main/CONTRIBUTING.md> to build and test documentation locally.
- [x] I've added tests (if relevant) corresponding to the changes introduced in this PR.
- [x] I've made sure all auto checks have passed.


Problem statement:
Enable `CodeExecutorAgent` to generate its own code and executing it.
### Confirmation

- [x] I confirm that I am a maintainer and so can use this template. If I am not, I understand this issue will be closed and I will be asked to use a different template.

### Issue body

Adding optional parameters: `system_message` and `model_client` to `CodeExecutorAgent` so it can generate and execute its own code -- achieving similar behavior as OpenAI Assistant Agent with code interpreter.

Hints from issue discussion:
Hi @ekzhu,  

I experimented with a baseline approach where I inherited from `AssistantAgent` instead of `BaseChatAgent`. I renamed `CodeExecutorAgent.on_messages` to `CodeExecutorAgent.execute_code_block` and added the following method:  

```python
# A very raw implementation
async def on_messages_with_code_exec(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
    response = await self.on_messages(messages, cancellation_token)
    result = await self.execute_code_block([response.chat_message], cancellation_token)
    agent_response = f"""
    Agent Message
    {response.chat_message.content}

    Tool Result
    {result.chat_message.content}
    """
    return Response(
        chat_message=TextMessage(content=agent_response, source=self.name)
    )
```

I had a few questions regarding this approach:
1. Is inheriting from `AssistantAgent` a good way to implement this feature?
2. For the official implementation, I'm considering renaming `on_messages` to `execute_code_block` and then overriding `AssistantAgent.on_messages` to first generate an LLM response, followed by executing the code using `CodeExecutorAgent.execute_code_block`. Does this seem like a good approach?
3. After generating and executing the code, would it be better to keep both as separate `TextMessage` instances, or should they be combined into a single `TextMessage` as shown above?

Looking forward to your thoughts!
Thanks for looking into this issue. As the first step, we can implement a reasoning --> execution --> reflection framework for the code executor agent. It has a model context and model client just like `AssistantAgent`, and it stores the incoming messages into model context, makes an inference call with the model client, and then if there are any code blocks, it executes them; lastly, it reflects on the output of the code block using the model client.

Do you want to start with a draft PR so we can iterate on that?

Let's not subclass `AssistantAgent`, which is quite complex, and we may need another refactor for it with new tasks like #5611 and #5931.

Commit message: Add code generation support to `CodeExecutorAgent` (#5824)

Signed-off-by: Abhijeetsingh Meena <abhijeet040403@gmail.com>

Patch to review:
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py
index 089daa2a15f0..afaf526f291e 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py
@@ -1,23 +1,59 @@
+import logging
 import re
-from typing import List, Sequence
+from typing import (
+    AsyncGenerator,
+    List,
+    Optional,
+    Sequence,
+    Union,
+)
 
 from autogen_core import CancellationToken, Component, ComponentModel
 from autogen_core.code_executor import CodeBlock, CodeExecutor
+from autogen_core.memory import Memory
+from autogen_core.model_context import (
+    ChatCompletionContext,
+    UnboundedChatCompletionContext,
+)
+from autogen_core.models import (
+    AssistantMessage,
+    ChatCompletionClient,
+    CreateResult,
+    LLMMessage,
+    SystemMessage,
+    UserMessage,
+)
 from pydantic import BaseModel
 from typing_extensions import Self
 
+from .. import EVENT_LOGGER_NAME
 from ..base import Response
-from ..messages import ChatMessage, TextMessage
+from ..messages import (
+    AgentEvent,
+    ChatMessage,
+    HandoffMessage,
+    MemoryQueryEvent,
+    ModelClientStreamingChunkEvent,
+    TextMessage,
+    ThoughtEvent,
+)
+from ..utils import remove_images
 from ._base_chat_agent import BaseChatAgent
 
+event_logger = logging.getLogger(EVENT_LOGGER_NAME)
+
 
 class CodeExecutorAgentConfig(BaseModel):
     """Configuration for CodeExecutorAgent"""
 
     name: str
     code_executor: ComponentModel
+    model_client: ComponentModel
     description: str = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks)."
     sources: List[str] | None = None
+    system_message: str | None = None
+    model_client_stream: bool = False
+    model_context: ComponentModel | None = None
 
 
 class CodeExecutorAgent(BaseChatAgent, Component[CodeExecutorAgentConfig]):
@@ -108,13 +144,31 @@ def __init__(
         self,
         name: str,
         code_executor: CodeExecutor,
+        model_client: ChatCompletionClient,
         *,
+        model_context: ChatCompletionContext | None = None,
+        model_client_stream: bool = False,
         description: str = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).",
+        system_message: str
+        | None = "You are a helpful AI assistant. Solve tasks using your tools. Reply with TERMINATE when the task has been completed.",
         sources: Sequence[str] | None = None,
     ) -> None:
         super().__init__(name=name, description=description)
         self._code_executor = code_executor
         self._sources = sources
+        self._model_client = model_client
+        self._model_client_stream = model_client_stream
+
+        if model_context is not None:
+            self._model_context = model_context
+        else:
+            self._model_context = UnboundedChatCompletionContext()
+
+        self._system_messaages: List[SystemMessage] = []
+        if system_message is None:
+            self._system_messages = []
+        else:
+            self._system_messages = [SystemMessage(content=system_message)]
 
     @property
     def produced_message_types(self) -> Sequence[type[ChatMessage]]:
@@ -122,6 +176,81 @@ def produced_message_types(self) -> Sequence[type[ChatMessage]]:
         return (TextMessage,)
 
     async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
+        async for message in self.on_messages_stream(messages, cancellation_token):
+            if isinstance(message, Response):
+                return message
+        raise AssertionError("The stream should have returned the final result.")
+
+    async def on_messages_stream(
+        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
+    ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:
+        """
+        Process the incoming messages with the assistant agent and yield events/responses as they happen.
+        """
+
+        # Gather all relevant state here
+        agent_name = self.name
+        model_context = self._model_context
+        system_messages = self._system_messages
+        model_client = self._model_client
+        model_client_stream = self._model_client_stream
+
+        # STEP 1: Add new user/handoff messages to the model context
+        await self._add_messages_to_context(
+            model_context=model_context,
+            messages=messages,
+        )
+
+        # STEP 2: Update model context with any relevant memory
+        inner_messages: List[AgentEvent | ChatMessage] = []
+        for event_msg in await self._update_model_context_with_memory(
+            memory=None,
+            model_context=model_context,
+            agent_name=agent_name,
+        ):
+            inner_messages.append(event_msg)
+            yield event_msg
+
+        # STEP 3: Run the first inference
+        model_result = None
+        async for inference_output in self._call_llm(
+            model_client=model_client,
+            model_client_stream=model_client_stream,
+            system_messages=system_messages,
+            model_context=model_context,
+            agent_name=agent_name,
+            cancellation_token=cancellation_token,
+        ):
+            if isinstance(inference_output, CreateResult):
+                model_result = inference_output
+            else:
+                # Streaming chunk event
+                yield inference_output
+
+        assert model_result is not None, "No model result was produced."
+
+        # --- NEW: If the model produced a hidden "thought," yield it as an event ---
+        if model_result.thought:
+            thought_event = ThoughtEvent(content=model_result.thought, source=agent_name)
+            yield thought_event
+            inner_messages.append(thought_event)
+
+        # Add the assistant message to the model context (including thought if present)
+        await model_context.add_message(
+            AssistantMessage(
+                content=model_result.content,
+                source=agent_name,
+                thought=getattr(model_result, "thought", None),
+            )
+        )
+
+        # # execute generated code if present
+        # execution_result: Response = await self.execute_code_block(self, messages, cancellation_token)
+        # await execution_result
+
+    async def execute_code_block(
+        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
+    ) -> Response:
         # Extract code blocks from the messages.
         code_blocks: List[CodeBlock] = []
         for msg in messages:
@@ -166,16 +295,104 @@ def _extract_markdown_code_blocks(self, markdown_text: str) -> List[CodeBlock]:
     def _to_config(self) -> CodeExecutorAgentConfig:
         return CodeExecutorAgentConfig(
             name=self.name,
+            model_client=self._model_client.dump_component(),
             code_executor=self._code_executor.dump_component(),
             description=self.description,
             sources=list(self._sources) if self._sources is not None else None,
+            system_message=self._system_messages[0].content
+            if self._system_messages and isinstance(self._system_messages[0].content, str)
+            else None,
+            model_client_stream=self._model_client_stream,
+            model_context=self._model_context.dump_component(),
         )
 
     @classmethod
     def _from_config(cls, config: CodeExecutorAgentConfig) -> Self:
         return cls(
             name=config.name,
+            model_client=ChatCompletionClient.load_component(config.model_client),
             code_executor=CodeExecutor.load_component(config.code_executor),
             description=config.description,
             sources=config.sources,
+            system_message=config.system_message,
+            model_client_stream=config.model_client_stream,
+            model_context=None,
         )
+
+    @staticmethod
+    def _get_compatible_context(model_client: ChatCompletionClient, messages: List[LLMMessage]) -> Sequence[LLMMessage]:
+        """Ensure that the messages are compatible with the underlying client, by removing images if needed."""
+        if model_client.model_info["vision"]:
+            return messages
+        else:
+            return remove_images(messages)
+
+    @classmethod
+    async def _call_llm(
+        cls,
+        model_client: ChatCompletionClient,
+        model_client_stream: bool,
+        system_messages: List[SystemMessage],
+        model_context: ChatCompletionContext,
+        agent_name: str,
+        cancellation_token: CancellationToken,
+    ) -> AsyncGenerator[Union[CreateResult, ModelClientStreamingChunkEvent], None]:
+        """
+        Perform a model inference and yield either streaming chunk events or the final CreateResult.
+        """
+        all_messages = await model_context.get_messages()
+        llm_messages = cls._get_compatible_context(model_client=model_client, messages=system_messages + all_messages)
+
+        if model_client_stream:
+            model_result: Optional[CreateResult] = None
+            async for chunk in model_client.create_stream(
+                llm_messages, tools=[], cancellation_token=cancellation_token
+            ):
+                if isinstance(chunk, CreateResult):
+                    model_result = chunk
+                elif isinstance(chunk, str):
+                    yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name)
+                else:
+                    raise RuntimeError(f"Invalid chunk type: {type(chunk)}")
+            if model_result is None:
+                raise RuntimeError("No final model result in streaming mode.")
+            yield model_result
+        else:
+            model_result = await model_client.create(llm_messages, tools=[], cancellation_token=cancellation_token)
+            yield model_result
+
+    @staticmethod
+    async def _update_model_context_with_memory(
+        memory: Optional[Sequence[Memory]],
+        model_context: ChatCompletionContext,
+        agent_name: str,
+    ) -> List[MemoryQueryEvent]:
+        """
+        If memory modules are present, update the model context and return the events produced.
+        """
+        events: List[MemoryQueryEvent] = []
+        if memory:
+            for mem in memory:
+                update_context_result = await mem.update_context(model_context)
+                if update_context_result and len(update_context_result.memories.results) > 0:
+                    memory_query_event_msg = MemoryQueryEvent(
+                        content=update_context_result.memories.results,
+                        source=agent_name,
+                    )
+                    events.append(memory_query_event_msg)
+        return events
+
+    @staticmethod
+    async def _add_messages_to_context(
+        model_context: ChatCompletionContext,
+        messages: Sequence[ChatMessage],
+    ) -> None:
+        """
+        Add incoming user (and possibly handoff) messages to the model context.
+        """
+        for msg in messages:
+            if isinstance(msg, HandoffMessage):
+                # Add handoff context to the model context.
+                for context_msg in msg.context:
+                    await model_context.add_message(context_msg)
+            await model_context.add_message(UserMessage(content=msg.content, source=msg.source))
diff --git a/python/packages/autogen-agentchat/tests/test_code_executor_agent.py b/python/packages/autogen-agentchat/tests/test_code_executor_agent.py
index 2ebf79feb4ae..b59611dfb149 100644
--- a/python/packages/autogen-agentchat/tests/test_code_executor_agent.py
+++ b/python/packages/autogen-agentchat/tests/test_code_executor_agent.py
@@ -1,16 +1,48 @@
+import json
+
 import pytest
 from autogen_agentchat.agents import CodeExecutorAgent
 from autogen_agentchat.base import Response
-from autogen_agentchat.messages import TextMessage
-from autogen_core import CancellationToken
+from autogen_agentchat.messages import (
+    TextMessage,
+)
+from autogen_core import CancellationToken, FunctionCall
+from autogen_core.models import (
+    CreateResult,
+    RequestUsage,
+)
+from autogen_core.models._model_client import ModelFamily
 from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor
+from autogen_ext.models.replay import ReplayChatCompletionClient
 
 
 @pytest.mark.asyncio
 async def test_basic_code_execution() -> None:
     """Test basic code execution"""
 
-    agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())
+    model_client = ReplayChatCompletionClient(
+        [
+            CreateResult(
+                finish_reason="function_calls",
+                content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")],
+                usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                thought="Calling pass function",
+                cached=False,
+            ),
+            "pass",
+            "TERMINATE",
+        ],
+        model_info={
+            "function_calling": True,
+            "vision": True,
+            "json_output": True,
+            "family": ModelFamily.GPT_4O,
+            "structured_output": True,
+        },
+    )
+    agent = CodeExecutorAgent(
+        name="code_executor", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client
+    )
 
     messages = [
         TextMessage(
@@ -38,7 +70,29 @@ async def test_basic_code_execution() -> None:
 async def test_code_execution_error() -> None:
     """Test basic code execution"""
 
-    agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())
+    model_client = ReplayChatCompletionClient(
+        [
+            CreateResult(
+                finish_reason="function_calls",
+                content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")],
+                usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                thought="Calling pass function",
+                cached=False,
+            ),
+            "pass",
+            "TERMINATE",
+        ],
+        model_info={
+            "function_calling": True,
+            "vision": True,
+            "json_output": True,
+            "family": ModelFamily.GPT_4O,
+            "structured_output": True,
+        },
+    )
+    agent = CodeExecutorAgent(
+        name="code_executor", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client
+    )
 
     messages = [
         TextMessage(
@@ -66,7 +120,29 @@ async def test_code_execution_error() -> None:
 async def test_code_execution_no_output() -> None:
     """Test basic code execution"""
 
-    agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())
+    model_client = ReplayChatCompletionClient(
+        [
+            CreateResult(
+                finish_reason="function_calls",
+                content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")],
+                usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                thought="Calling pass function",
+                cached=False,
+            ),
+            "pass",
+            "TERMINATE",
+        ],
+        model_info={
+            "function_calling": True,
+            "vision": True,
+            "json_output": True,
+            "family": ModelFamily.GPT_4O,
+            "structured_output": True,
+        },
+    )
+    agent = CodeExecutorAgent(
+        name="code_executor", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client
+    )
 
     messages = [
         TextMessage(
@@ -95,7 +171,29 @@ async def test_code_execution_no_output() -> None:
 async def test_code_execution_no_block() -> None:
     """Test basic code execution"""
 
-    agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())
+    model_client = ReplayChatCompletionClient(
+        [
+            CreateResult(
+                finish_reason="function_calls",
+                content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")],
+                usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                thought="Calling pass function",
+                cached=False,
+            ),
+            "pass",
+            "TERMINATE",
+        ],
+        model_info={
+            "function_calling": True,
+            "vision": True,
+            "json_output": True,
+            "family": ModelFamily.GPT_4O,
+            "structured_output": True,
+        },
+    )
+    agent = CodeExecutorAgent(
+        name="code_executor", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client
+    )
 
     messages = [
         TextMessage(
@@ -122,7 +220,29 @@ async def test_code_execution_no_block() -> None:
 async def test_code_execution_multiple_blocks() -> None:
     """Test basic code execution"""
 
-    agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())
+    model_client = ReplayChatCompletionClient(
+        [
+            CreateResult(
+                finish_reason="function_calls",
+                content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")],
+                usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                thought="Calling pass function",
+                cached=False,
+            ),
+            "pass",
+            "TERMINATE",
+        ],
+        model_info={
+            "function_calling": True,
+            "vision": True,
+            "json_output": True,
+            "family": ModelFamily.GPT_4O,
+            "structured_output": True,
+        },
+    )
+    agent = CodeExecutorAgent(
+        name="code_executor", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client
+    )
 
     messages = [
         TextMessage(
@@ -170,7 +290,29 @@ async def test_code_execution_multiple_blocks() -> None:
 async def test_code_execution_agent_serialization() -> None:
     """Test agent config serialization"""
 
-    agent = CodeExecutorAgent(name="code_executor", code_executor=LocalCommandLineCodeExecutor())
+    model_client = ReplayChatCompletionClient(
+        [
+            CreateResult(
+                finish_reason="function_calls",
+                content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")],
+                usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                thought="Calling pass function",
+                cached=False,
+            ),
+            "pass",
+            "TERMINATE",
+        ],
+        model_info={
+            "function_calling": True,
+            "vision": True,
+            "json_output": True,
+            "family": ModelFamily.GPT_4O,
+            "structured_output": True,
+        },
+    )
+    agent = CodeExecutorAgent(
+        name="code_executor", code_executor=LocalCommandLineCodeExecutor(), model_client=model_client
+    )
 
     # Serialize and deserialize the agent
     serialized_agent = agent.dump_component()
diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py
index 5ab1605a72b7..0ab8245030ea 100644
--- a/python/packages/autogen-agentchat/tests/test_group_chat.py
+++ b/python/packages/autogen-agentchat/tests/test_group_chat.py
@@ -143,7 +143,7 @@ async def test_round_robin_group_chat(runtime: AgentRuntime | None) -> None:
     )
     with tempfile.TemporaryDirectory() as temp_dir:
         code_executor_agent = CodeExecutorAgent(
-            "code_executor", code_executor=LocalCommandLineCodeExecutor(work_dir=temp_dir)
+            "code_executor", code_executor=LocalCommandLineCodeExecutor(work_dir=temp_dir), model_client=model_client
         )
         coding_assistant_agent = AssistantAgent(
             "coding_assistant",


Code review comments:
diff_hunk: @@ -108,20 +144,113 @@ def __init__(
         self,
         name: str,
         code_executor: CodeExecutor,
+        model_client: ChatCompletionClient,
path: python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py
text: @user1:
Let's make it optional to provide backward compatibility. So if `model_client` is provided it will enable the new behavior. Otherwise keep the same behavior as before. 

@author:
Done!

- I made `model_client` optional, if not provided, the `CodeExecutorAgent` will act as before.
- If provided it will generate the text and yield the inferred message.
- If the inferred message contains code it will execute it and yield its output.

diff_hunk: @@ -108,20 +144,113 @@ def __init__(
         self,
         name: str,
         code_executor: CodeExecutor,
+        model_client: ChatCompletionClient,
         *,
+        model_context: ChatCompletionContext | None = None,
+        model_client_stream: bool = False,
         description: str = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).",
+        system_message: str
path: python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py
text: @user1:
Make system message default to something related to code generation and execution.  

@author:
Done!

Updated the default system message to be more focussed on code generation and execution.