|
5 | 5 |
|
6 | 6 | from abc import ABC, abstractmethod |
7 | 7 | from collections.abc import AsyncGenerator, Callable |
8 | | -from typing import TYPE_CHECKING, Any |
| 8 | +from typing import TYPE_CHECKING, Any, Type, Union |
9 | 9 |
|
10 | | -from pydantic import ValidationError |
| 10 | +from pydantic import BaseModel, ValidationError |
11 | 11 |
|
12 | 12 | from a2a.auth.user import UnauthenticatedUser |
13 | 13 | from a2a.auth.user import User as A2AUser |
@@ -156,7 +156,21 @@ class JSONRPCApplication(ABC): |
156 | 156 | """ |
157 | 157 |
|
158 | 158 | # Method-to-model mapping for centralized routing |
159 | | - METHOD_TO_MODEL = { |
| 159 | + # Define the union type for all supported request models |
| 160 | + A2ARequestModel = Union[ |
| 161 | + SendMessageRequest, |
| 162 | + SendStreamingMessageRequest, |
| 163 | + GetTaskRequest, |
| 164 | + CancelTaskRequest, |
| 165 | + SetTaskPushNotificationConfigRequest, |
| 166 | + GetTaskPushNotificationConfigRequest, |
| 167 | + ListTaskPushNotificationConfigRequest, |
| 168 | + DeleteTaskPushNotificationConfigRequest, |
| 169 | + TaskResubscriptionRequest, |
| 170 | + GetAuthenticatedExtendedCardRequest, |
| 171 | + ] |
| 172 | + |
| 173 | + METHOD_TO_MODEL: dict[str, Type[A2ARequestModel]] = { |
160 | 174 | 'message/send': SendMessageRequest, |
161 | 175 | 'message/stream': SendStreamingMessageRequest, |
162 | 176 | 'tasks/get': GetTaskRequest, |
|
0 commit comments