# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.

from typing import List, Union, Optional
from typing_extensions import Literal, Annotated, TypeAlias

from ..._utils import PropertyInfo
from ..._models import BaseModel
from .user_message import UserMessage
from .system_message import SystemMessage
from ..stop_reason_type import StopReasonType
from .assistant_message import AssistantMessage
from .reasoning_message import ReasoningMessage
from .tool_call_message import ToolCallMessage
from ..tool_return_message import ToolReturnMessage
from .approval_request_message import ApprovalRequestMessage
from .hidden_reasoning_message import HiddenReasoningMessage
from .approval_response_message import ApprovalResponseMessage

__all__ = ["LettaStreamingResponse", "LettaPing", "LettaErrorMessage", "LettaStopReason", "LettaUsageStatistics"]


class LettaPing(BaseModel):
    message_type: Literal["ping"]
    """The type of the message."""


class LettaErrorMessage(BaseModel):
    error_type: str
    """The type of error."""

    message: str
    """The error message."""

    message_type: Literal["error_message"]
    """The type of the message."""

    run_id: str
    """The ID of the run."""

    detail: Optional[str] = None
    """An optional error detail."""


class LettaStopReason(BaseModel):
    stop_reason: StopReasonType
    """The reason why execution stopped."""

    message_type: Optional[Literal["stop_reason"]] = None
    """The type of the message."""


class LettaUsageStatistics(BaseModel):
    completion_tokens: Optional[int] = None
    """The number of tokens generated by the agent."""

    message_type: Optional[Literal["usage_statistics"]] = None

    prompt_tokens: Optional[int] = None
    """The number of tokens in the prompt."""

    run_ids: Optional[List[str]] = None
    """The background task run IDs associated with the agent interaction"""

    step_count: Optional[int] = None
    """The number of steps taken by the agent."""

    total_tokens: Optional[int] = None
    """The total number of tokens processed by the agent."""


LettaStreamingResponse: TypeAlias = Annotated[
    Union[
        SystemMessage,
        UserMessage,
        ReasoningMessage,
        HiddenReasoningMessage,
        ToolCallMessage,
        ToolReturnMessage,
        AssistantMessage,
        ApprovalRequestMessage,
        ApprovalResponseMessage,
        LettaPing,
        LettaErrorMessage,
        LettaStopReason,
        LettaUsageStatistics,
    ],
    PropertyInfo(discriminator="message_type"),
]
