Python generics for State & FwdProps; poetry -> uv

.github/workflows/test.yml

@@ -2,9 +2,13 @@ name: test
 
 on:
   push:
-    branches: main
+    branches: [
+      "main"
+    ]
   pull_request:
-    branches: main
+    branches: [
+      "main"
+    ]
 
 jobs:
   python:

python-sdk/.gitignore

@@ -63,3 +63,4 @@ venv.bak/
 
 # Project specific
 .DS_Store
+/uv.lock

python-sdk/.pre-commit-config.yaml

@@ -0,0 +1,30 @@
+repos:
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: v0.12.9
+    hooks:
+      - id: ruff
+        args: [
+          --fix
+        ]
+        files: ^python-sdk/
+      - id: ruff-format
+        files: ^python-sdk/
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.17.1
+    hooks:
+      - id: mypy
+        files: ^python-sdk/
+        args: [
+          --python-version=3.9,
+          --disallow-untyped-calls,
+          --disallow-untyped-defs,
+          --disallow-incomplete-defs,
+          --check-untyped-defs,
+          --no-implicit-optional,
+          --warn-redundant-casts,
+          --ignore-missing-imports,
+        ]
+        additional_dependencies:
+          - "types-pytz"
+        exclude_types: [ jupyter ]
+        exclude: "tests"

python-sdk/ag_ui/core/__init__.py

@@ -29,7 +29,7 @@
     RunErrorEvent,
     StepStartedEvent,
     StepFinishedEvent,
-    Event
+    Event,
 )
 
 from ag_ui.core.types import (
@@ -46,7 +46,6 @@
     Context,
     Tool,
     RunAgentInput,
-    State
 )
 
 __all__ = [
@@ -92,5 +91,4 @@
     "Context",
     "Tool",
     "RunAgentInput",
-    "State"
 ]

python-sdk/ag_ui/core/events.py

