amb_rag/backend/app/api/documents.py

"""文档 API 路由。"""

from fastapi import APIRouter, Depends, Query, UploadFile, File, Form
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from app.api.deps import get_current_user, get_db
from app.models.user import User
from app.schemas.document import (
    DocumentListResponse,
    DocumentResponse,
    DocumentUpdateRequest,
    DocumentUploadResponse,
)
from app.services.doc_service import DocumentService

router = APIRouter(prefix="/documents", tags=["documents"])


class CreateTextRequest(BaseModel):
    kb_id: str
    title: str = Field(min_length=1, max_length=512)
    content: str = Field(min_length=1)
    content_format: str = "markdown"  # markdown / text
    category_id: str | None = None


@router.post("/upload", response_model=DocumentUploadResponse, status_code=201)
async def upload_document(
    kb_id: str = Form(...),
    file: UploadFile = File(...),
    category_id: str = Form(None),
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
) -> DocumentUploadResponse:
    """上传文档到知识库。"""
    content = await file.read()
    svc = DocumentService(db)
    doc = svc.upload(
        user=user,
        kb_id=kb_id,
        filename=file.filename or "unknown",
        content=content,
        category_id=category_id,
    )
    return DocumentUploadResponse(
        id=doc.id,
        original_filename=doc.original_filename,
        file_size=doc.file_size,
        status=doc.status,
        message="文档上传成功，等待解析。",
    )


@router.post("/create-text", response_model=DocumentResponse, status_code=201)
def create_text_document(
    body: CreateTextRequest,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
) -> DocumentResponse:
    """直接创建文本内容文档（无需上传文件）。"""
    svc = DocumentService(db)
    doc = svc.create_text(
        user=user,
        kb_id=body.kb_id,
        title=body.title,
        content=body.content,
        content_format=body.content_format,
        category_id=body.category_id,
    )
    return _to_response(doc)


@router.get("", response_model=DocumentListResponse)
def list_documents(
    kb_id: str = Query(...),
    category_id: str = Query(None),
    page: int = Query(1, ge=1),
    page_size: int = Query(50, ge=1, le=200),
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
) -> DocumentListResponse:
    svc = DocumentService(db)
    items, total = svc.list_by_knowledge_base(kb_id, user, category_id=category_id, page=page, page_size=page_size)
    return DocumentListResponse(
        items=[_to_response(doc) for doc in items],
        total=total,
    )


@router.get("/{doc_id}", response_model=DocumentResponse)
def get_document(
    doc_id: str,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
) -> DocumentResponse:
    svc = DocumentService(db)
    doc = svc.get_or_404(doc_id, user)
    return _to_response(doc)


@router.put("/{doc_id}", response_model=DocumentResponse)
def update_document(
    doc_id: str,
    body: DocumentUpdateRequest,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
) -> DocumentResponse:
    svc = DocumentService(db)
    doc = svc.get_or_404(doc_id, user)
    update_fields = {}
    if body.title is not None:
        update_fields["title"] = body.title
    if body.description is not None:
        update_fields["description"] = body.description
    if body.keywords is not None:
        update_fields["keywords"] = body.keywords
    if body.category_id is not None:
        update_fields["category_id"] = body.category_id
    if update_fields:
        from app.repositories.doc_repo import DocumentRepository

        DocumentRepository(db).update(doc, **update_fields)
        db.commit()
    return _to_response(doc)


@router.delete("/{doc_id}", status_code=204)
def delete_document(
    doc_id: str,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
) -> None:
    svc = DocumentService(db)
    svc.delete(doc_id, user)


@router.post("/{doc_id}/reprocess", response_model=DocumentResponse)
def reprocess_document(
    doc_id: str,
    user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
) -> DocumentResponse:
    """重新解析文档。"""
    svc = DocumentService(db)
    doc = svc.get_or_404(doc_id, user)
    svc._process_document(doc)
    db.refresh(doc)
    return _to_response(doc)


def _to_response(doc) -> DocumentResponse:
    return DocumentResponse(
        id=doc.id,
        knowledge_base_id=doc.knowledge_base_id,
        original_filename=doc.original_filename,
        file_size=doc.file_size,
        mime_type=doc.mime_type,
        file_ext=doc.file_ext,
        sha256=doc.sha256,
        title=doc.title,
        description=doc.description,
        keywords=doc.keywords,
        content_summary=doc.content_summary,
        status=doc.status,
        error_code=doc.error_code,
        doc_token_hint=doc.doc_token_hint,
        category_id=doc.category_id,
        created_at=doc.created_at,
        updated_at=doc.updated_at,
    )