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 (
@@ -42,11 +42,9 @@
     UserMessage,
     ToolMessage,
     Message,
-    Role,
     Context,
     Tool,
     RunAgentInput,
-    State
 )
 
 __all__ = [
@@ -88,9 +86,8 @@
     "UserMessage",
     "ToolMessage",
     "Message",
-    "Role",
+    "TextMessageRole",
     "Context",
     "Tool",
     "RunAgentInput",
-    "State"
 ]

python-sdk/ag_ui/core/events.py

@@ -3,11 +3,11 @@
 """
 
 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, Role
+from .types import ConfiguredBaseModel, Message, StateT, JSONValue
 
 # Text messages can have any role except "tool"
 TextMessageRole = Literal["developer", "system", "assistant", "user"]
@@ -49,7 +49,7 @@ class BaseEvent(ConfiguredBaseModel):
     """
     type: EventType
     timestamp: Optional[int] = None
-    raw_event: Optional[Any] = None
+    raw_event: Optional[JSONValue] = None
 
 
 class TextMessageStartEvent(BaseEvent):
@@ -164,20 +164,20 @@ class ThinkingEndEvent(BaseEvent):
     """
     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: List[JSONValue]  # JSON Patch (RFC 6902)
 
 
 class MessagesSnapshotEvent(BaseEvent):
@@ -193,7 +193,7 @@ 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
 
 
@@ -203,7 +203,7 @@ class CustomEvent(BaseEvent):
     """
     type: Literal[EventType.CUSTOM] = EventType.CUSTOM  # pyright: ignore[reportIncompatibleVariableOverride]
     name: str
-    value: Any
+    value: JSONValue
 
 
 class RunStartedEvent(BaseEvent):
@@ -222,7 +222,7 @@ class RunFinishedEvent(BaseEvent):
     type: Literal[EventType.RUN_FINISHED] = EventType.RUN_FINISHED  # pyright: ignore[reportIncompatibleVariableOverride]
     thread_id: str
     run_id: str
-    result: Optional[Any] = None
+    result: Optional[JSONValue] = None
 
 
 class RunErrorEvent(BaseEvent):

python-sdk/ag_ui/core/types.py

@@ -2,16 +2,22 @@
 This module contains the types for the Agent User Interaction Protocol Python SDK.
 """
 
-from typing import Annotated, Any, List, Literal, Optional, Union
+from typing import Annotated, Any, List, Literal, Optional, Union, Generic
+from typing_extensions import TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic.alias_generators import to_camel
 
+JSONValue = Union[str, int, float, bool, None, dict[str, Any], list[Any]]
+StateT = TypeVar("StateT", default=JSONValue, contravariant=True)
+FwdPropsT = TypeVar("FwdPropsT", default=JSONValue, contravariant=True)
+
 
 class ConfiguredBaseModel(BaseModel):
     """
     A configurable base model.
     """
+
     model_config = ConfigDict(
         extra="forbid",
         alias_generator=to_camel,
@@ -23,6 +29,7 @@ class FunctionCall(ConfiguredBaseModel):
     """
     Name and arguments of a function call.
     """
+
     name: str
     arguments: str
 
@@ -31,6 +38,7 @@ class ToolCall(ConfiguredBaseModel):
     """
     A tool call, modelled after OpenAI tool calls.
     """
+
     id: str
     type: Literal["function"] = "function"  # pyright: ignore[reportIncompatibleVariableOverride]
     function: FunctionCall
@@ -40,6 +48,7 @@ class BaseMessage(ConfiguredBaseModel):
     """
     A base message, modelled after OpenAI messages.
     """
+
     id: str
     role: str
     content: Optional[str] = None
@@ -50,22 +59,23 @@ class DeveloperMessage(BaseMessage):
     """
     A developer message.
     """
+
     role: Literal["developer"] = "developer"  # pyright: ignore[reportIncompatibleVariableOverride]
-    content: str
 
 
 class SystemMessage(BaseMessage):
     """
     A system message.
     """
+
     role: Literal["system"] = "system"  # pyright: ignore[reportIncompatibleVariableOverride]
-    content: str
 
 
 class AssistantMessage(BaseMessage):
     """
     An assistant message.
     """
+
     role: Literal["assistant"] = "assistant"  # pyright: ignore[reportIncompatibleVariableOverride]
     tool_calls: Optional[List[ToolCall]] = None
 
@@ -74,14 +84,15 @@ class UserMessage(BaseMessage):
     """
     A user message.
     """
-    role: Literal["user"] = "user" # pyright: ignore[reportIncompatibleVariableOverride]
-    content: str
+
+    role: Literal["user"] = "user"  # pyright: ignore[reportIncompatibleVariableOverride]
 
 
 class ToolMessage(ConfiguredBaseModel):
     """
     A tool result message.
     """
+
     id: str
     role: Literal["tool"] = "tool"
     content: str
@@ -91,16 +102,15 @@ class ToolMessage(ConfiguredBaseModel):
 
 Message = Annotated[
     Union[DeveloperMessage, SystemMessage, AssistantMessage, UserMessage, ToolMessage],
-    Field(discriminator="role")
+    Field(discriminator="role"),
 ]
 
-Role = Literal["developer", "system", "assistant", "user", "tool"]
-
 
 class Context(ConfiguredBaseModel):
     """
     Additional context for the agent.
     """
+
     description: str
     value: str
 
@@ -109,23 +119,21 @@ class Tool(ConfiguredBaseModel):
     """
     A tool definition.
     """
+
     name: str
     description: str
     parameters: Any  # JSON Schema for the tool parameters
 
 
-class RunAgentInput(ConfiguredBaseModel):
+class RunAgentInput(ConfiguredBaseModel, Generic[StateT, FwdPropsT]):
     """
     Input for running an agent.
     """
+
     thread_id: str
     run_id: str
-    state: Any
+    state: StateT
     messages: List[Message]
     tools: List[Tool]
     context: List[Context]
-    forwarded_props: Any
-
-
-# State can be any type
-State = Any
+    forwarded_props: FwdPropsT

python-sdk/ag_ui/encoder/encoder.py

@@ -2,15 +2,19 @@
 This module contains the EventEncoder class
 """
 
+from typing import Union
+
 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):
+
+    def __init__(self, accept: Union[str, None] = None):
         pass
 
     def get_content_type(self) -> str:

python-sdk/pyproject.toml

@@ -1,15 +1,26 @@
-[tool.poetry]
-name = "ag-ui-protocol"
-version = "0.1.8"
+[project]
+name = "ag-ui"
+version = "0.1.9"
 description = ""
-authors = ["Markus Ecker <[email protected]>"]
+authors = [
+    { name = "Markus Ecker", email = "[email protected]" },
+]
 readme = "README.md"
-packages = [{include = "ag_ui", from = "."}]
-[tool.poetry.dependencies]
-python = "^3.9"
-pydantic = "^2.11.2"
-
+requires-python = ">=3.9,<4.0"
+dependencies = [
+    "pydantic>=2.11.2,<3.0.0",
+]
+packages = [
+    "ag_ui"
+]
 
 [build-system]
-requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "mypy>=1.17.1",
+    "pyright>=1.1.403",
+    "ruff>=0.12.9",
+]