from typing import List, Union

from pydantic import BaseModel, Field

from memgpt.schemas.enums import MessageStreamStatus
from memgpt.schemas.memgpt_message import MemGPTMessage
from memgpt.schemas.message import Message
from memgpt.schemas.usage import MemGPTUsageStatistics

# TODO: consider moving into own file


class MemGPTResponse(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 `MemGPTMessage`, depending on what was specified in the request.

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

    messages: Union[List[Message], List[MemGPTMessage]] = Field(..., description="The messages returned by the agent.")
    usage: MemGPTUsageStatistics = Field(..., description="The usage statistics of the agent.")


# The streaming response is either [DONE], [DONE_STEP], [DONE], an error, or a MemGPTMessage
MemGPTStreamingResponse = Union[MemGPTMessage, MessageStreamStatus]
