feat: expose document management api
This commit is contained in:
@@ -1,3 +1,6 @@
|
||||
from fastapi import APIRouter
|
||||
|
||||
from kbqa.documents.routes import router as documents_router
|
||||
|
||||
api_router = APIRouter(prefix="/api/v1")
|
||||
api_router.include_router(documents_router)
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from fastapi import APIRouter, Depends, File, Query, Request, Response, UploadFile, status
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from kbqa.api.errors import AppError
|
||||
from kbqa.database import get_session
|
||||
from kbqa.documents.repository import DocumentRepository
|
||||
from kbqa.documents.schemas import ChunkResponse, DocumentResponse, Page
|
||||
from kbqa.documents.service import DocumentService
|
||||
|
||||
router = APIRouter(prefix="/documents", tags=["documents"])
|
||||
|
||||
SessionDependency = Annotated[AsyncSession, Depends(get_session)]
|
||||
|
||||
|
||||
def get_document_service(request: Request) -> DocumentService:
|
||||
return request.app.state.document_service
|
||||
|
||||
|
||||
@router.post("", response_model=DocumentResponse, status_code=status.HTTP_202_ACCEPTED)
|
||||
async def upload_document(
|
||||
file: Annotated[UploadFile, File()],
|
||||
service: Annotated[DocumentService, Depends(get_document_service)],
|
||||
) -> DocumentResponse:
|
||||
document = await service.upload(file)
|
||||
return DocumentResponse.model_validate(document)
|
||||
|
||||
|
||||
@router.get("", response_model=Page[DocumentResponse])
|
||||
async def list_documents(
|
||||
session: SessionDependency,
|
||||
document_status: Annotated[
|
||||
Literal["pending", "indexing", "indexed", "failed", "deleting"] | None,
|
||||
Query(alias="status"),
|
||||
] = None,
|
||||
offset: Annotated[int, Query(ge=0)] = 0,
|
||||
limit: Annotated[int, Query(ge=1, le=100)] = 20,
|
||||
) -> Page[DocumentResponse]:
|
||||
items, total = await DocumentRepository(session).page(document_status, offset, limit)
|
||||
return Page(
|
||||
items=[DocumentResponse.model_validate(item) for item in items],
|
||||
total=total,
|
||||
offset=offset,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{document_id}", response_model=DocumentResponse)
|
||||
async def get_document(document_id: str, session: SessionDependency) -> DocumentResponse:
|
||||
document = await DocumentRepository(session).get(document_id)
|
||||
if document is None:
|
||||
raise AppError("DOCUMENT_NOT_FOUND", "文档不存在", 404)
|
||||
return DocumentResponse.model_validate(document)
|
||||
|
||||
|
||||
@router.get("/{document_id}/chunks", response_model=Page[ChunkResponse])
|
||||
async def list_document_chunks(
|
||||
document_id: str,
|
||||
session: SessionDependency,
|
||||
offset: Annotated[int, Query(ge=0)] = 0,
|
||||
limit: Annotated[int, Query(ge=1, le=100)] = 20,
|
||||
) -> Page[ChunkResponse]:
|
||||
repository = DocumentRepository(session)
|
||||
if await repository.get(document_id) is None:
|
||||
raise AppError("DOCUMENT_NOT_FOUND", "文档不存在", 404)
|
||||
items, total = await repository.page_chunks(document_id, offset, limit)
|
||||
return Page(
|
||||
items=[ChunkResponse.model_validate(item) for item in items],
|
||||
total=total,
|
||||
offset=offset,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{document_id}/retry",
|
||||
response_model=DocumentResponse,
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
)
|
||||
async def retry_document(
|
||||
document_id: str,
|
||||
service: Annotated[DocumentService, Depends(get_document_service)],
|
||||
) -> DocumentResponse:
|
||||
document = await service.retry(document_id)
|
||||
return DocumentResponse.model_validate(document)
|
||||
|
||||
|
||||
@router.delete("/{document_id}", status_code=status.HTTP_204_NO_CONTENT)
|
||||
async def delete_document(
|
||||
document_id: str,
|
||||
service: Annotated[DocumentService, Depends(get_document_service)],
|
||||
) -> Response:
|
||||
await service.delete(document_id)
|
||||
return Response(status_code=status.HTTP_204_NO_CONTENT)
|
||||
@@ -0,0 +1,46 @@
|
||||
from datetime import datetime
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class Page(BaseModel, Generic[T]):
|
||||
items: list[T]
|
||||
total: int
|
||||
offset: int
|
||||
limit: int
|
||||
|
||||
|
||||
class DocumentResponse(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
id: str
|
||||
filename: str
|
||||
file_type: str
|
||||
mime_type: str
|
||||
file_size: int
|
||||
status: str
|
||||
chunk_count: int
|
||||
error_code: str | None
|
||||
error_message: str | None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class ChunkResponse(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
id: str
|
||||
document_id: str
|
||||
order_index: int
|
||||
content: str
|
||||
page_number: int | None
|
||||
char_count: int
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class Pagination(BaseModel):
|
||||
offset: int = Field(default=0, ge=0)
|
||||
limit: int = Field(default=20, ge=1, le=100)
|
||||
@@ -21,6 +21,12 @@ ALLOWED_EXTENSIONS = {
|
||||
".txt": ("txt", "text/plain"),
|
||||
}
|
||||
|
||||
ALLOWED_MIME_TYPES = {
|
||||
"pdf": {"application/pdf"},
|
||||
"md": {"text/markdown", "text/plain", "text/x-markdown"},
|
||||
"txt": {"text/plain"},
|
||||
}
|
||||
|
||||
|
||||
class DocumentService:
|
||||
def __init__(
|
||||
@@ -41,6 +47,8 @@ class DocumentService:
|
||||
if suffix not in ALLOWED_EXTENSIONS:
|
||||
raise AppError("UNSUPPORTED_FILE_TYPE", "仅支持 PDF、Markdown 和 TXT 文件", 415)
|
||||
file_type, expected_mime = ALLOWED_EXTENSIONS[suffix]
|
||||
if file.content_type not in ALLOWED_MIME_TYPES[file_type]:
|
||||
raise AppError("INVALID_MIME_TYPE", "文件扩展名与内容类型不匹配", 415)
|
||||
document_id = str(uuid4())
|
||||
target = self.settings.raw_data_dir / f"{document_id}{suffix}"
|
||||
temporary = target.with_suffix(f"{suffix}.part")
|
||||
|
||||
Reference in New Issue
Block a user