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

from typing import List, Union, Optional
from datetime import datetime
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):
    """
    A ping message used as a keepalive to prevent SSE streams from timing out during long running requests.

    Args:
        id (str): The ID of the message
        date (datetime): The date the message was created in ISO format
    """

    id: str

    date: datetime

    is_err: Optional[bool] = None

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

    Ping messages are a keep-alive to prevent SSE streams from timing out during
    long running requests.
    """

    name: Optional[str] = None

    otid: Optional[str] = None

    run_id: Optional[str] = None

    sender_id: Optional[str] = None

    seq_id: Optional[int] = None

    step_id: Optional[str] = None


class LettaErrorMessage(BaseModel):
    """
    Error messages are used to notify the client of an error that occurred during the agent's execution.
    """

    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."""

    seq_id: Optional[int] = None
    """The sequence ID for cursor-based pagination."""


class LettaStopReason(BaseModel):
    """The stop reason from Letta indicating why agent loop stopped execution."""

    stop_reason: StopReasonType
    """The reason why execution stopped."""

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


class LettaUsageStatistics(BaseModel):
    """Usage statistics for the agent interaction.

    Attributes:
        completion_tokens (int): The number of tokens generated by the agent.
        prompt_tokens (int): The number of tokens in the prompt.
        total_tokens (int): The total number of tokens processed by the agent.
        step_count (int): The number of steps taken by the agent.
        cached_input_tokens (Optional[int]): The number of input tokens served from cache. None if not reported.
        cache_write_tokens (Optional[int]): The number of input tokens written to cache. None if not reported.
        reasoning_tokens (Optional[int]): The number of reasoning/thinking tokens generated. None if not reported.
    """

    cache_write_tokens: Optional[int] = None
    """The number of input tokens written to cache (Anthropic only).

    None if not reported by provider.
    """

    cached_input_tokens: Optional[int] = None
    """The number of input tokens served from cache. None if not reported by provider."""

    completion_tokens: Optional[int] = None
    """The number of tokens generated by the agent."""

    context_tokens: Optional[int] = None
    """Estimate of tokens currently in the context window."""

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

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

    reasoning_tokens: Optional[int] = None
    """The number of reasoning/thinking tokens generated.

    None if not reported by provider.
    """

    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"),
]
