Use class hierarchy to organize AgentChat message types and introduce StructuredMessage type

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py

@@ -31,7 +31,6 @@
     LLMMessage,
     ModelFamily,
     SystemMessage,
-    UserMessage,
 )
 from autogen_core.tools import BaseTool, FunctionTool
 from pydantic import BaseModel
@@ -814,14 +813,13 @@ async def _add_messages_to_context(
         messages: Sequence[ChatMessage],
     ) -> None:
         """
-        Add incoming user (and possibly handoff) messages to the model context.
+        Add incoming 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))
+                for llm_msg in msg.context:
+                    await model_context.add_message(llm_msg)
+            await model_context.add_message(msg.to_model_message())
 
     @staticmethod
     async def _update_model_context_with_memory(

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py

@@ -7,7 +7,6 @@
 from ..base import ChatAgent, Response, TaskResult
 from ..messages import (
     AgentEvent,
-    BaseChatMessage,
     ChatMessage,
     ModelClientStreamingChunkEvent,
     TextMessage,
@@ -121,15 +120,15 @@ async def run(
             text_msg = TextMessage(content=task, source="user")
             input_messages.append(text_msg)
             output_messages.append(text_msg)
-        elif isinstance(task, BaseChatMessage):
+        elif isinstance(task, ChatMessage):
             input_messages.append(task)
             output_messages.append(task)
         else:
             if not task:
                 raise ValueError("Task list cannot be empty.")
             # Task is a sequence of messages.
             for msg in task:
-                if isinstance(msg, BaseChatMessage):
+                if isinstance(msg, ChatMessage):
                     input_messages.append(msg)
                     output_messages.append(msg)
                 else:
@@ -159,15 +158,15 @@ async def run_stream(
             input_messages.append(text_msg)
             output_messages.append(text_msg)
             yield text_msg
-        elif isinstance(task, BaseChatMessage):
+        elif isinstance(task, ChatMessage):
             input_messages.append(task)
             output_messages.append(task)
             yield task
         else:
             if not task:
                 raise ValueError("Task list cannot be empty.")
             for msg in task:
-                if isinstance(msg, BaseChatMessage):
+                if isinstance(msg, ChatMessage):
                     input_messages.append(msg)
                     output_messages.append(msg)
                     yield msg

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py

@@ -21,7 +21,9 @@ class CodeExecutorAgentConfig(BaseModel):
 
 
 class CodeExecutorAgent(BaseChatAgent, Component[CodeExecutorAgentConfig]):
-    """An agent that extracts and executes code snippets found in received messages and returns the output.
+    """An agent that extracts and executes code snippets found in received
+    :class:`~autogen_agentchat.messages.TextMessage` messages and returns the output
+    of the code execution.
 
     It is typically used within a team with another agent that generates code snippets to be executed.
 

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py

@@ -1,7 +1,7 @@
 from typing import Any, AsyncGenerator, List, Mapping, Sequence
 
 from autogen_core import CancellationToken, Component, ComponentModel
-from autogen_core.models import ChatCompletionClient, LLMMessage, SystemMessage, UserMessage
+from autogen_core.models import ChatCompletionClient, LLMMessage, SystemMessage
 from pydantic import BaseModel
 from typing_extensions import Self
 
@@ -11,7 +11,6 @@
 from ..base import TaskResult, Team
 from ..messages import (
     AgentEvent,
-    BaseChatMessage,
     ChatMessage,
     ModelClientStreamingChunkEvent,
     TextMessage,
@@ -167,13 +166,9 @@ async def on_messages_stream(
         else:
             # Generate a response using the model client.
             llm_messages: List[LLMMessage] = [SystemMessage(content=self._instruction)]
-            llm_messages.extend(
-                [
-                    UserMessage(content=message.content, source=message.source)
-                    for message in inner_messages
-                    if isinstance(message, BaseChatMessage)
-                ]
-            )
+            for message in messages:
+                if isinstance(message, ChatMessage):
+                    llm_messages.append(message.to_model_message())
             llm_messages.append(SystemMessage(content=self._response_prompt))
             completion = await self._model_client.create(messages=llm_messages, cancellation_token=cancellation_token)
             assert isinstance(completion.content, str)

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py

@@ -82,6 +82,7 @@ async def simple_user_agent():
                         cancellation_token=CancellationToken(),
                     )
                 )
+                assert isinstance(response.chat_message, TextMessage)
                 print(f"Your name is {response.chat_message.content}")
 
     Example:
@@ -117,6 +118,7 @@ async def cancellable_user_agent():
                         )
                     )
                     response = await agent_task
+                    assert isinstance(response.chat_message, TextMessage)
                     print(f"Your name is {response.chat_message.content}")
                 except Exception as e:
                     print(f"Exception: {e}")

python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py

@@ -11,7 +11,6 @@
     BaseChatMessage,
     ChatMessage,
     HandoffMessage,
-    MultiModalMessage,
     StopMessage,
     TextMessage,
     ToolCallExecutionEvent,
@@ -137,18 +136,12 @@ async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMe
             if self._sources is not None and message.source not in self._sources:
                 continue
 
-            if isinstance(message.content, str) and self._termination_text in message.content:
+            content = message.to_text()
+            if self._termination_text in content:
                 self._terminated = True
                 return StopMessage(
                     content=f"Text '{self._termination_text}' mentioned", source="TextMentionTermination"
                 )
-            elif isinstance(message, MultiModalMessage):
-                for item in message.content:
-                    if isinstance(item, str) and self._termination_text in item:
-                        self._terminated = True
-                        return StopMessage(
-                            content=f"Text '{self._termination_text}' mentioned", source="TextMentionTermination"
-                        )
         return None
 
     async def reset(self) -> None: