"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""
# @generated-id: c6becbbd80bc

from __future__ import annotations
from .activitytaskcompletedresponse import (
    ActivityTaskCompletedResponse,
    ActivityTaskCompletedResponseTypedDict,
)
from .activitytaskfailedresponse import (
    ActivityTaskFailedResponse,
    ActivityTaskFailedResponseTypedDict,
)
from .activitytaskretryingresponse import (
    ActivityTaskRetryingResponse,
    ActivityTaskRetryingResponseTypedDict,
)
from .activitytaskstartedresponse import (
    ActivityTaskStartedResponse,
    ActivityTaskStartedResponseTypedDict,
)
from .customtaskcanceledresponse import (
    CustomTaskCanceledResponse,
    CustomTaskCanceledResponseTypedDict,
)
from .customtaskcompletedresponse import (
    CustomTaskCompletedResponse,
    CustomTaskCompletedResponseTypedDict,
)
from .customtaskfailedresponse import (
    CustomTaskFailedResponse,
    CustomTaskFailedResponseTypedDict,
)
from .customtaskinprogressresponse import (
    CustomTaskInProgressResponse,
    CustomTaskInProgressResponseTypedDict,
)
from .customtaskstartedresponse import (
    CustomTaskStartedResponse,
    CustomTaskStartedResponseTypedDict,
)
from .customtasktimedoutresponse import (
    CustomTaskTimedOutResponse,
    CustomTaskTimedOutResponseTypedDict,
)
from .streameventworkflowcontext import (
    StreamEventWorkflowContext,
    StreamEventWorkflowContextTypedDict,
)
from .workflowexecutioncanceledresponse import (
    WorkflowExecutionCanceledResponse,
    WorkflowExecutionCanceledResponseTypedDict,
)
from .workflowexecutioncompletedresponse import (
    WorkflowExecutionCompletedResponse,
    WorkflowExecutionCompletedResponseTypedDict,
)
from .workflowexecutioncontinuedasnewresponse import (
    WorkflowExecutionContinuedAsNewResponse,
    WorkflowExecutionContinuedAsNewResponseTypedDict,
)
from .workflowexecutionfailedresponse import (
    WorkflowExecutionFailedResponse,
    WorkflowExecutionFailedResponseTypedDict,
)
from .workflowexecutionstartedresponse import (
    WorkflowExecutionStartedResponse,
    WorkflowExecutionStartedResponseTypedDict,
)
from .workflowtaskfailedresponse import (
    WorkflowTaskFailedResponse,
    WorkflowTaskFailedResponseTypedDict,
)
from .workflowtasktimedoutresponse import (
    WorkflowTaskTimedOutResponse,
    WorkflowTaskTimedOutResponseTypedDict,
)
from datetime import datetime
from mistralai.client.types import BaseModel, UNSET_SENTINEL
from pydantic import model_serializer
from typing import Any, Dict, Optional, Union
from typing_extensions import NotRequired, TypeAliasType, TypedDict


StreamEventSsePayloadDataTypedDict = TypeAliasType(
    "StreamEventSsePayloadDataTypedDict",
    Union[
        WorkflowExecutionStartedResponseTypedDict,
        WorkflowExecutionCompletedResponseTypedDict,
        WorkflowExecutionFailedResponseTypedDict,
        WorkflowExecutionCanceledResponseTypedDict,
        WorkflowExecutionContinuedAsNewResponseTypedDict,
        WorkflowTaskTimedOutResponseTypedDict,
        WorkflowTaskFailedResponseTypedDict,
        CustomTaskStartedResponseTypedDict,
        CustomTaskInProgressResponseTypedDict,
        CustomTaskCompletedResponseTypedDict,
        CustomTaskFailedResponseTypedDict,
        CustomTaskTimedOutResponseTypedDict,
        CustomTaskCanceledResponseTypedDict,
        ActivityTaskStartedResponseTypedDict,
        ActivityTaskCompletedResponseTypedDict,
        ActivityTaskRetryingResponseTypedDict,
        ActivityTaskFailedResponseTypedDict,
    ],
)


StreamEventSsePayloadData = TypeAliasType(
    "StreamEventSsePayloadData",
    Union[
        WorkflowExecutionStartedResponse,
        WorkflowExecutionCompletedResponse,
        WorkflowExecutionFailedResponse,
        WorkflowExecutionCanceledResponse,
        WorkflowExecutionContinuedAsNewResponse,
        WorkflowTaskTimedOutResponse,
        WorkflowTaskFailedResponse,
        CustomTaskStartedResponse,
        CustomTaskInProgressResponse,
        CustomTaskCompletedResponse,
        CustomTaskFailedResponse,
        CustomTaskTimedOutResponse,
        CustomTaskCanceledResponse,
        ActivityTaskStartedResponse,
        ActivityTaskCompletedResponse,
        ActivityTaskRetryingResponse,
        ActivityTaskFailedResponse,
    ],
)


class StreamEventSsePayloadTypedDict(TypedDict):
    stream: str
    data: StreamEventSsePayloadDataTypedDict
    workflow_context: StreamEventWorkflowContextTypedDict
    broker_sequence: int
    timestamp: NotRequired[datetime]
    metadata: NotRequired[Dict[str, Any]]


class StreamEventSsePayload(BaseModel):
    stream: str

    data: StreamEventSsePayloadData

    workflow_context: StreamEventWorkflowContext

    broker_sequence: int

    timestamp: Optional[datetime] = None

    metadata: Optional[Dict[str, Any]] = None

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = set(["timestamp", "metadata"])
        serialized = handler(self)
        m = {}

        for n, f in type(self).model_fields.items():
            k = f.alias or n
            val = serialized.get(k, serialized.get(n))

            if val != UNSET_SENTINEL:
                if val is not None or k not in optional_fields:
                    m[k] = val

        return m