chore(python-sdk): modernise type hints

Leverage __future__ import annotations to enable support for streamlined
type hints using | vs Union from PEP 604.

Future annotations was enabled across the codebase so that any
additional use of Optional or Union would be flagged by type linters.

Removed unused imports, guarded import under TYPE_CHECKING where
necessary.

Fixed accept parameter definition for EventEncoder and added a missing
comment.

Fixes: #50

Author: stevenh

python-sdk/ag_ui/core/__init__.py

@@ -2,6 +2,8 @@
 This module contains the core types and events for the Agent User Interaction Protocol.
 """
 
+from __future__ import annotations
+
 from ag_ui.core.events import (
     EventType,
     BaseEvent,

python-sdk/ag_ui/core/events.py

@@ -2,8 +2,10 @@
 This module contains the event types for the Agent User Interaction Protocol Python SDK.
 """
 
+from __future__ import annotations
+
 from enum import Enum
-from typing import Any, List, Literal, Optional, Union, Annotated
+from typing import Any, List, Literal, Optional, Annotated
 from pydantic import Field
 
 from .types import Message, State, ConfiguredBaseModel
@@ -38,8 +40,8 @@ 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
+    timestamp: int | None = None
+    raw_event: Any = None
 
 
 class TextMessageStartEvent(BaseEvent):
@@ -76,9 +78,9 @@ class TextMessageChunkEvent(BaseEvent):
     Event containing a chunk of text message content.
     """
     type: Literal[EventType.TEXT_MESSAGE_CHUNK]
-    message_id: Optional[str] = None
+    message_id: str | None = None
     role: Optional[Literal["assistant"]] = None
-    delta: Optional[str] = None
+    delta: str | None = None
 
 class ToolCallStartEvent(BaseEvent):
     """
@@ -87,7 +89,7 @@ class ToolCallStartEvent(BaseEvent):
     type: Literal[EventType.TOOL_CALL_START]
     tool_call_id: str
     tool_call_name: str
-    parent_message_id: Optional[str] = None
+    parent_message_id: str | None = None
 
 
 class ToolCallArgsEvent(BaseEvent):
@@ -111,10 +113,10 @@ class ToolCallChunkEvent(BaseEvent):
     Event containing a chunk of tool call content.
     """
     type: Literal[EventType.TOOL_CALL_CHUNK]
-    tool_call_id: Optional[str] = None
-    tool_call_name: Optional[str] = None
-    parent_message_id: Optional[str] = None
-    delta: Optional[str] = None
+    tool_call_id: str | None = None
+    tool_call_name: str | None = None
+    parent_message_id: str | None = None
+    delta: str | None = None
 
 class StateSnapshotEvent(BaseEvent):
     """
@@ -146,7 +148,7 @@ class RawEvent(BaseEvent):
     """
     type: Literal[EventType.RAW]
     event: Any
-    source: Optional[str] = None
+    source: str | None = None
 
 
 class CustomEvent(BaseEvent):
@@ -182,7 +184,7 @@ class RunErrorEvent(BaseEvent):
     """
     type: Literal[EventType.RUN_ERROR]
     message: str
-    code: Optional[str] = None
+    code: str | None = None
 
 
 class StepStartedEvent(BaseEvent):
@@ -202,25 +204,23 @@ class StepFinishedEvent(BaseEvent):
 
 
 Event = Annotated[
-    Union[
-        TextMessageStartEvent,
-        TextMessageContentEvent,
-        TextMessageEndEvent,
-        TextMessageChunkEvent,
-        ToolCallStartEvent,
-        ToolCallArgsEvent,
-        ToolCallEndEvent,
-        ToolCallChunkEvent,
-        StateSnapshotEvent,
-        StateDeltaEvent,
-        MessagesSnapshotEvent,
-        RawEvent,
-        CustomEvent,
-        RunStartedEvent,
-        RunFinishedEvent,
-        RunErrorEvent,
-        StepStartedEvent,
-        StepFinishedEvent,
-    ],
+    TextMessageStartEvent |
+    TextMessageContentEvent |
+    TextMessageEndEvent |
+    TextMessageChunkEvent |
+    ToolCallStartEvent |
+    ToolCallArgsEvent |
+    ToolCallEndEvent |
+    ToolCallChunkEvent |
+    StateSnapshotEvent |
+    StateDeltaEvent |
+    MessagesSnapshotEvent |
+    RawEvent |
+    CustomEvent |
+    RunStartedEvent |
+    RunFinishedEvent |
+    RunErrorEvent |
+    StepStartedEvent |
+    StepFinishedEvent,
     Field(discriminator="type")
 ]

python-sdk/ag_ui/core/types.py

@@ -2,7 +2,9 @@
 This module contains the types for the Agent User Interaction Protocol Python SDK.
 """
 
-from typing import Any, List, Literal, Optional, Union, Annotated
+from __future__ import annotations
+
+from typing import Any, List, Literal, Optional, Annotated
 from pydantic import BaseModel, Field, ConfigDict
 from pydantic.alias_generators import to_camel
 
@@ -41,8 +43,8 @@ class BaseMessage(ConfiguredBaseModel):
     """
     id: str
     role: str
-    content: Optional[str] = None
-    name: Optional[str] = None
+    content: str | None = None
+    name: str | None = None
 
 
 class DeveloperMessage(BaseMessage):
@@ -88,7 +90,7 @@ class ToolMessage(ConfiguredBaseModel):
 
 
 Message = Annotated[
-    Union[DeveloperMessage, SystemMessage, AssistantMessage, UserMessage, ToolMessage],
+    DeveloperMessage | SystemMessage | AssistantMessage | UserMessage | ToolMessage,
     Field(discriminator="role")
 ]
 

python-sdk/ag_ui/encoder/__init__.py

@@ -2,6 +2,8 @@
 This module contains the EventEncoder class.
 """
 
+from __future__ import annotations
+
 from ag_ui.encoder.encoder import EventEncoder, AGUI_MEDIA_TYPE
 
 __all__ = ["EventEncoder", "AGUI_MEDIA_TYPE"]

python-sdk/ag_ui/encoder/encoder.py

@@ -2,16 +2,25 @@
 This module contains the EventEncoder class
 """
 
-from ag_ui.core.events import BaseEvent
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ag_ui.core.events import BaseEvent
+
 
 AGUI_MEDIA_TYPE = "application/vnd.ag-ui.event+proto"
 
 class EventEncoder:
     """
     Encodes Agent User Interaction events.
     """
-    def __init__(self, accept: str = None):
-        pass
+    def __init__(self, accept: str | None = None) -> None:
+        """
+        Initializes the EventEncoder.
+        """
+        self.accept = accept
 
     def get_content_type(self) -> str:
         """

python-sdk/tests/test_encoder.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import unittest
 import json
 from datetime import datetime

python-sdk/tests/test_events.py

@@ -1,9 +1,10 @@
+from __future__ import annotations
+
 import unittest
-import json
 from datetime import datetime
 from pydantic import ValidationError, TypeAdapter
 
-from ag_ui.core.types import Message, UserMessage, AssistantMessage, FunctionCall, ToolCall
+from ag_ui.core.types import UserMessage, AssistantMessage, FunctionCall, ToolCall
 from ag_ui.core.events import (
     EventType,
     BaseEvent,

python-sdk/tests/test_types.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import unittest
 from pydantic import ValidationError
 from pydantic import TypeAdapter