Rebased

Author: wdoppenberg

python-sdk/ag_ui/core/__init__.py

@@ -42,7 +42,6 @@
     UserMessage,
     ToolMessage,
     Message,
-    Role,
     Context,
     Tool,
     RunAgentInput,
@@ -87,7 +86,7 @@
     "UserMessage",
     "ToolMessage",
     "Message",
-    "Role",
+    "TextMessageRole",
     "Context",
     "Tool",
     "RunAgentInput",

python-sdk/ag_ui/core/events.py

@@ -7,11 +7,10 @@
 
 from pydantic import Field
 
-from .types import ConfiguredBaseModel, Message, State, Role
+from .types import ConfiguredBaseModel, Message, StateT, JSONValue
 
 # Text messages can have any role except "tool"
 TextMessageRole = Literal["developer", "system", "assistant", "user"]
-from .types import ConfiguredBaseModel, Message, AgentStateT, JSONValue
 
 
 class EventType(str, Enum):
@@ -197,13 +196,13 @@ class ThinkingEndEvent(BaseEvent):
     type: Literal[EventType.THINKING_END] = EventType.THINKING_END  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
-class StateSnapshotEvent(BaseEvent, Generic[AgentStateT]):
+class StateSnapshotEvent(BaseEvent, Generic[StateT]):
     """
     Event containing a snapshot of the state.
     """
 
     type: Literal[EventType.STATE_SNAPSHOT] = EventType.STATE_SNAPSHOT  # pyright: ignore[reportIncompatibleVariableOverride]
-    snapshot: AgentStateT
+    snapshot: StateT
 
 
 class StateDeltaEvent(BaseEvent):

python-sdk/ag_ui/core/types.py

@@ -105,8 +105,6 @@ class ToolMessage(ConfiguredBaseModel):
     Field(discriminator="role"),
 ]
 
-Role = Literal["developer", "system", "assistant", "user", "tool"]
-
 
 class Context(ConfiguredBaseModel):
     """

python-sdk/tests/test_text_roles.py

@@ -8,7 +8,6 @@
     TextMessageContentEvent,
     TextMessageEndEvent,
     TextMessageChunkEvent,
-    Role,
 )
 
 # Test all available roles for text messages (excluding "tool")