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 <ekzhu@users.noreply.github.com>
This commit is contained in:
EeS 2025-04-05 05:50:50 +09:00 committed by GitHub
parent 687946258f
commit 39321266f9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 137 additions and 9 deletions

View File

@ -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 ..base import TaskResult, Team
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,6 +132,7 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
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
@ -121,10 +140,22 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
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 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
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,12 +209,14 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
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))
@ -174,14 +224,36 @@ class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
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()

View File

@ -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"