from typing import List, Optional, Union
from pydantic import BaseModel, Field


# region 入库
class AddDocDto(BaseModel):
    class DocDataObj(BaseModel):
        UserName: str
        DbName: str
        FileID: Union[int, str]
        FileName: str

    class ChunkData(BaseModel):
        Question: str
        Answer: str

    DocType: str = "qa"
    DocData: DocDataObj
    ChunkDatas: List[ChunkData] = []


class AddDocVo(BaseModel):
    """添加QA响应VO模型"""

    Status: str
    ChunkIds: List[int]

    def success(self) -> bool:
        return self.Status == "success"


# endregion


# region 搜索


class SearchQaDto(BaseModel):
    class DocData(BaseModel):
        UserName: str
        DbName: str

    DocType: Optional[str] = "qa"
    DocDatas: List[DocData] = []
    Query: Optional[str] = None
    Limit: Optional[int] = 100
    SearchLimit: Optional[int] = 100


class SearchQaVo(BaseModel):
    """搜索响应VO模型"""

    class MatchItem(BaseModel):
        """匹配项模型"""

        Answer: Optional[str] = None
        AnswerType: Optional[str] = None
        DbName: Optional[str] = None
        FileId: Optional[str] = None
        FileName: Optional[str] = None
        From: Optional[str] = None
        Id: Optional[int] = None
        Question: Optional[str] = None
        Score: Optional[float] = None
        SimilarFlag: Optional[int] = None
        UserName: Optional[str] = Field(..., description="用户名")

    Count: Optional[int] = Field(..., description="匹配结果数量")
    MatchList: Optional[List[MatchItem]] = Field(..., description="匹配结果列表")


# endregion
