Relevant A2A types

relevant-a2a-types.py

# ------------------------
# Base JSON-RPC Components
# ------------------------

class JSONRPCMessage(BaseModel):
    jsonrpc: Literal["2.0"] = "2.0"
    id: int | str | None = Field(default_factory=lambda: uuid4().hex)

class JSONRPCRequest(JSONRPCMessage):
    pass

class JSONRPCError(BaseModel):
    code: int
    message: str
    data: Any | None = None

class JSONRPCResponse(JSONRPCMessage):
    result: Any | None = None
    error: JSONRPCError | None = None

# ------------------------
# Enums
# ------------------------

class Role(Enum):
    user = "user"
    agent = "agent"

class TaskState(str, Enum):
    submitted = "submitted"
    working = "working"
    input_required = "input-required"
    completed = "completed"
    canceled = "canceled"
    failed = "failed"
    unknown = "unknown"
    rejected = "rejected"
    auth_required = "auth_required"

# ------------------------
# Message Parts
# ------------------------

class TextPart(BaseModel):
    kind: Literal["text"] = "text"
    text: str
    metadata: dict[str, Any] | None = None

class FileContent(BaseModel):
    model_config = ConfigDict(populate_by_name=True)
    name: str | None = None
    mime_type: str | None = Field(default=None, alias="mimeType")
    bytes: str | None = None
    uri: str | None = None

    @model_validator(mode="after")
    def check_content(self) -> Self:
        if not (self.bytes or self.uri):
            raise ValueError("Either 'bytes' or 'uri' must be present in the file data")
        if self.bytes and self.uri:
            raise ValueError("Only one of 'bytes' or 'uri' can be present")
        return self

class FilePart(BaseModel):
    kind: Literal["file"] = "file"
    file: FileContent
    metadata: dict[str, Any] | None = None

class DataPart(BaseModel):
    kind: Literal["data"] = "data"
    data: dict[str, Any]
    metadata: dict[str, Any] | None = None

Part = Annotated[TextPart | FilePart | DataPart, Field(discriminator="kind")]

# ------------------------
# Message & Configuration
# ------------------------

class PushNotificationAuthenticationInfo(BaseModel):
    schemes: list[str]
    credentials: str | None = None

class PushNotificationConfig(BaseModel):
    url: str
    token: str | None = None
    authentication: PushNotificationAuthenticationInfo | None = None

class MessageSendConfiguration(BaseModel):
    accepted_output_modes: list[str] = Field(alias="acceptedOutputModes")
    history_length: int | None = Field(default=0, alias="historyLength")
    push_notification_config: PushNotificationConfig | None = Field(
        default=None, alias="pushNotificationConfig"
    )
    blocking: bool | None = None

class Message(BaseModel):
    model_config = ConfigDict(populate_by_name=True)
    kind: Literal["message"] = "message"
    role: Literal["user", "agent"]
    parts: list[Part]
    metadata: dict[str, Any] | None = None
    message_id: str = Field(alias="messageId")
    context_id: str | None = Field(default=None, alias="contextId")
    task_id: str | None = Field(default=None, alias="taskId")

class MessageSendParams(BaseModel):
    model_config = ConfigDict(populate_by_name=True)
    message: Message
    configuration: MessageSendConfiguration | None = None
    metadata: dict[str, Any] | None = None

# ------------------------
# Artifacts and Task State
# ------------------------

class Artifact(BaseModel):
    name: str | None = None
    description: str | None = None
    parts: list[Part]
    metadata: dict[str, Any] | None = None
    index: int = 0
    append: bool | None = None
    last_chunk: bool | None = Field(default=None, alias="lastChunk")

class TaskStatus(BaseModel):
    state: TaskState
    message: Message | None = None
    timestamp: datetime = Field(default_factory=datetime.now)

    @field_serializer("timestamp")
    def serialize_dt(self, dt: datetime) -> str:
        return dt.isoformat()

class Task(BaseModel):
    id: str
    kind: Literal["task"] = "task"
    context_id: str | None = Field(default=None, alias="contextId")
    status: TaskStatus
    artifacts: list[Artifact] | None = None
    history: list[Message] | None = None
    metadata: dict[str, Any] | None = None

# ------------------------
# Event Updates (streaming)
# ------------------------

class TaskStatusUpdateEvent(BaseModel):
    id: str
    status: TaskStatus
    final: bool = False
    metadata: dict[str, Any] | None = None

class TaskArtifactUpdateEvent(BaseModel):
    id: str
    artifact: Artifact
    metadata: dict[str, Any] | None = None

# ------------------------
# Final Request & Response Types
# ------------------------

class SendMessageRequest(JSONRPCRequest):
    method: Literal["message/send"] = "message/send"
    params: MessageSendParams

class StreamMessageRequest(JSONRPCRequest):
    method: Literal["message/stream"] = "message/stream"
    params: MessageSendParams

class SendMessageResponse(JSONRPCResponse):
    result: Task | Message | None = None

class SendStreamingMessageResponse(JSONRPCResponse):
    result: Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None = None