Python: added generic types to ChatOptions and ChatResponse/AgentResponse for Response Format (#3305)

* added generic types to ChatOptions and ChatResponse/AgentResponse for response format

* fix typevar import

* fix for older python versions

* fix missing import

* fixed imports

* fixed mypy

* mypy fix

python/packages/ag-ui/agent_framework_ag_ui/_client.py

@@ -28,25 +28,21 @@ from ._http_service import AGUIHttpService
 from ._message_adapters import agent_framework_messages_to_agui
 from ._utils import convert_tools_to_agui_format
 
-if TYPE_CHECKING:
-    from ._types import AGUIChatOptions
-
-from typing import TypedDict
-
 if sys.version_info >= (3, 13):
-    from typing import TypeVar
+    from typing import TypeVar  # type: ignore # pragma: no cover
 else:
-    from typing_extensions import TypeVar
-
+    from typing_extensions import TypeVar  # type: ignore # pragma: no cover
 if sys.version_info >= (3, 12):
     from typing import override  # type: ignore # pragma: no cover
 else:
     from typing_extensions import override  # type: ignore[import] # pragma: no cover
-
 if sys.version_info >= (3, 11):
-    from typing import Self  # pragma: no cover
+    from typing import Self, TypedDict  # pragma: no cover
 else:
-    from typing_extensions import Self  # pragma: no cover
+    from typing_extensions import Self, TypedDict  # pragma: no cover
+
+if TYPE_CHECKING:
+    from ._types import AGUIChatOptions
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -85,7 +81,7 @@ def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseCha
 
     @wraps(original_get_response)
     async def response_wrapper(self: Any, *args: Any, **kwargs: Any) -> ChatResponse:
-        response = await original_get_response(self, *args, **kwargs)
+        response: ChatResponse[Any] = await original_get_response(self, *args, **kwargs)  # type: ignore[var-annotated]
         if response.messages:
             for message in response.messages:
                 _unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], message.contents))

python/packages/ag-ui/agent_framework_ag_ui/_types.py

@@ -3,15 +3,23 @@
 """Type definitions for AG-UI integration."""
 
 import sys
-from typing import Any, TypedDict
+from typing import Any, Generic
 
 from agent_framework import ChatOptions
 from pydantic import BaseModel, Field
 
 if sys.version_info >= (3, 13):
-    from typing import TypeVar
+    from typing import TypeVar  # type: ignore # pragma: no cover
 else:
-    from typing_extensions import TypeVar
+    from typing_extensions import TypeVar  # type: ignore # pragma: no cover
+if sys.version_info >= (3, 11):
+    from typing import TypedDict  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import TypedDict  # type: ignore # pragma: no cover
+
+
+TAGUIChatOptions = TypeVar("TAGUIChatOptions", bound=TypedDict, default="AGUIChatOptions", covariant=True)  # type: ignore[valid-type]
+TResponseModel = TypeVar("TResponseModel", bound=BaseModel | None, default=None)
 
 
 class PredictStateConfig(TypedDict):
@@ -76,7 +84,7 @@ class AGUIRequest(BaseModel):
 # region AG-UI Chat Options TypedDict
 
 
-class AGUIChatOptions(ChatOptions, total=False):
+class AGUIChatOptions(ChatOptions[TResponseModel], Generic[TResponseModel], total=False):
     """AG-UI protocol-specific chat options dict.
 
     Extends base ChatOptions for the AG-UI (Agent-UI) protocol.
@@ -140,7 +148,5 @@ class AGUIChatOptions(ChatOptions, total=False):
 AGUI_OPTION_TRANSLATIONS: dict[str, str] = {}
 """Maps ChatOptions keys to AG-UI parameter names (protocol uses standard names)."""
 
-TAGUIChatOptions = TypeVar("TAGUIChatOptions", bound=TypedDict, default="AGUIChatOptions", covariant=True)  # type: ignore[valid-type]
-
 
 # endregion

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py

@@ -12,6 +12,10 @@ if sys.version_info >= (3, 13):
     from typing import TypeVar  # type: ignore # pragma: no cover
 else:
     from typing_extensions import TypeVar  # type: ignore # pragma: no cover
+if sys.version_info >= (3, 11):
+    from typing import TypedDict  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import TypedDict  # type: ignore # pragma: no cover
 
 if TYPE_CHECKING:
     from agent_framework import ChatOptions