from typing import Optional

from pydantic import BaseModel, Field

from .chat_messages import AssistantMessage


class Choice(BaseModel):
    index: int = Field(
        ..., description="The index of the choice in the list of choices."
    )
    message: AssistantMessage = Field(
        ..., description="The message generated by the model."
    )
    logprobs: Optional[dict] = Field(
        None, description="The log probabilities of the tokens in the completion."
    )
    finish_reason: Optional[str] = Field(
        None, description="The reason the model stopped generating tokens."
    )


class Usage(BaseModel):
    prompt_tokens: int = Field(..., description="The number of tokens in the prompt.")
    completion_tokens: int = Field(
        ..., description="The number of tokens in the completion."
    )
    total_tokens: int = Field(..., description="The total number of tokens used.")


class ChatCompletionResponse(BaseModel):
    id: str = Field(..., description="The ID of the chat completion.")
    object: str = Field(..., description="The type of the object.")
    created: int = Field(..., description="The timestamp of the chat completion.")
    model: str = Field(
        ..., description="The ID of the model used to generate the chat completion."
    )
    choices: list[Choice] = Field(
        ..., description="The chat completion choices generated by the model."
    )
    usage: Usage = Field(
        ..., description="The token usage statistics for the chat completion."
    )
    system_fingerprint: Optional[str] = Field(
        None,
        description="A fingerprint that can be used to monitor changes in the backend.",
    )
