from pydantic import BaseModel, Field
from typing import List, Optional
from huggingface_hub import ChatCompletionInputMessage, ChatCompletionInputGrammarType, ChatCompletionInputStreamOptions, ChatCompletionInputToolChoiceClass, ChatCompletionInputTool

class ChatRequest(BaseModel):
    model: str = Field(..., description="The model to use for chat-completion. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed Inference Endpoint. If not provided, the default recommended model for chat-based text-generation will be used. See https://huggingface.co/tasks/text-generation for more details.")
    messages: List[ChatCompletionInputMessage] = Field(..., description="Conversation history consisting of roles and content pairs.")
    frequency_penalty: Optional[float] = Field(0.0, ge=-2.0, le=2.0, description="Penalizes new tokens based on their existing frequency in the text so far. Range: [-2.0, 2.0]. Defaults to 0.0.")
    logit_bias: Optional[dict] = Field(None, description="Modify the likelihood of specified tokens appearing in the completion. Accepts a JSON object that maps tokens to an associated bias value from -100 to 100.")
    logprobs: Optional[bool] = Field(None, description="Whether to return log probabilities of the output tokens or not.")
    max_tokens: Optional[int] = Field(8192, description="Maximum number of tokens allowed in the response. Defaults to 100.")
    n: Optional[int] = Field(None, description="UNUSED.")
    presence_penalty: Optional[float] = Field(None, ge=-2.0, le=2.0, description="Positive values penalize new tokens based on whether they appear in the text so far.")
    response_format: Optional[ChatCompletionInputGrammarType] = Field(None, description="Grammar constraints. Can be either a JSONSchema or a regex.")
    seed: Optional[int] = Field(None, description="Seed for reproducible control flow.")
    stop: Optional[str] = Field(None, description="Up to four strings which trigger the end of the response.")
    stream: Optional[bool] = Field(False, description="Enable realtime streaming of responses. Defaults to False.")
    stream_options: Optional[ChatCompletionInputStreamOptions] = Field(None, description="Options for streaming completions.")
    temperature: Optional[float] = Field(1.0, ge=0.0, le=2.0, description="Controls randomness of the generations. Lower values ensure less random completions.")
    top_logprobs: Optional[int] = Field(None, ge=0, le=5, description="Specifying the number of most likely tokens to return at each token position.")
    top_p: Optional[float] = Field(0.95, gt=0.0, lt=1.0, description="Fraction of the most likely next words to sample from.")
    tool_choice: Optional[ChatCompletionInputToolChoiceClass] = Field("auto", description="The tool to use for the completion. Defaults to 'auto'.")
    tool_prompt: Optional[str] = Field(None, description="A prompt to be appended before the tools.")
    tools: Optional[List] = Field(None, description="A list of tools the model may call.")