feat: add list input support for chat messages with unit tests

Author: iamarunbrahma

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

@@ -4,7 +4,14 @@
 from autogen_core import CancellationToken
 
 from ..base import ChatAgent, Response, TaskResult
-from ..messages import AgentMessage, ChatMessage, HandoffMessage, MultiModalMessage, StopMessage, TextMessage
+from ..messages import (
+    AgentMessage,
+    ChatMessage,
+    HandoffMessage,
+    MultiModalMessage,
+    StopMessage,
+    TextMessage,
+)
 from ..state import BaseState
 
 
@@ -45,8 +52,9 @@ async def on_messages_stream(
         self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
     ) -> AsyncGenerator[AgentMessage | 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 the messages in the response."""
+        and the final item is the response. The base implementation in
+        :class:`BaseChatAgent` simply calls :meth:`on_messages` and yields
+        the messages in the response."""
         response = await self.on_messages(messages, cancellation_token)
         for inner_message in response.inner_messages or []:
             yield inner_message
@@ -69,6 +77,13 @@ async def run(
             text_msg = TextMessage(content=task, source="user")
             input_messages.append(text_msg)
             output_messages.append(text_msg)
+        elif isinstance(task, list):
+            for msg in task:
+                if isinstance(msg, (TextMessage, MultiModalMessage, StopMessage, HandoffMessage)):
+                    input_messages.append(msg)
+                    output_messages.append(msg)
+                else:
+                    raise ValueError(f"Invalid message type in list: {type(msg)}")
         elif isinstance(task, (TextMessage, MultiModalMessage, StopMessage, HandoffMessage)):
             input_messages.append(task)
             output_messages.append(task)

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py

@@ -22,7 +22,6 @@
     ToolCallMessage,
     ToolCallResultMessage,
 )
-
 from ....state import MagenticOneOrchestratorState
 from .._base_group_chat_manager import BaseGroupChatManager
 from .._events import (

python/packages/autogen-agentchat/tests/test_assistant_agent.py

@@ -8,6 +8,7 @@
 from autogen_agentchat.agents import AssistantAgent
 from autogen_agentchat.base import Handoff, TaskResult
 from autogen_agentchat.messages import (
+    ChatMessage,
     HandoffMessage,
     MultiModalMessage,
     TextMessage,
@@ -21,7 +22,10 @@
 from openai.types.chat.chat_completion import ChatCompletion, Choice
 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 from openai.types.chat.chat_completion_message import ChatCompletionMessage
-from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
+from openai.types.chat.chat_completion_message_tool_call import (
+    ChatCompletionMessageToolCall,
+    Function,
+)
 from openai.types.completion_usage import CompletionUsage
 from utils import FileLogHandler
 
@@ -33,14 +37,14 @@
 class _MockChatCompletion:
     def __init__(self, chat_completions: List[ChatCompletion]) -> None:
         self._saved_chat_completions = chat_completions
-        self._curr_index = 0
+        self.curr_index = 0
 
     async def mock_create(
         self, *args: Any, **kwargs: Any
     ) -> ChatCompletion | AsyncGenerator[ChatCompletionChunk, None]:
         await asyncio.sleep(0.1)
-        completion = self._saved_chat_completions[self._curr_index]
-        self._curr_index += 1
+        completion = self._saved_chat_completions[self.curr_index]
+        self.curr_index += 1
         return completion
 
 
@@ -90,7 +94,11 @@ async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None:
         ChatCompletion(
             id="id2",
             choices=[
-                Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(content="Hello", role="assistant"))
+                Choice(
+                    finish_reason="stop",
+                    index=0,
+                    message=ChatCompletionMessage(content="Hello", role="assistant"),
+                )
             ],
             created=0,
             model=model,
@@ -101,7 +109,9 @@ async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None:
             id="id2",
             choices=[
                 Choice(
-                    finish_reason="stop", index=0, message=ChatCompletionMessage(content="TERMINATE", role="assistant")
+                    finish_reason="stop",
+                    index=0,
+                    message=ChatCompletionMessage(content="TERMINATE", role="assistant"),
                 )
             ],
             created=0,
@@ -115,7 +125,11 @@ async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None:
     agent = AssistantAgent(
         "tool_use_agent",
         model_client=OpenAIChatCompletionClient(model=model, api_key=""),
-        tools=[_pass_function, _fail_function, FunctionTool(_echo_function, description="Echo")],
+        tools=[
+            _pass_function,
+            _fail_function,
+            FunctionTool(_echo_function, description="Echo"),
+        ],
     )
     result = await agent.run(task="task")
     assert len(result.messages) == 4
@@ -133,7 +147,7 @@ async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None:
     assert result.messages[3].models_usage.prompt_tokens == 10
 
     # Test streaming.
-    mock._curr_index = 0  # pyright: ignore
+    mock.curr_index = 0  # pyright: ignore
     index = 0
     async for message in agent.run_stream(task="task"):
         if isinstance(message, TaskResult):
@@ -147,7 +161,11 @@ async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None:
     agent2 = AssistantAgent(
         "tool_use_agent",
         model_client=OpenAIChatCompletionClient(model=model, api_key=""),
-        tools=[_pass_function, _fail_function, FunctionTool(_echo_function, description="Echo")],
+        tools=[
+            _pass_function,
+            _fail_function,
+            FunctionTool(_echo_function, description="Echo"),
+        ],
     )
     await agent2.load_state(state)
     state2 = await agent2.save_state()
@@ -192,7 +210,11 @@ async def test_handoffs(monkeypatch: pytest.MonkeyPatch) -> None:
     tool_use_agent = AssistantAgent(
         "tool_use_agent",
         model_client=OpenAIChatCompletionClient(model=model, api_key=""),
-        tools=[_pass_function, _fail_function, FunctionTool(_echo_function, description="Echo")],
+        tools=[
+            _pass_function,
+            _fail_function,
+            FunctionTool(_echo_function, description="Echo"),
+        ],
         handoffs=[handoff],
     )
     assert HandoffMessage in tool_use_agent.produced_message_types
@@ -212,7 +234,7 @@ async def test_handoffs(monkeypatch: pytest.MonkeyPatch) -> None:
     assert result.messages[3].models_usage is None
 
     # Test streaming.
-    mock._curr_index = 0  # pyright: ignore
+    mock.curr_index = 0  # pyright: ignore
     index = 0
     async for message in tool_use_agent.run_stream(task="task"):
         if isinstance(message, TaskResult):
@@ -229,7 +251,11 @@ async def test_multi_modal_task(monkeypatch: pytest.MonkeyPatch) -> None:
         ChatCompletion(
             id="id2",
             choices=[
-                Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(content="Hello", role="assistant"))
+                Choice(
+                    finish_reason="stop",
+                    index=0,
+                    message=ChatCompletionMessage(content="Hello", role="assistant"),
+                )
             ],
             created=0,
             model=model,
@@ -239,7 +265,10 @@ async def test_multi_modal_task(monkeypatch: pytest.MonkeyPatch) -> None:
     ]
     mock = _MockChatCompletion(chat_completions)
     monkeypatch.setattr(AsyncCompletions, "create", mock.mock_create)
-    agent = AssistantAgent(name="assistant", model_client=OpenAIChatCompletionClient(model=model, api_key=""))
+    agent = AssistantAgent(
+        name="assistant",
+        model_client=OpenAIChatCompletionClient(model=model, api_key=""),
+    )
     # Generate a random base64 image.
     img_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"
     result = await agent.run(task=MultiModalMessage(source="user", content=["Test", Image.from_base64(img_base64)]))
@@ -250,14 +279,24 @@ async def test_multi_modal_task(monkeypatch: pytest.MonkeyPatch) -> None:
 async def test_invalid_model_capabilities() -> None:
     model = "random-model"
     model_client = OpenAIChatCompletionClient(
-        model=model, api_key="", model_capabilities={"vision": False, "function_calling": False, "json_output": False}
+        model=model,
+        api_key="",
+        model_capabilities={
+            "vision": False,
+            "function_calling": False,
+            "json_output": False,
+        },
     )
 
     with pytest.raises(ValueError):
         agent = AssistantAgent(
             name="assistant",
             model_client=model_client,
-            tools=[_pass_function, _fail_function, FunctionTool(_echo_function, description="Echo")],
+            tools=[
+                _pass_function,
+                _fail_function,
+                FunctionTool(_echo_function, description="Echo"),
+            ],
         )
 
     with pytest.raises(ValueError):
@@ -268,3 +307,62 @@ async def test_invalid_model_capabilities() -> None:
         # Generate a random base64 image.
         img_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"
         await agent.run(task=MultiModalMessage(source="user", content=["Test", Image.from_base64(img_base64)]))
+
+
+@pytest.mark.asyncio
+async def test_list_chat_messages(monkeypatch: pytest.MonkeyPatch) -> None:
+    model = "gpt-4o-2024-05-13"
+    chat_completions = [
+        ChatCompletion(
+            id="id1",
+            choices=[
+                Choice(
+                    finish_reason="stop",
+                    index=0,
+                    message=ChatCompletionMessage(content="Response to message 1", role="assistant"),
+                )
+            ],
+            created=0,
+            model=model,
+            object="chat.completion",
+            usage=CompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
+        ),
+    ]
+    mock = _MockChatCompletion(chat_completions)
+    monkeypatch.setattr(AsyncCompletions, "create", mock.mock_create)
+    agent = AssistantAgent(
+        "test_agent",
+        model_client=OpenAIChatCompletionClient(model=model, api_key=""),
+    )
+
+    # Create a list of chat messages
+    messages: List[ChatMessage] = [
+        TextMessage(content="Message 1", source="user"),
+        TextMessage(content="Message 2", source="user"),
+    ]
+
+    # Test run method with list of messages
+    result = await agent.run(task=messages)
+    assert len(result.messages) == 3  # 2 input messages + 1 response message
+    assert isinstance(result.messages[0], TextMessage)
+    assert result.messages[0].content == "Message 1"
+    assert result.messages[0].source == "user"
+    assert isinstance(result.messages[1], TextMessage)
+    assert result.messages[1].content == "Message 2"
+    assert result.messages[1].source == "user"
+    assert isinstance(result.messages[2], TextMessage)
+    assert result.messages[2].content == "Response to message 1"
+    assert result.messages[2].source == "test_agent"
+    assert result.messages[2].models_usage is not None
+    assert result.messages[2].models_usage.completion_tokens == 5
+    assert result.messages[2].models_usage.prompt_tokens == 10
+
+    # Test run_stream method with list of messages
+    mock.curr_index = 0  # Reset mock index using public attribute
+    index = 0
+    async for message in agent.run_stream(task=messages):
+        if isinstance(message, TaskResult):
+            assert message == result
+        else:
+            assert message == result.messages[index]
+        index += 1