Added generics for state & forwarded props; poetry -> uv; linting

Changelog: changed

Author: wdoppenberg

python-sdk/.gitignore

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

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

@@ -0,0 +1,41 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: trailing-whitespace
+        exclude_types: [ jupyter ]
+      - id: end-of-file-fixer
+        exclude_types: [ jupyter ]
+      - id: check-docstring-first
+      - id: debug-statements
+      - id: check-ast
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: v0.11.8
+    hooks:
+      - id: ruff
+        args: [
+          --fix
+        ]
+      - id: ruff-format
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.15.0
+    hooks:
+      - id: mypy
+        args: [
+          --python-version=3.12,
+          --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"
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.8.1
+    hooks:
+      - id: nbstripout

python-sdk/ag_ui/__init__.py

@@ -46,7 +46,6 @@
     Context,
     Tool,
     RunAgentInput,
-    State
 )
 
 __all__ = [
@@ -92,5 +91,4 @@
     "Context",
     "Tool",
     "RunAgentInput",
-    "State"
 ]

python-sdk/ag_ui/core/__init__.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
+from .types import ConfiguredBaseModel, Message, AgentStateT, JSONValue
 
 
 class EventType(str, Enum):
@@ -46,7 +46,7 @@ class BaseEvent(ConfiguredBaseModel):
     """
     type: EventType
     timestamp: Optional[int] = None
-    raw_event: Optional[Any] = None
+    raw_event: Optional[JSONValue] = None
 
 
 class TextMessageStartEvent(BaseEvent):
@@ -161,20 +161,20 @@ class ThinkingEndEvent(BaseEvent):
     """
     type: Literal[EventType.THINKING_END] = EventType.THINKING_END  # pyright: ignore[reportIncompatibleVariableOverride]
 
-class StateSnapshotEvent(BaseEvent):
+class StateSnapshotEvent(BaseEvent, Generic[AgentStateT]):
     """
     Event containing a snapshot of the state.
     """
     type: Literal[EventType.STATE_SNAPSHOT] = EventType.STATE_SNAPSHOT  # pyright: ignore[reportIncompatibleVariableOverride]
-    snapshot: State
+    snapshot: AgentStateT
 
 
 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):
@@ -190,7 +190,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
 
 
@@ -200,7 +200,7 @@ class CustomEvent(BaseEvent):
     """
     type: Literal[EventType.CUSTOM] = EventType.CUSTOM  # pyright: ignore[reportIncompatibleVariableOverride]
     name: str
-    value: Any
+    value: JSONValue
 
 
 class RunStartedEvent(BaseEvent):
@@ -219,7 +219,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: JSONValue = None
 
 
 class RunErrorEvent(BaseEvent):

python-sdk/ag_ui/core/events.py

@@ -2,11 +2,16 @@
 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]]
+AgentStateT = TypeVar('AgentStateT', default=JSONValue, contravariant=True)
+FwdPropsT = TypeVar('FwdPropsT', default=JSONValue, contravariant=True)
+
 
 class ConfiguredBaseModel(BaseModel):
     """
@@ -51,15 +56,13 @@ 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):
@@ -75,7 +78,6 @@ class UserMessage(BaseMessage):
     A user message.
     """
     role: Literal["user"] = "user" # pyright: ignore[reportIncompatibleVariableOverride]
-    content: str
 
 
 class ToolMessage(ConfiguredBaseModel):
@@ -114,18 +116,14 @@ class Tool(ConfiguredBaseModel):
     parameters: Any  # JSON Schema for the tool parameters
 
 
-class RunAgentInput(ConfiguredBaseModel):
+class RunAgentInput(ConfiguredBaseModel, Generic[AgentStateT, FwdPropsT]):
     """
     Input for running an agent.
     """
     thread_id: str
     run_id: str
-    state: Any
+    state: AgentStateT
     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/core/types.py

@@ -10,7 +10,7 @@ class EventEncoder:
     """
     Encodes Agent User Interaction events.
     """
-    def __init__(self, accept: str = None):
+    def __init__(self, accept: str | None = None):
         pass
 
     def get_content_type(self) -> str:

python-sdk/ag_ui/encoder/encoder.py

@@ -1,15 +1,30 @@
-[tool.poetry]
+[project]
 name = "ag-ui-protocol"
-version = "0.1.8"
+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 = [
+    { include = "ag_ui", from = "ag_ui" }
+]
+
+[tool.hatch.build.targets.wheel]
+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",
+]

python-sdk/pyproject.toml

@@ -1,9 +1,8 @@
 import unittest
-import json
 from datetime import datetime
-from pydantic import ValidationError, TypeAdapter
+from pydantic import 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,