@@ -3,17 +3,18 @@
 """
 
 from enum import Enum
-from typing import Annotated, Any, List, Literal, Optional, Union
+from typing import Annotated, List, Literal, Optional, Union, Generic
 
 from pydantic import Field
 
-from .types import ConfiguredBaseModel, Message, State
+from .types import ConfiguredBaseModel, Message, StateT, JSONValue
 
 
 class EventType(str, Enum):
     """
     The type of event.
     """
+
     TEXT_MESSAGE_START = "TEXT_MESSAGE_START"
     TEXT_MESSAGE_CONTENT = "TEXT_MESSAGE_CONTENT"
     TEXT_MESSAGE_END = "TEXT_MESSAGE_END"
@@ -44,15 +45,17 @@ class BaseEvent(ConfiguredBaseModel):
     """
     Base event for all events in the Agent User Interaction Protocol.
     """
+
     type: EventType
     timestamp: Optional[int] = None
-    raw_event: Optional[Any] = None
+    raw_event: Optional[JSONValue] = None
 
 
 class TextMessageStartEvent(BaseEvent):
     """
     Event indicating the start of a text message.
     """
+
     type: Literal[EventType.TEXT_MESSAGE_START] = EventType.TEXT_MESSAGE_START  # pyright: ignore[reportIncompatibleVariableOverride]
     message_id: str
     role: Literal["assistant"] = "assistant"
@@ -62,6 +65,7 @@ class TextMessageContentEvent(BaseEvent):
     """
     Event containing a piece of text message content.
     """
+
     type: Literal[EventType.TEXT_MESSAGE_CONTENT] = EventType.TEXT_MESSAGE_CONTENT  # pyright: ignore[reportIncompatibleVariableOverride]
     message_id: str
     delta: str = Field(min_length=1)
@@ -71,41 +75,58 @@ class TextMessageEndEvent(BaseEvent):
     """
     Event indicating the end of a text message.
     """
+
     type: Literal[EventType.TEXT_MESSAGE_END] = EventType.TEXT_MESSAGE_END  # pyright: ignore[reportIncompatibleVariableOverride]
     message_id: str
 
+
 class TextMessageChunkEvent(BaseEvent):
     """
     Event containing a chunk of text message content.
     """
+
     type: Literal[EventType.TEXT_MESSAGE_CHUNK] = EventType.TEXT_MESSAGE_CHUNK  # pyright: ignore[reportIncompatibleVariableOverride]
     message_id: Optional[str] = None
     role: Optional[Literal["assistant"]] = None
     delta: Optional[str] = None
 
+
 class ThinkingTextMessageStartEvent(BaseEvent):
     """
     Event indicating the start of a thinking text message.
     """
-    type: Literal[EventType.THINKING_TEXT_MESSAGE_START] = EventType.THINKING_TEXT_MESSAGE_START  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    type: Literal[EventType.THINKING_TEXT_MESSAGE_START] = (
+        EventType.THINKING_TEXT_MESSAGE_START
+    )  # pyright: ignore[reportIncompatibleVariableOverride]
+
 
 class ThinkingTextMessageContentEvent(BaseEvent):
     """
     Event indicating a piece of a thinking text message.
     """
-    type: Literal[EventType.THINKING_TEXT_MESSAGE_CONTENT] = EventType.THINKING_TEXT_MESSAGE_CONTENT  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    type: Literal[EventType.THINKING_TEXT_MESSAGE_CONTENT] = (
+        EventType.THINKING_TEXT_MESSAGE_CONTENT
+    )  # pyright: ignore[reportIncompatibleVariableOverride]
     delta: str = Field(min_length=1)
 
+
 class ThinkingTextMessageEndEvent(BaseEvent):
     """
     Event indicating the end of a thinking text message.
     """
-    type: Literal[EventType.THINKING_TEXT_MESSAGE_END] = EventType.THINKING_TEXT_MESSAGE_END  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    type: Literal[EventType.THINKING_TEXT_MESSAGE_END] = (
+        EventType.THINKING_TEXT_MESSAGE_END
+    )  # pyright: ignore[reportIncompatibleVariableOverride]
+
 
 class ToolCallStartEvent(BaseEvent):
     """
     Event indicating the start of a tool call.
     """
+
     type: Literal[EventType.TOOL_CALL_START] = EventType.TOOL_CALL_START  # pyright: ignore[reportIncompatibleVariableOverride]
     tool_call_id: str
     tool_call_name: str
@@ -116,6 +137,7 @@ class ToolCallArgsEvent(BaseEvent):
     """
     Event containing tool call arguments.
     """
+
     type: Literal[EventType.TOOL_CALL_ARGS] = EventType.TOOL_CALL_ARGS  # pyright: ignore[reportIncompatibleVariableOverride]
     tool_call_id: str
     delta: str
@@ -125,62 +147,75 @@ class ToolCallEndEvent(BaseEvent):
     """
     Event indicating the end of a tool call.
     """
+
     type: Literal[EventType.TOOL_CALL_END] = EventType.TOOL_CALL_END  # pyright: ignore[reportIncompatibleVariableOverride]
     tool_call_id: str
 
+
 class ToolCallChunkEvent(BaseEvent):
     """
     Event containing a chunk of tool call content.
     """
+
     type: Literal[EventType.TOOL_CALL_CHUNK] = EventType.TOOL_CALL_CHUNK  # pyright: ignore[reportIncompatibleVariableOverride]
     tool_call_id: Optional[str] = None
     tool_call_name: Optional[str] = None
     parent_message_id: Optional[str] = None
     delta: Optional[str] = None
 
+
 class ToolCallResultEvent(BaseEvent):
     """
     Event containing the result of a tool call.
     """
+
     message_id: str
     type: Literal[EventType.TOOL_CALL_RESULT] = EventType.TOOL_CALL_RESULT  # pyright: ignore[reportIncompatibleVariableOverride]
     tool_call_id: str
     content: str
     role: Optional[Literal["tool"]] = None
 
+
 class ThinkingStartEvent(BaseEvent):
     """
     Event indicating the start of a thinking step event.
     """
+
     type: Literal[EventType.THINKING_START] = EventType.THINKING_START  # pyright: ignore[reportIncompatibleVariableOverride]
     title: Optional[str] = None
 
+
 class ThinkingEndEvent(BaseEvent):
     """
     Event indicating the end of a thinking step event.
     """
+
     type: Literal[EventType.THINKING_END] = EventType.THINKING_END  # pyright: ignore[reportIncompatibleVariableOverride]
 
-class StateSnapshotEvent(BaseEvent):
+
+class StateSnapshotEvent(BaseEvent, Generic[StateT]):
     """
     Event containing a snapshot of the state.
     """
+
     type: Literal[EventType.STATE_SNAPSHOT] = EventType.STATE_SNAPSHOT  # pyright: ignore[reportIncompatibleVariableOverride]
-    snapshot: State
+    snapshot: StateT
 
 
 class StateDeltaEvent(BaseEvent):
     """
     Event containing a delta of the state.
     """
+
     type: Literal[EventType.STATE_DELTA] = EventType.STATE_DELTA  # pyright: ignore[reportIncompatibleVariableOverride]
-    delta: List[Any]  # JSON Patch (RFC 6902)
+    delta: JSONValue  # JSON Patch (RFC 6902)
 
 
 class MessagesSnapshotEvent(BaseEvent):
     """
     Event containing a snapshot of the messages.
     """
+
     type: Literal[EventType.MESSAGES_SNAPSHOT] = EventType.MESSAGES_SNAPSHOT  # pyright: ignore[reportIncompatibleVariableOverride]
     messages: List[Message]
 
@@ -189,24 +224,27 @@ class RawEvent(BaseEvent):
     """
     Event containing a raw event.
     """
+
     type: Literal[EventType.RAW] = EventType.RAW  # pyright: ignore[reportIncompatibleVariableOverride]
-    event: Any
+    event: JSONValue
     source: Optional[str] = None
 
 
 class CustomEvent(BaseEvent):
     """
     Event containing a custom event.
     """
+
     type: Literal[EventType.CUSTOM] = EventType.CUSTOM  # pyright: ignore[reportIncompatibleVariableOverride]
     name: str
-    value: Any
+    value: JSONValue
 
 
 class RunStartedEvent(BaseEvent):
     """
     Event indicating that a run has started.
     """
+
     type: Literal[EventType.RUN_STARTED] = EventType.RUN_STARTED  # pyright: ignore[reportIncompatibleVariableOverride]
     thread_id: str
     run_id: str
@@ -216,16 +254,18 @@ class RunFinishedEvent(BaseEvent):
     """
     Event indicating that a run has finished.
     """
+
     type: Literal[EventType.RUN_FINISHED] = EventType.RUN_FINISHED  # pyright: ignore[reportIncompatibleVariableOverride]
     thread_id: str
     run_id: str
-    result: Optional[Any] = None
+    result: JSONValue = None
 
 
 class RunErrorEvent(BaseEvent):
     """
     Event indicating that a run has encountered an error.
     """
+
     type: Literal[EventType.RUN_ERROR] = EventType.RUN_ERROR  # pyright: ignore[reportIncompatibleVariableOverride]
     message: str
     code: Optional[str] = None
@@ -235,6 +275,7 @@ class StepStartedEvent(BaseEvent):
     """
     Event indicating that a step has started.
     """
+
     type: Literal[EventType.STEP_STARTED] = EventType.STEP_STARTED  # pyright: ignore[reportIncompatibleVariableOverride]
     step_name: str
 
@@ -243,6 +284,7 @@ class StepFinishedEvent(BaseEvent):
     """
     Event indicating that a step has finished.
     """
+
     type: Literal[EventType.STEP_FINISHED] = EventType.STEP_FINISHED  # pyright: ignore[reportIncompatibleVariableOverride]
     step_name: str
 
@@ -269,5 +311,5 @@ class StepFinishedEvent(BaseEvent):
         StepStartedEvent,
         StepFinishedEvent,
     ],
-    Field(discriminator="type")
+    Field(discriminator="type"),
 ]