from pydantic import BaseModel, Field
from typing import Optional, List, Literal, Dict


class LogProbs(BaseModel):
    text_offset: List[int] = Field(default_factory=list)
    token_logprobs: List[float] = Field(default_factory=list)
    tokens: List[str] = Field(default_factory=list)
    top_logprobs: Optional[List[Dict[str, float]]] = Field(default_factory=list)


class UsageInfo(BaseModel):
    """Usage statistics.

    Attributes:
      prompt_tokens (int): The number of tokens in the prompt.
      total_tokens (int): The total number of tokens used in the request (prompt + completion).
      completion_tokens (Optional[int]): The number of tokens in the generated completion.
    """

    prompt_tokens: int
    total_tokens: int
    completion_tokens: Optional[int] = None


class Choice(BaseModel):
    """A completion choice.

    Attributes:
      index (int): The index of the completion choice.
      text (str): The completion response.
      logprobs (float, optional): The log probabilities of the most likely tokens.
      finish_reason (str): The reason the model stopped generating tokens. This will be "stop" if
        the model hit a natural stop point or a provided stop sequence, or
        "length" if the maximum number of tokens specified in the request was
        reached.
    """

    index: int
    text: str
    logprobs: Optional[LogProbs] = None
    finish_reason: Optional[Literal["stop", "length", "error"]] = None


class CompletionResponse(BaseModel):
    """The response message from a /v1/completions call.

    Attributes:
      id (str): A unique identifier of the response.
      object (str): The object type, which is always "text_completion".
      created (int): The Unix time in seconds when the response was generated.
      choices (List[Choice]): The list of generated completion choices.
    """

    id: str
    object: str
    created: int
    model: str
    choices: List[Choice]
    usage: UsageInfo


class CompletionResponseStreamChoice(BaseModel):
    """A streamed completion choice.

    Attributes:
      index (int): The index of the completion choice.
      text (str): The completion response.
      logprobs (float, optional): The log probabilities of the most likely tokens.
      finish_reason (str): The reason the model stopped generating tokens. This will be "stop" if
        the model hit a natural stop point or a provided stop sequence, or
        "length" if the maximum number of tokens specified in the request was
        reached.
    """

    index: int
    text: str
    logprobs: Optional[LogProbs] = None
    finish_reason: Optional[Literal["stop", "length", "error"]] = None


class CompletionStreamResponse(BaseModel):
    """The streamed response message from a /v1/completions call.

    Attributes:
      id (str): A unique identifier of the response.
      object (str): The object type, which is always "text_completion".
      created (int): The Unix time in seconds when the response was generated.
      model (str): The model used for the chat completion.
      choices (List[CompletionResponseStreamChoice]):
        The list of streamed completion choices.
    """

    id: str
    object: str
    created: int
    model: str
    choices: List[CompletionResponseStreamChoice]
    usage: Optional[UsageInfo] = None


class Model(BaseModel):
    """A model deployed to the Fireworks platform.

    Attributes:
      id (str): The model name.
      object (str): The object type, which is always "model".
      created (int): The Unix time in seconds when the model was generated.
    """

    id: str
    object: str
    created: int


class ListModelsResponse(BaseModel):
    """The response message from a /v1/models call.

    Attributes:
      object (str): The object type, which is always "list".
      data (List[Model]): The list of models.
    """

    object: str
    data: List[Model]


class ChatMessage(BaseModel):
    """A chat completion message.

    Attributes:
      role (str): The role of the author of this message.
      content (str): The contents of the message.
    """

    role: str
    content: str


class ChatCompletionResponseChoice(BaseModel):
    """A chat completion choice generated by a chat model.

    Attributes:
      index (int): The index of the chat completion choice.
      message (ChatMessage): The chat completion message.
      finish_reason (Optional[str]): The reason the model stopped generating tokens. This will be "stop" if
        the model hit a natural stop point or a provided stop sequence, or
        "length" if the maximum number of tokens specified in the request was
        reached.
    """

    index: int
    message: ChatMessage
    finish_reason: Optional[str] = None


class ChatCompletionResponse(BaseModel):
    """The response message from a /v1/chat/completions call.

    Attributes:
      id (str): A unique identifier of the response.
      object (str): The object type, which is always "chat.completion".
      created (int): The Unix time in seconds when the response was generated.
      model (str): The model used for the chat completion.
      choices (List[ChatCompletionResponseChoice]): The list of chat completion choices.
      usage (UsageInfo): Usage statistics for the chat completion.
    """

    id: str
    object: str
    created: int
    model: str
    choices: List[ChatCompletionResponseChoice]
    usage: Optional[UsageInfo] = None


class DeltaMessage(BaseModel):
    """A message delta.

    Attributes:
      role (str): The role of the author of this message.
      content (str): The contents of the chunk message.
    """

    role: Optional[str] = None
    content: Optional[str] = None


class ChatCompletionResponseStreamChoice(BaseModel):
    """A streamed chat completion choice.

    Attributes:
      index (int): The index of the chat completion choice.
      delta (DeltaMessage): The message delta.
      finish_reason (str): The reason the model stopped generating tokens. This will be "stop" if
        the model hit a natural stop point or a provided stop sequence, or
        "length" if the maximum number of tokens specified in the request was
        reached.
    """

    index: int
    delta: DeltaMessage
    finish_reason: Optional[Literal["stop", "length"]] = None


class ChatCompletionStreamResponse(BaseModel):
    """The streamed response message from a /v1/chat/completions call.

    Attributes:
      id (str): A unique identifier of the response.
      object (str): The object type, which is always "chat.completion".
      created (int): The Unix time in seconds when the response was generated.
      model (str): The model used for the chat completion.
      choices (List[ChatCompletionResponseStreamChoice]):
        The list of streamed chat completion choices.
    """

    id: str
    object: str
    created: int
    model: str
    choices: List[ChatCompletionResponseStreamChoice]
    usage: Optional[UsageInfo] = None
