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


class GenerateQuery(BaseModel):
    task: str = Field(
        title="Task",
        description="The task for which the generation is being requested",
    )
    input: Dict[str, str] = Field(
        title="Input", description="The input variables for the generation"
    )
    prompts: Optional[List[str]] = Field(
        title="Prompt IDs",
        description="The unique IDs of the prompts to generate from"
        + ". The prompts need to be associated with the task.",
    )
    model_id: Optional[str] = Field(
        title="Model ID",
        description="The unique ID of the fine-tuned model to generate from",
    )
    best_of: Optional[int] = Field(
        title="Best Of", description="The number of generations to be returned"
    )
    metric: Optional[str] = Field(
        title="Metric",
        description="The metric used for sorting the generations",
    )
    sampling: Optional[str] = Field(
        title="Sampling",
        description="The sampling method used for generating the generations"
        + " - works for open source models only. OpenAI defaults to nucleus sampling.",
    )
    user_properties: Optional[Dict[str, Any]] = Field(
        title="User Properties",
        description="The user properties to track for the generation",
    )
    source: Optional[str] = Field(
        title="Source", description="The source of the generation"
    )
    stream: Optional[bool] = Field(
        title="Streaming",
        description="Whether the generation is to be streamed",
    )


class GenerationResponse(BaseModel):
    generation_id: str = Field(
        title="Generation ID", description="The unique ID of the generation"
    )
    version: str = Field(
        title="Prompt ID", description="The unique ID of the prompt"
    )
    generation: str = Field(
        title="Completion", description="The generated completion"
    )


class ListGenerationResponse(BaseModel):
    task: str = Field(
        title="Task",
        description="The task for which the generation is being requested",
    )
    data: List[GenerationResponse] = Field(
        title="Generations", description="The list of generations"
    )
    created_at: str = Field(
        title="Created At",
        description="The timestamp at which the generations were created",
    )


class GenerationLoggingQuery(BaseModel):
    task: str = Field(
        title="Task",
        description="The task for which the generation is being requested",
    )
    model: str = Field(
        title="Model",
        description="The model that was used to generate the text",
    )
    hyperparameters: Dict[str, Any] = Field(
        title="Hyperparameters",
        description="The hyperparameters that were used to generate the text",
    )
    generation: str = Field(
        title="Generation",
        description="The text that was generated by your model provider",
    )
    prompt_template: Optional[str] = Field(
        title="Prompt Template",
        description="The prompt template that was used to generate the text",
    )
    inputs: Optional[Dict[str, Any]] = Field(
        title="Inputs", description="The inputs filled into the prompt"
    )
    prompt: Optional[str] = Field(
        title="Prompt", description="Prompt used to generate the text"
    )
    version: Optional[str] = Field(
        title="Version", description="The version of the prompt"
    )
    usage: Optional[Dict[str, Any]] = Field(
        title="Usage", description="The token usage for the generation"
    )
    source: Optional[str] = Field(
        title="Source", description="The source of the generation"
    )
    version: Optional[str] = Field(
        title="Version", description="The version of the prompt"
    )
    latency: Optional[float] = Field(
        title="Latency",
        description="The latency of the generation in milliseconds",
    )
    user_properties: Optional[Dict[str, Any]] = Field(
        title="User Properties",
        description="The user properties to track for the generation",
    )
    metadata: Optional[Dict[str, Any]] = Field(
        title="Metadata",
        description="The metadata to track for the generation",
    )
    ground_truth: Optional[str] = Field(
        title="Ground Truth", description="The ground truth for the generation"
    )
    feedback: Optional[Dict[str, Any]] = Field(
        title="Feedback",
        description="The feedback associated with this generation",
    )


class Generation(BaseModel):
    generation_id: str = Field(
        title="Generation ID", description="The unique ID of the generation"
    )
    version: str = Field(
        title="Prompt ID", description="The unique ID of the prompt"
    )
    task: str = Field(
        title="Task",
        description="The task for which the generation is being requested",
    )
    model: str = Field(
        title="Model",
        description="The model that was used to generate the text",
    )
    hyperparameters: Dict[str, Any] = Field(
        title="Hyperparameters",
        description="The hyperparameters that were used to generate the text",
    )
    generation: str = Field(
        title="Generation",
        description="The text that was generated by your model provider",
    )
    total_tokens: int = Field(
        title="Total Tokens",
        description="The total number of tokens generated",
    )
    completion_tokens: int = Field(
        title="Completion Tokens",
        description="The number of tokens generated for the completion",
    )
    cost: float = Field(title="Cost", description="The cost of the generation")
    latency: float = Field(
        title="Latency",
        description="The latency of the generation in milliseconds",
    )
    source: str = Field(
        title="Source", description="The source of the generation"
    )
    feedback: Optional[Dict[str, Any]] = Field(
        title="Feedback",
        description="The feedback associated with this generation",
    )
