"""知识库 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