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


class ChatQuery(BaseModel):
    project: str = Field(
        title="Project",
        description="The project that the chat is associated with",
    )
    version: Optional[str] = Field(
        title="Version",
        description="The version of the chat template",
    )
    inputs: Optional[Dict[str, Any]] = Field(
        title="Inputs",
        description="The inputs to fill into the chat template",
    )
    provider: str = Field(
        title="Provider",
        description="The provider that the chat is associated with",
    )
    model: str = Field(
        title="Model",
        description="The model that is used to generate the text",
    )
    hyperparameters: Optional[Dict[str, Any]] = Field(
        title="Hyperparameters",
        description="The hyperparameters that are used to generate the text",
    )
    messages: List[Dict[str, Any]] = Field(
        title="Messages",
        description="The messages that are used to generate the text",
    )
    functions: Optional[List[Dict[str, Any]]] = Field(
        title="Functions",
        description="The functions that are provided to the model to use",
    )
    function_call: Optional[str] = Field(
        title="Function Call",
        description="The function call mode that is used to generate the text",
    )
    plugins: Optional[List[str]] = Field(
        title="Plugins",
        description="The plugins that are used to generate the text",
    )
    source: Optional[str] = Field(
        title="Source",
        description="The source of the chat completion",
    )
    num_samples: Optional[int] = Field(
        title="Number of Samples",
        description="The number of samples to generate",
    )
    stream: Optional[bool] = Field(
        title="Streaming",
        description="Whether the chat completion is to be streamed",
    )
    user: Optional[str] = Field(
        title="User",
        description="The user that generated the text",
    )
    metadata: Optional[Dict[str, Any]] = Field(
        title="Metadata",
        description="The metadata to track for the chat completion",
    )
    ground_truth: Optional[str] = Field(
        title="Ground Truth",
        description="The ground truth for the chat completion",
    )

class ChatResponse(BaseModel):
    timestamp: str = Field(
        title="Timestamp",
        description="The timestamp of the chat completion",
    )
    generation_id: str = Field(
        title="Generation ID",
        description="The unique ID of the generation",
    )
    generation: str = Field(
        title="Generation",
        description="The text that was generated by your model provider",
    )
    metrics: Optional[Dict[str, Any]] = Field(
        title="Metrics",
        description="The metrics associated with the generation",
    )
    completion_tokens: int = Field(
        title="Completion Tokens",
        description="The number of tokens generated for the completion",
    )
    choices: List[Dict[str, Any]] = Field(
        title="Choices",
        description="The choices associated with the generation",
    )
    cost: float = Field(
        title="Cost",
        description="The cost of the generation",
    )
    latency: float = Field(
        title="Latency",
        description="The latency of the generation in milliseconds",
    )