ai-video-fullstack/backend/routes/knowledge_bases.py

"""知识库 CRUD。前端助手编辑页的"知识库"下拉对接这里。

KB 自身引用一个 Embedding 模型资源。被助手引用时禁止删除
(DB 层 ON DELETE RESTRICT),这里把外键冲突翻译成 409。
"""

import uuid

from db.models import Assistant, KnowledgeBase, KnowledgeChunk, KnowledgeDocument, ModelResource
from db.session import get_session
from fastapi import APIRouter, BackgroundTasks, Depends, File, Form, HTTPException, UploadFile
from schemas import (
    KnowledgeBaseOut, KnowledgeBaseUpsert, KnowledgeChunkOut, KnowledgeDocumentOut,
    KnowledgeSearchIn, KnowledgeTextIn,
)
from services.auth import require_admin
from services.knowledge import create_document, delete_storage_object, process_document, search
import settings
from sqlalchemy import select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

router = APIRouter(
    prefix="/api/knowledge-bases",
    tags=["knowledge-bases"],
    dependencies=[Depends(require_admin)],
)


async def _validate_embedding_resource(
    session: AsyncSession, resource_id: str | None
) -> None:
    if not resource_id:
        return
    resource = await session.get(ModelResource, resource_id)
    if not resource or resource.capability != "Embedding":
        raise HTTPException(400, "知识库必须引用 Embedding 模型资源")


def _to_out(kb: KnowledgeBase) -> KnowledgeBaseOut:
    return KnowledgeBaseOut(
        id=kb.id,
        name=kb.name,
        description=kb.description,
        embedding_model_resource_id=kb.embedding_model_resource_id,
        status=kb.status,
        updated_at=kb.updated_at.isoformat() if kb.updated_at else None,
    )


def _document_out(document: KnowledgeDocument) -> KnowledgeDocumentOut:
    return KnowledgeDocumentOut(
        id=document.id, knowledge_base_id=document.knowledge_base_id,
        name=document.name, source_type=document.source_type, mime_type=document.mime_type,
        size_bytes=document.size_bytes, status=document.status,
        error_message=document.error_message, chunk_count=document.chunk_count,
        created_at=document.created_at.isoformat() if document.created_at else None,
    )


@router.get("", response_model=list[KnowledgeBaseOut])
async def list_knowledge_bases(session: AsyncSession = Depends(get_session)):
    rows = (
        await session.execute(select(KnowledgeBase).order_by(KnowledgeBase.name))
    ).scalars().all()
    return [_to_out(kb) for kb in rows]


@router.post("", response_model=KnowledgeBaseOut)
async def create_knowledge_base(
    body: KnowledgeBaseUpsert, session: AsyncSession = Depends(get_session)
):
    await _validate_embedding_resource(session, body.embedding_model_resource_id)
    kb = KnowledgeBase(id=f"kb_{uuid.uuid4().hex[:12]}", **body.model_dump())
    session.add(kb)
    await session.commit()
    await session.refresh(kb)
    return _to_out(kb)


@router.get("/{kb_id}", response_model=KnowledgeBaseOut)
async def get_knowledge_base(
    kb_id: str, session: AsyncSession = Depends(get_session)
):
    kb = await session.get(KnowledgeBase, kb_id)
    if not kb:
        raise HTTPException(404, "知识库不存在")
    return _to_out(kb)


@router.put("/{kb_id}", response_model=KnowledgeBaseOut)
async def update_knowledge_base(
    kb_id: str,
    body: KnowledgeBaseUpsert,
    session: AsyncSession = Depends(get_session),
):
    kb = await session.get(KnowledgeBase, kb_id)
    if not kb:
        raise HTTPException(404, "知识库不存在")
    await _validate_embedding_resource(session, body.embedding_model_resource_id)
    if kb.embedding_model_resource_id != body.embedding_model_resource_id:
        has_document = (await session.execute(
            select(KnowledgeDocument.id).where(KnowledgeDocument.knowledge_base_id == kb_id).limit(1)
        )).scalar_one_or_none()
        if has_document:
            raise HTTPException(409, "知识库已有文档，不能更换 Embedding 模型；请新建知识库")
    for k, v in body.model_dump().items():
        setattr(kb, k, v)
    await session.commit()
    await session.refresh(kb)
    return _to_out(kb)


@router.delete("/{kb_id}")
async def delete_knowledge_base(
    kb_id: str, session: AsyncSession = Depends(get_session)
):
    kb = await session.get(KnowledgeBase, kb_id)
    if not kb:
        raise HTTPException(404, "知识库不存在")
    referenced = (await session.execute(
        select(Assistant.id).where(Assistant.knowledge_base_id == kb_id).limit(1)
    )).scalar_one_or_none()
    if referenced:
        raise HTTPException(409, "知识库正被助手引用,无法删除")
    documents = (await session.execute(
        select(KnowledgeDocument).where(KnowledgeDocument.knowledge_base_id == kb_id)
    )).scalars().all()
    for document in documents:
        try:
            await delete_storage_object(document)
        except Exception:
            pass
    try:
        await session.delete(kb)
        await session.commit()
    except IntegrityError:
        # 被助手引用(ON DELETE RESTRICT):先解绑再删
        await session.rollback()
        raise HTTPException(409, "知识库正被助手引用,无法删除")
    return {"ok": True}


