easy-rag/models/__init__.py

60 lines
1.2 KiB
Python

from pydantic import BaseModel
from typing import Optional, List
from datetime import datetime
class DocumentUpload(BaseModel):
"""文档上传请求模型"""
filename: str
content_type: str
class DocumentInfo(BaseModel):
"""文档信息模型"""
id: str
filename: str
upload_time: datetime
size: int
chunks_count: int
class ChatRequest(BaseModel):
"""聊天请求模型"""
question: str
top_k: Optional[int] = 3
temperature: Optional[float] = 0.7
class ChatResponse(BaseModel):
"""聊天响应模型"""
answer: str
sources: List[dict]
processing_time: float
class ErrorResponse(BaseModel):
"""错误响应模型"""
error: str
detail: Optional[str] = None
class SuccessResponse(BaseModel):
"""成功响应模型"""
message: str
data: Optional[dict] = None
class StreamChatRequest(BaseModel):
"""流式聊天请求模型"""
question: str
top_k: Optional[int] = 3
temperature: Optional[float] = 0.7
class StreamChatChunk(BaseModel):
"""流式聊天数据块模型"""
content: str
is_final: bool = False
sources: Optional[List[dict]] = None
processing_time: Optional[float] = None