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
@@ -23,15 +25,15 @@ def test_encode_method(self):
         # Create a test event
         timestamp = int(datetime.now().timestamp() * 1000)
         event = BaseEvent(type=EventType.RAW, timestamp=timestamp)
-        
+
         # Create encoder and encode event
         encoder = EventEncoder()
         encoded = encoder.encode(event)
-        
+
         # The encode method calls encode_sse, so the result should be in SSE format
         expected = f"data: {event.model_dump_json(by_alias=True, exclude_none=True)}\n\n"
         self.assertEqual(encoded, expected)
-        
+
         # Verify that camelCase is used in the encoded output
         self.assertIn('"type":', encoded)
         self.assertIn('"timestamp":', encoded)
@@ -48,25 +50,25 @@ def test_encode_sse_method(self):
             delta="Hello, world!",
             timestamp=1648214400000
         )
-        
+
         # Create encoder and encode event to SSE
         encoder = EventEncoder()
         encoded_sse = encoder._encode_sse(event)
-        
+
         # Verify the format is correct for SSE (data: [json]\n\n)
         self.assertTrue(encoded_sse.startswith("data: "))
         self.assertTrue(encoded_sse.endswith("\n\n"))
-        
+
         # Extract and verify the JSON content
         json_content = encoded_sse[6:-2]  # Remove "data: " prefix and "\n\n" suffix
         decoded = json.loads(json_content)
-        
+
         # Check that all fields were properly encoded
         self.assertEqual(decoded["type"], "TEXT_MESSAGE_CONTENT")
         self.assertEqual(decoded["messageId"], "msg_123")  # Check snake_case converted to camelCase
         self.assertEqual(decoded["delta"], "Hello, world!")
         self.assertEqual(decoded["timestamp"], 1648214400000)
-        
+
         # Verify that snake_case has been converted to camelCase
         self.assertIn("messageId", decoded)  # camelCase key exists
         self.assertNotIn("message_id", decoded)  # snake_case key doesn't exist
@@ -75,12 +77,12 @@ def test_encode_with_different_event_types(self):
         """Test encoding different types of events"""
         # Create encoder
         encoder = EventEncoder()
-        
+
         # Test with a basic BaseEvent
         base_event = BaseEvent(type=EventType.RAW, timestamp=1648214400000)
         encoded_base = encoder.encode(base_event)
         self.assertIn('"type":"RAW"', encoded_base)
-        
+
         # Test with a more complex event
         content_event = TextMessageContentEvent(
             type=EventType.TEXT_MESSAGE_CONTENT,
@@ -89,20 +91,20 @@ def test_encode_with_different_event_types(self):
             timestamp=1648214400000
         )
         encoded_content = encoder.encode(content_event)
-        
+
         # Verify correct encoding and camelCase conversion
         self.assertIn('"type":"TEXT_MESSAGE_CONTENT"', encoded_content)
         self.assertIn('"messageId":"msg_456"', encoded_content)  # Check snake_case converted to camelCase
         self.assertIn('"delta":"Testing different events"', encoded_content)
-        
+
         # Extract JSON and verify camelCase conversion
         json_content = encoded_content.split("data: ")[1].rstrip("\n\n")
         decoded = json.loads(json_content)
-        
+
         # Verify messageId is camelCase (not message_id)
         self.assertIn("messageId", decoded)
         self.assertNotIn("message_id", decoded)
-        
+
     def test_null_value_exclusion(self):
         """Test that fields with None values are excluded from the JSON output"""
         # Create an event with some fields set to None
@@ -111,22 +113,22 @@ def test_null_value_exclusion(self):
             timestamp=1648214400000,
             raw_event=None  # Explicitly set to None
         )
-        
+
         # Create encoder and encode event
         encoder = EventEncoder()
         encoded = encoder.encode(event)
-        
+
         # Extract JSON
         json_content = encoded.split("data: ")[1].rstrip("\n\n")
         decoded = json.loads(json_content)
-        
+
         # Verify fields that are present
         self.assertIn("type", decoded)
         self.assertIn("timestamp", decoded)
-        
+
         # Verify null fields are excluded
         self.assertNotIn("rawEvent", decoded)
-        
+
         # Test with another event that has optional fields
         # Create event with some optional fields set to None
         event_with_optional = ToolCallStartEvent(
@@ -136,18 +138,18 @@ def test_null_value_exclusion(self):
             parent_message_id=None,  # Optional field explicitly set to None
             timestamp=1648214400000
         )
-        
+
         encoded_optional = encoder.encode(event_with_optional)
         json_content_optional = encoded_optional.split("data: ")[1].rstrip("\n\n")
         decoded_optional = json.loads(json_content_optional)
-        
+
         # Required fields should be present
         self.assertIn("toolCallId", decoded_optional)
         self.assertIn("toolCallName", decoded_optional)
-        
+
         # Optional field with None value should be excluded
         self.assertNotIn("parentMessageId", decoded_optional)
-        
+
     def test_round_trip_serialization(self):
         """Test that events can be serialized to JSON with camelCase and deserialized back correctly"""
         # Create a complex event with multiple fields
@@ -158,10 +160,10 @@ def test_round_trip_serialization(self):
             parent_message_id="msg_parent_456",
             timestamp=1648214400000
         )
-        
+
         # Serialize to JSON with camelCase fields
         json_str = original_event.model_dump_json(by_alias=True)
-        
+
         # Verify JSON uses camelCase
         json_data = json.loads(json_str)
         self.assertIn("toolCallId", json_data)
@@ -170,19 +172,19 @@ def test_round_trip_serialization(self):
         self.assertNotIn("tool_call_id", json_data)
         self.assertNotIn("tool_call_name", json_data)
         self.assertNotIn("parent_message_id", json_data)
-        
+
         # Deserialize back to an event
         deserialized_event = ToolCallStartEvent.model_validate_json(json_str)
-        
+
         # Verify the deserialized event is equivalent to the original
         self.assertEqual(deserialized_event.type, original_event.type)
         self.assertEqual(deserialized_event.tool_call_id, original_event.tool_call_id)
         self.assertEqual(deserialized_event.tool_call_name, original_event.tool_call_name)
         self.assertEqual(deserialized_event.parent_message_id, original_event.parent_message_id)
         self.assertEqual(deserialized_event.timestamp, original_event.timestamp)
-        
+
         # Verify complete equality using model_dump
         self.assertEqual(
-            original_event.model_dump(), 
+            original_event.model_dump(),
             deserialized_event.model_dump()
         )