101 lines
2.2 KiB
Python
101 lines
2.2 KiB
Python
from datetime import datetime
|
|
from typing import Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
# ---------- Request Schemas ----------
|
|
|
|
class KnowledgeBaseCreate(BaseModel):
|
|
name: str = Field(..., max_length=200)
|
|
type: str = Field(..., pattern="^(industry|enterprise)$")
|
|
description: Optional[str] = None
|
|
|
|
|
|
class DocumentUpload(BaseModel):
|
|
title: str = Field(..., max_length=500)
|
|
source_type: str = Field(..., pattern="^(text|url|markdown)$")
|
|
content: Optional[str] = None # for text/markdown
|
|
source_url: Optional[str] = None # for url type
|
|
|
|
|
|
class KnowledgeSearchRequest(BaseModel):
|
|
query: str = Field(..., min_length=1)
|
|
knowledge_base_ids: list[str]
|
|
top_k: int = Field(default=10, ge=1, le=50)
|
|
|
|
|
|
# ---------- Response Schemas ----------
|
|
|
|
class KnowledgeBaseResponse(BaseModel):
|
|
id: str
|
|
name: str
|
|
type: str
|
|
description: Optional[str]
|
|
document_count: int
|
|
status: str
|
|
created_at: datetime
|
|
|
|
model_config = {"from_attributes": True}
|
|
|
|
|
|
class DocumentResponse(BaseModel):
|
|
id: str
|
|
title: str
|
|
source_type: str
|
|
source_url: Optional[str]
|
|
chunk_count: int
|
|
status: str
|
|
error_message: Optional[str]
|
|
created_at: datetime
|
|
|
|
model_config = {"from_attributes": True}
|
|
|
|
|
|
class SearchResultItem(BaseModel):
|
|
chunk_id: str
|
|
content: str
|
|
score: float
|
|
document_id: str
|
|
document_title: str
|
|
metadata: dict
|
|
|
|
|
|
class SearchResponse(BaseModel):
|
|
results: list[SearchResultItem]
|
|
total: int
|
|
latency_ms: int
|
|
|
|
|
|
class ChunkPreview(BaseModel):
|
|
id: str
|
|
content: str
|
|
chunk_index: int
|
|
token_count: int
|
|
|
|
model_config = {"from_attributes": True}
|
|
|
|
|
|
# ---------- 增量索引 Schemas ----------
|
|
|
|
class UpdateDocumentRequest(BaseModel):
|
|
"""更新文档内容请求"""
|
|
content: str
|
|
|
|
|
|
class RetrieveRequest(BaseModel):
|
|
"""增强检索请求"""
|
|
query: str = Field(..., min_length=1)
|
|
top_k: Optional[int] = Field(default=5, ge=1, le=50)
|
|
use_rerank: Optional[bool] = Field(default=True)
|
|
use_compression: Optional[bool] = Field(default=False)
|
|
|
|
|
|
class RebuildResponse(BaseModel):
|
|
"""重建索引响应"""
|
|
total: int
|
|
processed: int
|
|
skipped: int
|
|
failed: int
|
|
errors: list[dict]
|