from pydantic import BaseModel, Field

from ..data_models import OllamaEndpointResponseBody
from .message_models import Message


class Function(BaseModel):
    name: str = Field(None, description="The name of the function to call.")
    arguments: str = Field(None, description="The arguments to pass to the function.")


class ToolCall(BaseModel):
    function: Function = Field(None, description="The function call details.")


class OllamaStreamChatCompletionResponseBody(OllamaEndpointResponseBody):
    model: str = Field(None, description="The model name")

    created_at: str = Field(
        None, description="The timestamp when the response was created"
    )

    message: Message = Field(
        None, description="he partial or full response generated by the model"
    )

    done: bool = Field(
        None,
        description="A flag indicating whether the response generation is complete",
    )

    done_reason: str = Field(
        None, description="The response generation complete reason"
    )


class OllamaChatCompletionResponseBody(OllamaStreamChatCompletionResponseBody):
    total_duration: int = Field(None, description="Time spent generating the response")

    load_duration: int = Field(
        None, description="Time spent in nanoseconds loading the model"
    )

    prompt_eval_count: int = Field(None, description="number of tokens in the prompt")

    prompt_eval_duration: int = Field(
        None, description="time spent in nanoseconds evaluating the prompt"
    )

    eval_count: int = Field(None, description="number of tokens in the response")

    eval_duration: int = Field(
        None, description="time in nanoseconds spent generating the response"
    )
