#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

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


class SetConversationRequest(BaseModel):
    """设置对话请求"""
    conversation_id: Optional[str] = None
    is_new: bool
    name: Optional[str] = Field(default="New conversation", max_length=255)
    dialog_id: str


class DeleteConversationsRequest(BaseModel):
    """删除对话请求"""
    conversation_ids: List[str]


class CompletionRequest(BaseModel):
    """完成请求（聊天完成）"""
    conversation_id: str
    messages: List[Dict[str, Any]]
    llm_id: Optional[str] = None
    stream: Optional[bool] = True
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    frequency_penalty: Optional[float] = None
    presence_penalty: Optional[float] = None
    max_tokens: Optional[int] = None


class TTSRequest(BaseModel):
    """文本转语音请求"""
    text: str


class DeleteMessageRequest(BaseModel):
    """删除消息请求"""
    conversation_id: str
    message_id: str


class ThumbupRequest(BaseModel):
    """点赞/点踩请求"""
    conversation_id: str
    message_id: str
    thumbup: Optional[bool] = None
    feedback: Optional[str] = ""


class AskRequest(BaseModel):
    """提问请求"""
    question: str
    kb_ids: List[str]
    search_id: Optional[str] = ""


class MindmapRequest(BaseModel):
    """思维导图请求"""
    question: str
    kb_ids: List[str]
    search_id: Optional[str] = ""


class RelatedQuestionsRequest(BaseModel):
    """相关问题请求"""
    question: str
    search_id: Optional[str] = ""