Improve SocietyOfMindAgent message handling (#6142)

Please refer to #6123  for full context.

That issue outlines several design and behavioral problems with
`SocietyOfMindAgent`.
This DRAFT PR focuses on resolving the most critical and broken
behaviors first.

Here is the error list
🔍 SocietyOfMindAgent: Design Issues and Historical Comparison (v0.2 vs
v0.4+)

### ✅ P1–P4 Regression Issue Table (Updated with Fixes in PR #6142)

| ID | Description | Current v0.4+ Issue | Resolution in PR #6142 | Was
it a problem in v0.2? | Notes |

|-----|-------------|----------------------|--------------------------|----------------------------|-------|
| **P1** | `inner_messages` leaks into outer team termination evaluation
| `Response.inner_messages` is appended to the outer team's
`_message_thread`, affecting termination conditions. Violates
encapsulation. | ✅ `inner_messages` is excluded from `_message_thread`,
avoiding contamination of outer termination logic. | ❌ No | Structural
boundary is now enforced |
| **P2** | Inner team does not execute when outer message history is
empty | In chained executions, if no new outer message exists, no task
is created and the inner team is skipped entirely | ✅ Detects absence of
new outer message and reuses the previous task, passing it via a handoff
message. This ensures the inner team always receives a valid task to
execute | ❌ No | The issue was silent task omission, not summary
failure. Summary succeeds as a downstream effect |
| **P3** | Summary LLM prompt is built from external input only | Prompt
is constructed using external message history, ignoring internal
reasoning | ✅ Prompt construction now uses
`final_response.inner_messages`, restoring internal reasoning as the
source of summarization | ❌ No | Matches v0.2 internal monologue
behavior |
| **P4** | External input is included in summary prompt (possibly
incorrectly) | Outer messages are used in the final LLM summarization
prompt | ✅ Resolved via the same fix as P3; outer messages are no longer
used for summary | ❌ No | Redundant with P3, now fully addressed |


<!-- Thank you for your contribution! Please review
https://microsoft.github.io/autogen/docs/Contribute before opening a
pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this
solves. -->

## Related issue number
resolve #6123 
Blocked #6168 (Sometimes SoMA send last whitespace message)
related #6187
<!-- For example: "Closes #1234" -->

## Checks

- [ ] 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.
- [ ] I've added tests (if relevant) corresponding to the changes
introduced in this PR.
- [ ] I've made sure all auto checks have passed.

---------

Co-authored-by: Eric Zhu <[email protected]>

Author: SongChiYoung

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

@@ -1,6 +1,10 @@
 from typing import Any, AsyncGenerator, List, Mapping, Sequence
 
 from autogen_core import CancellationToken, Component, ComponentModel
+from autogen_core.model_context import (
+    ChatCompletionContext,
+    UnboundedChatCompletionContext,
+)
 from autogen_core.models import ChatCompletionClient, LLMMessage, SystemMessage
 from pydantic import BaseModel
 from typing_extensions import Self
@@ -12,6 +16,7 @@
 from ..messages import (
     BaseAgentEvent,
     BaseChatMessage,
+    HandoffMessage,
     ModelClientStreamingChunkEvent,
     TextMessage,
 )
@@ -27,6 +32,7 @@ class SocietyOfMindAgentConfig(BaseModel):
     description: str | None = None
     instruction: str | None = None
     response_prompt: str | None = None
+    model_context: ComponentModel | None = None
 
 
 class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
@@ -38,6 +44,16 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
     Once the response is generated, the agent resets the inner team by
     calling :meth:`Team.reset`.
 
+    Limit context size sent to the model:
+
+    You can limit the number of messages sent to the model by setting
+    the `model_context` parameter to a :class:`~autogen_core.model_context.BufferedChatCompletionContext`.
+    This will limit the number of recent messages sent to the model and can be useful
+    when the model has a limit on the number of tokens it can process.
+    You can also create your own model context by subclassing
+    :class:`~autogen_core.model_context.ChatCompletionContext`.
+
+
     Args:
         name (str): The name of the agent.
         team (Team): The team of agents to use.
@@ -47,6 +63,8 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
             Defaults to :attr:`DEFAULT_INSTRUCTION`. It assumes the role of 'system'.
         response_prompt (str, optional): The response prompt to use when generating a response using the inner team's messages.
             Defaults to :attr:`DEFAULT_RESPONSE_PROMPT`. It assumes the role of 'system'.
+        model_context (ChatCompletionContext | None, optional): The model context for storing and retrieving :class:`~autogen_core.models.LLMMessage`. It can be preloaded with initial messages. The initial messages will be cleared when the agent is reset.
+
 
 
     Example:
@@ -114,17 +132,30 @@ def __init__(
         description: str = DEFAULT_DESCRIPTION,
         instruction: str = DEFAULT_INSTRUCTION,
         response_prompt: str = DEFAULT_RESPONSE_PROMPT,
+        model_context: ChatCompletionContext | None = None,
     ) -> None:
         super().__init__(name=name, description=description)
         self._team = team
         self._model_client = model_client
         self._instruction = instruction
         self._response_prompt = response_prompt
 
+        if model_context is not None:
+            self._model_context = model_context
+        else:
+            self._model_context = UnboundedChatCompletionContext()
+
     @property
     def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
         return (TextMessage,)
 
+    @property
+    def model_context(self) -> ChatCompletionContext:
+        """
+        The model context in use by the agent.
+        """
+        return self._model_context
+
     async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
         # Call the stream method and collect the messages.
         response: Response | None = None
@@ -138,18 +169,35 @@ async def on_messages_stream(
         self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
     ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
         # Prepare the task for the team of agents.
-        task = list(messages)
+        task_messages = list(messages)
 
         # Run the team of agents.
         result: TaskResult | None = None
         inner_messages: List[BaseAgentEvent | BaseChatMessage] = []
+        model_context = self._model_context
         count = 0
+
+        prev_content = await model_context.get_messages()
+        if len(prev_content) > 0:
+            prev_message = HandoffMessage(
+                content="relevant previous messages",
+                source=self.name,
+                target="",
+                context=prev_content,
+            )
+            task_messages = [prev_message] + task_messages
+
+        if len(task_messages) == 0:
+            task = None
+        else:
+            task = task_messages
+
         async for inner_msg in self._team.run_stream(task=task, cancellation_token=cancellation_token):
             if isinstance(inner_msg, TaskResult):
                 result = inner_msg
             else:
                 count += 1
-                if count <= len(task):
+                if count <= len(task_messages):
                     # Skip the task messages.
                     continue
                 yield inner_msg
@@ -161,27 +209,51 @@ async def on_messages_stream(
 
         if len(inner_messages) == 0:
             yield Response(
-                chat_message=TextMessage(source=self.name, content="No response."), inner_messages=inner_messages
+                chat_message=TextMessage(source=self.name, content="No response."),
+                inner_messages=[],
+                # Response's inner_messages should be empty. Cause that mean is response to outer world.
             )
         else:
             # Generate a response using the model client.
             llm_messages: List[LLMMessage] = [SystemMessage(content=self._instruction)]
-            for message in messages:
+            for message in inner_messages:
                 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)
             assert isinstance(completion.content, str)
             yield Response(
                 chat_message=TextMessage(source=self.name, content=completion.content, models_usage=completion.usage),
-                inner_messages=inner_messages,
+                inner_messages=[],
+                # Response's inner_messages should be empty. Cause that mean is response to outer world.
             )
 
+        # Add new user/handoff messages to the model context
+        await self._add_messages_to_context(
+            model_context=model_context,
+            messages=messages,
+        )
+
         # Reset the team.
         await self._team.reset()
 
+    @staticmethod
+    async def _add_messages_to_context(
+        model_context: ChatCompletionContext,
+        messages: Sequence[BaseChatMessage],
+    ) -> None:
+        """
+        Add incoming messages to the model context.
+        """
+        for msg in messages:
+            if isinstance(msg, HandoffMessage):
+                for llm_msg in msg.context:
+                    await model_context.add_message(llm_msg)
+            await model_context.add_message(msg.to_model_message())
+
     async def on_reset(self, cancellation_token: CancellationToken) -> None:
         await self._team.reset()
+        await self._model_context.clear()
 
     async def save_state(self) -> Mapping[str, Any]:
         team_state = await self._team.save_state()

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

@@ -31,11 +31,9 @@ async def test_society_of_mind_agent(runtime: AgentRuntime | None) -> None:
     inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime)
     society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client)
     response = await society_of_mind_agent.run(task="Count to 10.")
