Fix pydantic warnings when saving state (#4801)

microsoft · Dec 23, 2024 · 8e116fd · 8e116fd
1 parent 9c8877e
commit 8e116fd
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/state/_states.py b/python/packages/autogen-agentchat/src/autogen_agentchat/state/_states.py
@@ -1,4 +1,4 @@
-from typing import Any, List, Mapping, Optional
+from typing import Annotated, Any, List, Mapping, Optional
 
 from pydantic import BaseModel, Field
 
@@ -7,6 +7,9 @@
     ChatMessage,
 )
 
+# Ensures pydantic can distinguish between types of events & messages.
+_AgentMessage = Annotated[AgentEvent | ChatMessage, Field(discriminator="type")]
+
 
 class BaseState(BaseModel):
     """Base class for all saveable state"""
@@ -33,7 +36,7 @@ class TeamState(BaseState):
 class BaseGroupChatManagerState(BaseState):
     """Base state for all group chat managers."""
 
-    message_thread: List[AgentEvent | ChatMessage] = Field(default_factory=list)
+    message_thread: List[_AgentMessage] = Field(default_factory=list)
     current_turn: int = Field(default=0)
     type: str = Field(default="BaseGroupChatManagerState")