67 lines
2.9 KiB
Python
67 lines
2.9 KiB
Python
"""Define API models for document."""
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing import List, Optional, Dict, Any
|
|
from datetime import datetime
|
|
# Group related schema definitions so validation rules stay consistent.
|
|
|
|
|
|
|
|
class DocumentUploadRequest(BaseModel):
|
|
"""Define the Document Upload Request API model."""
|
|
doc_name: Optional[str] = Field(None, description="文档名称")
|
|
regulation_type: Optional[str] = Field(None, description="法规类型")
|
|
version: Optional[str] = Field(None, description="文档版本")
|
|
|
|
|
|
class DocumentUploadResponse(BaseModel):
|
|
"""Define the Document Upload Response API model."""
|
|
doc_id: str = Field(..., description="文档ID")
|
|
doc_name: str = Field(..., description="文档名称")
|
|
status: str = Field(..., description="处理状态")
|
|
message: str = Field(default="", description="状态消息")
|
|
num_chunks: int = Field(default=0, description="分块数量")
|
|
summary: str = Field(default="", description="LLM生成的文档摘要")
|
|
summary_latency_ms: int = Field(default=0, description="摘要生成耗时(ms)")
|
|
regulation_type: str = Field(default="", description="法规类型")
|
|
version: str = Field(default="", description="文档版本号")
|
|
timestamp: datetime = Field(default_factory=datetime.now, description="时间戳")
|
|
|
|
|
|
class SearchRequest(BaseModel):
|
|
"""Define the Search Request API model."""
|
|
query: str = Field(..., description="查询文本")
|
|
top_k: int = Field(default=10, description="返回结果数量")
|
|
filters: Optional[str] = Field(None, description="过滤条件")
|
|
|
|
|
|
class SearchResultItem(BaseModel):
|
|
"""Define the Search Result Item API model."""
|
|
id: int = Field(..., description="记录ID")
|
|
content: str = Field(..., description="内容")
|
|
score: float = Field(..., description="相似度分数")
|
|
metadata: Dict[str, Any] = Field(default_factory=dict, description="元数据")
|
|
|
|
|
|
class SearchResponse(BaseModel):
|
|
"""Define the Search Response API model."""
|
|
query: str = Field(..., description="查询文本")
|
|
total: int = Field(..., description="结果总数")
|
|
results: List[SearchResultItem] = Field(default_factory=list, description="结果列表")
|
|
timestamp: datetime = Field(default_factory=datetime.now, description="时间戳")
|
|
|
|
|
|
class DocumentStatusResponse(BaseModel):
|
|
"""Define the Document Status Response API model."""
|
|
doc_id: str = Field(..., description="文档ID")
|
|
status: str = Field(..., description="状态")
|
|
num_chunks: Optional[int] = Field(None, description="分块数量")
|
|
timestamp: datetime = Field(default_factory=datetime.now, description="时间戳")
|
|
|
|
|
|
class ErrorResponse(BaseModel):
|
|
"""Define the Error Response API model."""
|
|
error: str = Field(..., description="错误类型")
|
|
message: str = Field(..., description="错误消息")
|
|
timestamp: datetime = Field(default_factory=datetime.now, description="时间戳")
|