-    assert len(response.messages) == 4
+    assert len(response.messages) == 2
     assert response.messages[0].source == "user"
-    assert response.messages[1].source == "assistant1"
-    assert response.messages[2].source == "assistant2"
-    assert response.messages[3].source == "society_of_mind"
+    assert response.messages[1].source == "society_of_mind"
 
     # Test save and load state.
     state = await society_of_mind_agent.save_state()
@@ -57,3 +55,61 @@ async def test_society_of_mind_agent(runtime: AgentRuntime | None) -> None:
     loaded_soc_agent = SocietyOfMindAgent.load_component(soc_agent_config)
     assert isinstance(loaded_soc_agent, SocietyOfMindAgent)
     assert loaded_soc_agent.name == "society_of_mind"
+
+
+@pytest.mark.asyncio
+async def test_society_of_mind_agent_empty_messges(runtime: AgentRuntime | None) -> None:
+    model_client = ReplayChatCompletionClient(
+        ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"],
+    )
+    agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.")
+    agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.")
+    inner_termination = MaxMessageTermination(3)
+    inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime)
+    society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client)
+    response = await society_of_mind_agent.run()
+    assert len(response.messages) == 1
+    assert response.messages[0].source == "society_of_mind"
+
+
+@pytest.mark.asyncio
+async def test_society_of_mind_agent_no_response(runtime: AgentRuntime | None) -> None:
+    model_client = ReplayChatCompletionClient(
+        ["1", "2", "3"],
+    )
+    agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.")
+    agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.")
+    inner_termination = MaxMessageTermination(1)  # Set to 1 to force no response.
+    inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime)
+    society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client)
+    response = await society_of_mind_agent.run(task="Count to 10.")
+    assert len(response.messages) == 2
+    assert response.messages[0].source == "user"
+    assert response.messages[1].source == "society_of_mind"
+    assert response.messages[1].to_text() == "No response."
+
+
+@pytest.mark.asyncio
+async def test_society_of_mind_agent_multiple_rounds(runtime: AgentRuntime | None) -> None:
+    model_client = ReplayChatCompletionClient(
+        ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"],
+    )
+    agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.")
+    agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.")
+    inner_termination = MaxMessageTermination(3)
+    inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination, runtime=runtime)
+    society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client)
+    response = await society_of_mind_agent.run(task="Count to 10.")
+    assert len(response.messages) == 2
+    assert response.messages[0].source == "user"
+    assert response.messages[1].source == "society_of_mind"
+
+    # Continue.
+    response = await society_of_mind_agent.run()
+    assert len(response.messages) == 1
+    assert response.messages[0].source == "society_of_mind"
+
+    # Continue.
+    response = await society_of_mind_agent.run()
+    assert len(response.messages) == 1
+    assert response.messages[0].source == "society_of_mind"