Enhance knowledge base functionality and integrate S3 storage support

- Add new models for `KnowledgeDocument` and `KnowledgeChunk` to manage document ingestion and chunking. - Implement S3-compatible storage integration for knowledge documents, allowing for file uploads and retrieval. - Introduce API endpoints for managing knowledge bases and documents, including creation, deletion, and searching. - Update frontend components to support knowledge base configuration and document management, improving user interaction. - Enhance backend services for knowledge processing and retrieval, ensuring robust handling of document statuses and errors.
2026-07-12 13:58:47 +08:00
parent 01c563a3e7
commit de58f30014
21 changed files with 995 additions and 34 deletions
--- a/backend/routes/knowledge_bases.py
+++ b/backend/routes/knowledge_bases.py
@@ -6,11 +6,16 @@ KB 自身引用一个 Embedding 模型资源。被助手引用时禁止删除

 import uuid

-from db.models import KnowledgeBase, ModelResource
+from db.models import Assistant, KnowledgeBase, KnowledgeChunk, KnowledgeDocument, ModelResource
 from db.session import get_session
-from fastapi import APIRouter, Depends, HTTPException
-from schemas import KnowledgeBaseOut, KnowledgeBaseUpsert
+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
@@ -43,6 +48,16 @@ def _to_out(kb: KnowledgeBase) -> KnowledgeBaseOut:
    )


+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 = (
@@ -83,6 +98,12 @@ async def update_knowledge_base(
    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()
@@ -97,6 +118,19 @@ async def delete_knowledge_base(
    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()
@@ -105,3 +139,119 @@ async def delete_knowledge_base(
        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