# This file was auto-generated by Fern from our API Definition.

from __future__ import annotations
from ....core.unchecked_base_model import UncheckedBaseModel
import typing
import datetime as dt
from ....core.pydantic_utilities import IS_PYDANTIC_V2
import pydantic
from ....types.tool_call_message_tool_call import ToolCallMessageToolCall
from ....types.tool_return_message_status import ToolReturnMessageStatus
import typing_extensions
from ....core.unchecked_base_model import UnionMetadata


class LettaStreamingResponse_SystemMessage(UncheckedBaseModel):
    message_type: typing.Literal["system_message"] = "system_message"
    id: str
    date: dt.datetime
    message: str

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow


class LettaStreamingResponse_UserMessage(UncheckedBaseModel):
    message_type: typing.Literal["user_message"] = "user_message"
    id: str
    date: dt.datetime
    message: str

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow


class LettaStreamingResponse_ReasoningMessage(UncheckedBaseModel):
    message_type: typing.Literal["reasoning_message"] = "reasoning_message"
    id: str
    date: dt.datetime
    reasoning: str

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow


class LettaStreamingResponse_ToolCallMessage(UncheckedBaseModel):
    message_type: typing.Literal["tool_call_message"] = "tool_call_message"
    id: str
    date: dt.datetime
    tool_call: ToolCallMessageToolCall

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow


class LettaStreamingResponse_ToolReturnMessage(UncheckedBaseModel):
    message_type: typing.Literal["tool_return_message"] = "tool_return_message"
    id: str
    date: dt.datetime
    tool_return: str
    status: ToolReturnMessageStatus
    tool_call_id: str
    stdout: typing.Optional[typing.List[str]] = None
    stderr: typing.Optional[typing.List[str]] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow


class LettaStreamingResponse_AssistantMessage(UncheckedBaseModel):
    message_type: typing.Literal["assistant_message"] = "assistant_message"
    id: str
    date: dt.datetime
    assistant_message: str

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow


class LettaStreamingResponse_UsageStatistics(UncheckedBaseModel):
    message_type: typing.Literal["usage_statistics"] = "usage_statistics"
    completion_tokens: typing.Optional[int] = None
    prompt_tokens: typing.Optional[int] = None
    total_tokens: typing.Optional[int] = None
    step_count: typing.Optional[int] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow


LettaStreamingResponse = typing_extensions.Annotated[
    typing.Union[
        LettaStreamingResponse_SystemMessage,
        LettaStreamingResponse_UserMessage,
        LettaStreamingResponse_ReasoningMessage,
        LettaStreamingResponse_ToolCallMessage,
        LettaStreamingResponse_ToolReturnMessage,
        LettaStreamingResponse_AssistantMessage,
        LettaStreamingResponse_UsageStatistics,
    ],
    UnionMetadata(discriminant="message_type"),
]
