diff --git a/backend/src/kbqa/api/router.py b/backend/src/kbqa/api/router.py index ca144ef..493fb7b 100644 --- a/backend/src/kbqa/api/router.py +++ b/backend/src/kbqa/api/router.py @@ -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) diff --git a/backend/src/kbqa/documents/routes.py b/backend/src/kbqa/documents/routes.py new file mode 100644 index 0000000..03ef3d4 --- /dev/null +++ b/backend/src/kbqa/documents/routes.py @@ -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) diff --git a/backend/src/kbqa/documents/schemas.py b/backend/src/kbqa/documents/schemas.py new file mode 100644 index 0000000..f343b6c --- /dev/null +++ b/backend/src/kbqa/documents/schemas.py @@ -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) diff --git a/backend/src/kbqa/documents/service.py b/backend/src/kbqa/documents/service.py index 90d5f8e..5271e24 100644 --- a/backend/src/kbqa/documents/service.py +++ b/backend/src/kbqa/documents/service.py @@ -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")