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

from typing import List, Optional
from typing_extensions import Literal

from .message import Message
from ..._models import BaseModel
from ..stop_reason_type import StopReasonType

__all__ = [
    "LettaResponse",
    "StopReason",
    "Usage",
    "Logprobs",
    "LogprobsContent",
    "LogprobsContentTopLogprob",
    "LogprobsRefusal",
    "LogprobsRefusalTopLogprob",
    "Turn",
]


class StopReason(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 Usage(BaseModel):
    """The usage statistics of the agent."""

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


class LogprobsContentTopLogprob(BaseModel):
    token: str

    logprob: float

    bytes: Optional[List[int]] = None


class LogprobsContent(BaseModel):
    token: str

    logprob: float

    top_logprobs: List[LogprobsContentTopLogprob]

    bytes: Optional[List[int]] = None


class LogprobsRefusalTopLogprob(BaseModel):
    token: str

    logprob: float

    bytes: Optional[List[int]] = None


class LogprobsRefusal(BaseModel):
    token: str

    logprob: float

    top_logprobs: List[LogprobsRefusalTopLogprob]

    bytes: Optional[List[int]] = None


class Logprobs(BaseModel):
    """Log probabilities of the output tokens from the last LLM call.

    Only present if return_logprobs was enabled.
    """

    content: Optional[List[LogprobsContent]] = None

    refusal: Optional[List[LogprobsRefusal]] = None


class Turn(BaseModel):
    """Token data for a single LLM generation turn in a multi-turn agent interaction.

    Used for RL training to track token IDs and logprobs across all LLM calls,
    not just the final one. Tool results are included so the client can tokenize
    them with loss_mask=0 (non-trainable).
    """

    role: Literal["assistant", "tool"]
    """
    Role of this turn: 'assistant' for LLM generations (trainable), 'tool' for tool
    results (non-trainable).
    """

    content: Optional[str] = None
    """Text content. For tool turns, client tokenizes this with loss_mask=0."""

    output_ids: Optional[List[int]] = None
    """Token IDs from SGLang native endpoint. Only present for assistant turns."""

    output_token_logprobs: Optional[List[List[object]]] = None
    """Logprobs from SGLang: [[logprob, token_id, top_logprob_or_null], ...].

    Only present for assistant turns.
    """

    tool_name: Optional[str] = None
    """Name of the tool called. Only present for tool turns."""


class LettaResponse(BaseModel):
    """
    Response object from an agent interaction, consisting of the new messages generated by the agent and usage statistics.
    The type of the returned messages can be either `Message` or `LettaMessage`, depending on what was specified in the request.

    Attributes:
        messages (List[Union[Message, LettaMessage]]): The messages returned by the agent.
        usage (LettaUsageStatistics): The usage statistics
    """

    messages: List[Message]
    """The messages returned by the agent."""

    stop_reason: StopReason
    """The stop reason from Letta indicating why agent loop stopped execution."""

    usage: Usage
    """The usage statistics of the agent."""

    logprobs: Optional[Logprobs] = None
    """Log probabilities of the output tokens from the last LLM call.

    Only present if return_logprobs was enabled.
    """

    turns: Optional[List[Turn]] = None
    """Token data for all LLM generations in multi-turn agent interaction.

    Includes token IDs and logprobs for each assistant turn, plus tool result
    content. Only present if return_token_ids was enabled. Used for RL training with
    loss masking.
    """
