mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-17 14:01:06 +00:00
025490a1bd
This PR refactored `AgentEvent` and `ChatMessage` union types to abstract base classes. This allows for user-defined message types that subclass one of the base classes to be used in AgentChat. To support a unified interface for working with the messages, the base classes added abstract methods for: - Convert content to string - Convert content to a `UserMessage` for model client - Convert content for rendering in console. - Dump into a dictionary - Load and create a new instance from a dictionary This way, all agents such as `AssistantAgent` and `SocietyOfMindAgent` can utilize the unified interface to work with any built-in and user-defined message type. This PR also introduces a new message type, `StructuredMessage` for AgentChat (Resolves #5131), which is a generic type that requires a user-specified content type. You can create a `StructuredMessage` as follow: ```python class MessageType(BaseModel): data: str references: List[str] message = StructuredMessage[MessageType](content=MessageType(data="data", references=["a", "b"]), source="user") # message.content is of type `MessageType`. ``` This PR addresses the receving side of this message type. To produce this message type from `AssistantAgent`, the work continue in #5934. Added unit tests to verify this message type works with agents and teams.
94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
import pytest
|
|
from autogen_agentchat.messages import HandoffMessage, MessageFactory, StructuredMessage, TextMessage
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class TestContent(BaseModel):
|
|
"""Test content model."""
|
|
|
|
field1: str
|
|
field2: int
|
|
|
|
|
|
def test_structured_message() -> None:
|
|
# Create a structured message with the test content
|
|
message = StructuredMessage[TestContent](
|
|
source="test_agent",
|
|
content=TestContent(field1="test", field2=42),
|
|
)
|
|
|
|
# Check that the message type is correct
|
|
assert message.type == "StructuredMessage[TestContent]" # type: ignore
|
|
|
|
# Check that the content is of the correct type
|
|
assert isinstance(message.content, TestContent)
|
|
|
|
# Check that the content fields are set correctly
|
|
assert message.content.field1 == "test"
|
|
assert message.content.field2 == 42
|
|
|
|
# Check that model_dump works correctly
|
|
dumped_message = message.model_dump()
|
|
assert dumped_message["source"] == "test_agent"
|
|
assert dumped_message["content"]["field1"] == "test"
|
|
assert dumped_message["content"]["field2"] == 42
|
|
assert dumped_message["type"] == "StructuredMessage[TestContent]"
|
|
|
|
|
|
def test_message_factory() -> None:
|
|
factory = MessageFactory()
|
|
|
|
# Text message data
|
|
text_data = {
|
|
"type": "TextMessage",
|
|
"source": "test_agent",
|
|
"content": "Hello, world!",
|
|
}
|
|
|
|
# Create a TextMessage instance
|
|
text_message = factory.create(text_data)
|
|
assert isinstance(text_message, TextMessage)
|
|
assert text_message.source == "test_agent"
|
|
assert text_message.content == "Hello, world!"
|
|
assert text_message.type == "TextMessage" # type: ignore
|
|
|
|
# Handoff message data
|
|
handoff_data = {
|
|
"type": "HandoffMessage",
|
|
"source": "test_agent",
|
|
"content": "handoff to another agent",
|
|
"target": "target_agent",
|
|
}
|
|
|
|
# Create a HandoffMessage instance
|
|
handoff_message = factory.create(handoff_data)
|
|
assert isinstance(handoff_message, HandoffMessage)
|
|
assert handoff_message.source == "test_agent"
|
|
assert handoff_message.content == "handoff to another agent"
|
|
assert handoff_message.target == "target_agent"
|
|
assert handoff_message.type == "HandoffMessage" # type: ignore
|
|
|
|
# Structured message data
|
|
structured_data = {
|
|
"type": "StructuredMessage[TestContent]",
|
|
"source": "test_agent",
|
|
"content": {
|
|
"field1": "test",
|
|
"field2": 42,
|
|
},
|
|
}
|
|
# Create a StructuredMessage instance -- this will fail because the type
|
|
# is not registered in the factory.
|
|
with pytest.raises(ValueError):
|
|
structured_message = factory.create(structured_data)
|
|
# Register the StructuredMessage type in the factory
|
|
factory.register(StructuredMessage[TestContent])
|
|
# Create a StructuredMessage instance
|
|
structured_message = factory.create(structured_data)
|
|
assert isinstance(structured_message, StructuredMessage)
|
|
assert isinstance(structured_message.content, TestContent) # type: ignore
|
|
assert structured_message.source == "test_agent"
|
|
assert structured_message.content.field1 == "test"
|
|
assert structured_message.content.field2 == 42
|
|
assert structured_message.type == "StructuredMessage[TestContent]" # type: ignore
|