Rename to use BaseChatMessage and BaseAgentEvent. Bring back union types. (#6144)

Rename the `ChatMessage` and `AgentEvent` base classes to `BaseChatMessage` and `BaseAgentEvent`. 

Bring back the `ChatMessage` and `AgentEvent` as union of built-in concrete types to avoid breaking existing applications that depends on Pydantic serialization. 

Why?

Many existing code uses containers like this:

```python
class AppMessage(BaseModel):
   name: str
   message: ChatMessage 

# Serialization is this:
m = AppMessage(...)
m.model_dump_json()

# Fields like HandoffMessage.target will be lost because it is now treated as a base class without content or target fields.
```

The assumption on `ChatMessage` or `AgentEvent` to be a union of concrete types could be in many existing code bases. So this PR brings back the union types, while keep method type hints such as those on `on_messages` to use the `BaseChatMessage` and `BaseAgentEvent` base classes for flexibility.

Author: ekzhu

python/packages/agbench/benchmarks/GAIA/Templates/SelectorGroupChat/scenario.py

@@ -16,7 +16,7 @@
 from autogen_ext.agents.web_surfer import MultimodalWebSurfer
 from autogen_ext.agents.file_surfer import FileSurfer
 from autogen_agentchat.agents import CodeExecutorAgent
-from autogen_agentchat.messages import TextMessage, AgentEvent, ChatMessage, HandoffMessage, MultiModalMessage, StopMessage
+from autogen_agentchat.messages import TextMessage, BaseAgentEvent, BaseChatMessage, HandoffMessage, MultiModalMessage, StopMessage
 from autogen_core.models import LLMMessage, UserMessage, AssistantMessage
 
 # Suppress warnings about the requests.Session() not being closed
@@ -141,7 +141,7 @@ def __init__(self, prompt: str, model_client: ChatCompletionClient, termination_
     def terminated(self) -> bool:
         return self._terminated
 
-    async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None:
+    async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
         if self._terminated:
             raise TerminatedException("Termination condition has already been reached")
 

python/packages/agbench/src/agbench/linter/__init__.py

@@ -1,4 +1,4 @@
 # __init__.py
-from ._base import Code, Document, CodedDocument, BaseQualitativeCoder
+from ._base import BaseQualitativeCoder, Code, CodedDocument, Document
 
 __all__ = ["Code", "Document", "CodedDocument", "BaseQualitativeCoder"]

python/packages/agbench/src/agbench/linter/_base.py

@@ -1,7 +1,8 @@
-import json
 import hashlib
+import json
 import re
-from typing import Protocol, List, Set, Optional
+from typing import List, Optional, Protocol, Set
+
 from pydantic import BaseModel, Field
 
 

python/packages/agbench/src/agbench/linter/cli.py

@@ -1,8 +1,10 @@
-import os
 import argparse
-from typing import List, Sequence, Optional
+import os
+from typing import List, Optional, Sequence
+
 from openai import OpenAI
-from ._base import Document, CodedDocument
+
+from ._base import CodedDocument, Document
 from .coders.oai_coder import OAIQualitativeCoder
 
 

python/packages/agbench/src/agbench/linter/coders/oai_coder.py

@@ -1,13 +1,11 @@
 import os
 import re
-
-from typing import List, Set, Optional
-from pydantic import BaseModel
+from typing import List, Optional, Set
 
 from openai import OpenAI
+from pydantic import BaseModel
 
-from .._base import CodedDocument, Document, Code
-from .._base import BaseQualitativeCoder
+from .._base import BaseQualitativeCoder, Code, CodedDocument, Document
 
 
 class CodeList(BaseModel):

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

@@ -40,8 +40,8 @@
 from ..base import Handoff as HandoffBase
 from ..base import Response
 from ..messages import (
-    AgentEvent,
-    ChatMessage,
+    BaseAgentEvent,
+    BaseChatMessage,
     HandoffMessage,
     MemoryQueryEvent,
     ModelClientStreamingChunkEvent,
@@ -697,8 +697,8 @@ def __init__(
         self._is_running = False
 
     @property
-    def produced_message_types(self) -> Sequence[type[ChatMessage]]:
-        message_types: List[type[ChatMessage]] = [TextMessage]
+    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
+        message_types: List[type[BaseChatMessage]] = [TextMessage]
         if self._handoffs:
             message_types.append(HandoffMessage)
         if self._tools:
@@ -712,15 +712,15 @@ def model_context(self) -> ChatCompletionContext:
         """
         return self._model_context
 
-    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
+    async def on_messages(self, messages: Sequence[BaseChatMessage], 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]:
+        self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
+    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
         """
         Process the incoming messages with the assistant agent and yield events/responses as they happen.
         """
@@ -745,7 +745,7 @@ async def on_messages_stream(
         )
 
         # STEP 2: Update model context with any relevant memory
-        inner_messages: List[AgentEvent | ChatMessage] = []
+        inner_messages: List[BaseAgentEvent | BaseChatMessage] = []
         for event_msg in await self._update_model_context_with_memory(
             memory=memory,
             model_context=model_context,
@@ -810,7 +810,7 @@ async def on_messages_stream(
     @staticmethod
     async def _add_messages_to_context(
         model_context: ChatCompletionContext,
-        messages: Sequence[ChatMessage],
+        messages: Sequence[BaseChatMessage],
     ) -> None:
         """
         Add incoming messages to the model context.
@@ -886,7 +886,7 @@ async def _call_llm(
     async def _process_model_result(
         cls,
         model_result: CreateResult,
-        inner_messages: List[AgentEvent | ChatMessage],
+        inner_messages: List[BaseAgentEvent | BaseChatMessage],
         cancellation_token: CancellationToken,
         agent_name: str,
         system_messages: List[SystemMessage],
@@ -898,7 +898,7 @@ async def _process_model_result(
         model_client_stream: bool,
         reflect_on_tool_use: bool,
         tool_call_summary_format: str,
-    ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:
+    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
         """
         Handle final or partial responses from model_result, including tool calls, handoffs,
         and reflection if needed.
@@ -992,7 +992,7 @@ async def _process_model_result(
     def _check_and_handle_handoff(
         model_result: CreateResult,
         executed_calls_and_results: List[Tuple[FunctionCall, FunctionExecutionResult]],
-        inner_messages: List[AgentEvent | ChatMessage],
+        inner_messages: List[BaseAgentEvent | BaseChatMessage],
         handoffs: Dict[str, HandoffBase],
         agent_name: str,
     ) -> Optional[Response]:
@@ -1061,7 +1061,7 @@ async def _reflect_on_tool_use_flow(
         model_client_stream: bool,
         model_context: ChatCompletionContext,
         agent_name: str,
-        inner_messages: List[AgentEvent | ChatMessage],
+        inner_messages: List[BaseAgentEvent | BaseChatMessage],
     ) -> AsyncGenerator[Response | ModelClientStreamingChunkEvent | ThoughtEvent, None]:
         """
         If reflect_on_tool_use=True, we do another inference based on tool results
@@ -1113,7 +1113,7 @@ async def _reflect_on_tool_use_flow(
     @staticmethod
     def _summarize_tool_use(
         executed_calls_and_results: List[Tuple[FunctionCall, FunctionExecutionResult]],
-        inner_messages: List[AgentEvent | ChatMessage],
+        inner_messages: List[BaseAgentEvent | BaseChatMessage],
         handoffs: Dict[str, HandoffBase],
         tool_call_summary_format: str,
         agent_name: str,

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

@@ -6,8 +6,8 @@
 
 from ..base import ChatAgent, Response, TaskResult
 from ..messages import (
-    AgentEvent,
-    ChatMessage,
+    BaseAgentEvent,
+    BaseChatMessage,
     ModelClientStreamingChunkEvent,
     TextMessage,
 )
@@ -59,13 +59,13 @@ def description(self) -> str:
 
     @property
     @abstractmethod
-    def produced_message_types(self) -> Sequence[type[ChatMessage]]:
+    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
         """The types of messages that the agent produces in the
-        :attr:`Response.chat_message` field. They must be :class:`ChatMessage` types."""
+        :attr:`Response.chat_message` field. They must be :class:`BaseChatMessage` types."""
         ...
 
     @abstractmethod
-    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
+    async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
         """Handles incoming messages and returns a response.
 
         .. note::
@@ -81,8 +81,8 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token:
         ...
 
     async def on_messages_stream(
-        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
-    ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:
+        self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
+    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
         """Handles incoming messages and returns a stream of messages and
         and the final item is the response. The base implementation in
         :class:`BaseChatAgent` simply calls :meth:`on_messages` and yields
@@ -106,29 +106,29 @@ async def on_messages_stream(
     async def run(
         self,
         *,
-        task: str | ChatMessage | Sequence[ChatMessage] | None = None,
+        task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
         cancellation_token: CancellationToken | None = None,
     ) -> TaskResult:
         """Run the agent with the given task and return the result."""
         if cancellation_token is None:
             cancellation_token = CancellationToken()
-        input_messages: List[ChatMessage] = []
-        output_messages: List[AgentEvent | ChatMessage] = []
+        input_messages: List[BaseChatMessage] = []
+        output_messages: List[BaseAgentEvent | BaseChatMessage] = []
         if task is None:
             pass
         elif isinstance(task, str):
             text_msg = TextMessage(content=task, source="user")
             input_messages.append(text_msg)
             output_messages.append(text_msg)
-        elif isinstance(task, ChatMessage):
+        elif isinstance(task, BaseChatMessage):
             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, ChatMessage):
+                if isinstance(msg, BaseChatMessage):
                     input_messages.append(msg)
                     output_messages.append(msg)
                 else:
@@ -142,31 +142,31 @@ async def run(
     async def run_stream(
         self,
         *,
-        task: str | ChatMessage | Sequence[ChatMessage] | None = None,
+        task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
         cancellation_token: CancellationToken | None = None,
-    ) -> AsyncGenerator[AgentEvent | ChatMessage | TaskResult, None]:
+    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]:
         """Run the agent with the given task and return a stream of messages
         and the final task result as the last item in the stream."""
         if cancellation_token is None:
             cancellation_token = CancellationToken()
-        input_messages: List[ChatMessage] = []
-        output_messages: List[AgentEvent | ChatMessage] = []
+        input_messages: List[BaseChatMessage] = []
+        output_messages: List[BaseAgentEvent | BaseChatMessage] = []
         if task is None:
             pass
         elif isinstance(task, str):
             text_msg = TextMessage(content=task, source="user")
             input_messages.append(text_msg)
             output_messages.append(text_msg)
             yield text_msg
-        elif isinstance(task, ChatMessage):
+        elif isinstance(task, BaseChatMessage):
             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, ChatMessage):
+                if isinstance(msg, BaseChatMessage):
                     input_messages.append(msg)
                     output_messages.append(msg)
                     yield msg

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

@@ -7,7 +7,7 @@
 from typing_extensions import Self
 
 from ..base import Response
-from ..messages import ChatMessage, TextMessage
+from ..messages import BaseChatMessage, TextMessage
 from ._base_chat_agent import BaseChatAgent
 
 
@@ -119,11 +119,11 @@ def __init__(
         self._sources = sources
 
     @property
-    def produced_message_types(self) -> Sequence[type[ChatMessage]]:
+    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
         """The types of messages that the code executor agent produces."""
         return (TextMessage,)
 
-    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
+    async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
         # Extract code blocks from the messages.
         code_blocks: List[CodeBlock] = []
         for msg in messages:

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

@@ -10,8 +10,8 @@
 
 from ..base import TaskResult, Team
 from ..messages import (
-    AgentEvent,
-    ChatMessage,
+    BaseAgentEvent,
+    BaseChatMessage,
     ModelClientStreamingChunkEvent,
     TextMessage,
 )
@@ -122,10 +122,10 @@ def __init__(
         self._response_prompt = response_prompt
 
     @property
-    def produced_message_types(self) -> Sequence[type[ChatMessage]]:
+    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
         return (TextMessage,)
 
-    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
+    async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
         # Call the stream method and collect the messages.
         response: Response | None = None
         async for msg in self.on_messages_stream(messages, cancellation_token):
@@ -135,14 +135,14 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token:
         return response
 
     async def on_messages_stream(
-        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
-    ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:
+        self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
+    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
         # Prepare the task for the team of agents.
         task = list(messages)
 
         # Run the team of agents.
         result: TaskResult | None = None
-        inner_messages: List[AgentEvent | ChatMessage] = []
+        inner_messages: List[BaseAgentEvent | BaseChatMessage] = []
         count = 0
         async for inner_msg in self._team.run_stream(task=task, cancellation_token=cancellation_token):
             if isinstance(inner_msg, TaskResult):
@@ -167,7 +167,7 @@ async def on_messages_stream(
             # Generate a response using the model client.
             llm_messages: List[LLMMessage] = [SystemMessage(content=self._instruction)]
             for message in messages:
-                if isinstance(message, ChatMessage):
+                if isinstance(message, BaseChatMessage):
                     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)

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

@@ -10,7 +10,7 @@
 from typing_extensions import Self
 
 from ..base import Response
-from ..messages import AgentEvent, ChatMessage, HandoffMessage, TextMessage, UserInputRequestedEvent
+from ..messages import BaseAgentEvent, BaseChatMessage, HandoffMessage, TextMessage, UserInputRequestedEvent
 from ._base_chat_agent import BaseChatAgent
 
 SyncInputFunc = Callable[[str], str]
@@ -170,11 +170,11 @@ def __init__(
         self._is_async = iscoroutinefunction(self.input_func)
 
     @property
-    def produced_message_types(self) -> Sequence[type[ChatMessage]]:
+    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
         """Message types this agent can produce."""
         return (TextMessage, HandoffMessage)
 
-    def _get_latest_handoff(self, messages: Sequence[ChatMessage]) -> Optional[HandoffMessage]:
+    def _get_latest_handoff(self, messages: Sequence[BaseChatMessage]) -> Optional[HandoffMessage]:
         """Find the HandoffMessage in the message sequence that addresses this agent."""
         if len(messages) > 0 and isinstance(messages[-1], HandoffMessage):
             if messages[-1].target == self.name:
@@ -201,15 +201,15 @@ async def _get_input(self, prompt: str, cancellation_token: Optional[Cancellatio
         except Exception as e:
             raise RuntimeError(f"Failed to get user input: {str(e)}") from e
 
-    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
+    async def on_messages(self, messages: Sequence[BaseChatMessage], 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]:
+        self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
+    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
         """Handle incoming messages by requesting user input."""
         try:
             # Check for handoff first