@router.get("/{kb_id}/documents", response_model=list[KnowledgeDocumentOut])
async def list_documents(kb_id: str, session: AsyncSession = Depends(get_session)):
    if not await session.get(KnowledgeBase, kb_id):
        raise HTTPException(404, "知识库不存在")
    rows = (await session.execute(
        select(KnowledgeDocument).where(KnowledgeDocument.knowledge_base_id == kb_id)
        .order_by(KnowledgeDocument.created_at.desc())
    )).scalars().all()
    return [_document_out(row) for row in rows]


@router.post("/{kb_id}/documents/text", response_model=KnowledgeDocumentOut, status_code=202)
async def add_text(
    kb_id: str, body: KnowledgeTextIn, background_tasks: BackgroundTasks,
    session: AsyncSession = Depends(get_session),
):
    kb = await session.get(KnowledgeBase, kb_id)
    if not kb:
        raise HTTPException(404, "知识库不存在")
    try:
        document = await create_document(
            session, kb, name=body.name, source_type="text",
            raw_data=body.content.encode("utf-8"), mime_type="text/plain",
        )
        background_tasks.add_task(process_document, document.id)
        return _document_out(document)
    except Exception as exc:
        raise HTTPException(400, f"文字入库失败: {exc}") from exc


@router.post("/{kb_id}/documents/file", response_model=KnowledgeDocumentOut, status_code=202)
async def add_file(
    kb_id: str,
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    name: str = Form(default=""),
    session: AsyncSession = Depends(get_session),
):
    kb = await session.get(KnowledgeBase, kb_id)
    if not kb:
        raise HTTPException(404, "知识库不存在")
    data = await file.read(settings.KNOWLEDGE_MAX_FILE_BYTES + 1)
    if len(data) > settings.KNOWLEDGE_MAX_FILE_BYTES:
        limit_mb = settings.KNOWLEDGE_MAX_FILE_BYTES // 1024 // 1024
        raise HTTPException(413, f"文件不能超过 {limit_mb} MB")
    filename = file.filename or "未命名文件"
    try:
        document = await create_document(
            session, kb, name=name.strip() or filename, source_type="file",
            raw_data=data, mime_type=file.content_type or "application/octet-stream",
        )
        background_tasks.add_task(process_document, document.id)
        return _document_out(document)
    except Exception as exc:
        raise HTTPException(400, f"文件入库失败: {exc}") from exc


@router.delete("/{kb_id}/documents/{document_id}")
async def delete_document(kb_id: str, document_id: str, session: AsyncSession = Depends(get_session)):
    document = await session.get(KnowledgeDocument, document_id)
    if not document or document.knowledge_base_id != kb_id:
        raise HTTPException(404, "文档不存在")
    try:
        await delete_storage_object(document)
    finally:
        await session.delete(document)
        await session.commit()
    return {"ok": True}


@router.post(
    "/{kb_id}/documents/{document_id}/retry",
    response_model=KnowledgeDocumentOut,
    status_code=202,
)
async def retry_document(
    kb_id: str, document_id: str, background_tasks: BackgroundTasks,
    session: AsyncSession = Depends(get_session),
):
    document = await session.get(KnowledgeDocument, document_id)
    if not document or document.knowledge_base_id != kb_id:
        raise HTTPException(404, "文档不存在")
    if document.status in {"pending", "processing"}:
        raise HTTPException(409, "文档正在处理中")
    document.status = "pending"
    document.error_message = ""
    await session.commit()
    await session.refresh(document)
    background_tasks.add_task(process_document, document.id)
    return _document_out(document)


@router.get("/{kb_id}/documents/{document_id}/chunks", response_model=list[KnowledgeChunkOut])
async def list_document_chunks(
    kb_id: str, document_id: str, session: AsyncSession = Depends(get_session)
):
    document = await session.get(KnowledgeDocument, document_id)
    if not document or document.knowledge_base_id != kb_id:
        raise HTTPException(404, "文档不存在")
    chunks = (await session.execute(
        select(KnowledgeChunk).where(KnowledgeChunk.document_id == document_id)
        .order_by(KnowledgeChunk.chunk_index).limit(100)
    )).scalars().all()
    return [KnowledgeChunkOut(id=row.id, chunk_index=row.chunk_index, content=row.content) for row in chunks]


@router.post("/{kb_id}/search")
async def search_knowledge_base(
    kb_id: str, body: KnowledgeSearchIn, session: AsyncSession = Depends(get_session)
):
    try:
        return await search(session, kb_id, body.query, body.top_k)
    except Exception as exc:
        raise HTTPException(400, f"检索失败: {exc}") from exc