ai-contract-platform/backend/app/api/knowledge_bases.py

from typing import Optional

from fastapi import APIRouter, Depends, File, Form, Request, UploadFile
from sqlalchemy.orm import Session

from app.api.reviews import require_current_user
from app.clients.shared_ai_platform import SharedAIPlatformClient, get_shared_ai_platform_client
from app.core.response import success_response
from app.db import get_db_session
from app.services.audit_log_service import AuditLogService

router = APIRouter(prefix="/knowledge-bases", tags=["knowledge-bases"])
ROUTE_PERMISSION_SCOPE = "knowledge_base_management"
ROUTE_PERMISSION_ACTIONS = {
    "list": "knowledge_bases:list",
    "create": "knowledge_bases:create",
    "documents": "knowledge_bases:documents:list",
    "upload_document": "knowledge_bases:documents:upload",
    "retry_document": "knowledge_bases:documents:retry",
}


def get_audit_log_service(db: Session = Depends(get_db_session)) -> AuditLogService:
    return AuditLogService(db)


@router.get("")
def list_knowledge_bases(
    page: int = 1,
    page_size: int = 20,
    shared_ai_platform_client: SharedAIPlatformClient = Depends(get_shared_ai_platform_client),
    current_user: dict = Depends(require_current_user),
):
    return success_response(shared_ai_platform_client.list_knowledge_bases(current_user, page, page_size))


@router.post("")
def create_knowledge_base(
    name: str = Form(...),
    description: Optional[str] = Form(default=None),
    shared_ai_platform_client: SharedAIPlatformClient = Depends(get_shared_ai_platform_client),
    audit_log_service: AuditLogService = Depends(get_audit_log_service),
    current_user: dict = Depends(require_current_user),
):
    result = shared_ai_platform_client.create_knowledge_base(current_user, name, description)
    audit_log_service.record_event(
        tenant_id=current_user.get("tenant_id"),
        user_id=current_user.get("id"),
        action="create_knowledge_base",
        resource_type="knowledge_base",
        resource_id=result["id"],
        detail_json={
            "name": result["name"],
            "description": result["description"],
            "status": result["status"],
            "document_count": result["document_count"],
        },
    )
    return success_response(result)


@router.get("/{knowledge_base_id}/documents")
def list_knowledge_base_documents(
    knowledge_base_id: str,
    page: int = 1,
    page_size: int = 20,
    shared_ai_platform_client: SharedAIPlatformClient = Depends(get_shared_ai_platform_client),
    current_user: dict = Depends(require_current_user),
):
    return success_response(shared_ai_platform_client.list_documents(current_user, knowledge_base_id, page, page_size))


@router.post("/{knowledge_base_id}/documents")
async def upload_knowledge_base_document(
    knowledge_base_id: str,
    request: Request,
    file: UploadFile = File(...),
    note: Optional[str] = Form(default=None),
    shared_ai_platform_client: SharedAIPlatformClient = Depends(get_shared_ai_platform_client),
    audit_log_service: AuditLogService = Depends(get_audit_log_service),
    current_user: dict = Depends(require_current_user),
):
    content = await file.read()
    result = shared_ai_platform_client.upload_document(current_user, knowledge_base_id, file.filename, file.content_type, content)
    audit_log_service.record_event(
        tenant_id=current_user.get("tenant_id"),
        user_id=current_user.get("id"),
        action="upload_knowledge_document",
        resource_type="knowledge_base_document",
        resource_id=result["id"],
        detail_json={
            "knowledge_base_id": knowledge_base_id,
            "file_name": result["file_name"],
            "mime_type": result["mime_type"],
            "file_size": result["file_size"],
            "status": result.get("status"),
            "current_stage": result.get("current_stage"),
            "parse_status": result.get("parse_status"),
            "index_status": result.get("index_status"),
            "error_message": result.get("error_message"),
            "note": note,
        },
        ip_address=request.client.host if request.client else None,
        user_agent=request.headers.get("user-agent"),
    )
    return success_response(result, message="文档已入库" if result["status"] == "indexed" else "文档入库失败")


@router.post("/{knowledge_base_id}/documents/{document_id}/retry")
async def retry_knowledge_base_document(
    knowledge_base_id: str,
    document_id: str,
    request: Request,
    file: Optional[UploadFile] = File(default=None),
    note: Optional[str] = Form(default=None),
    shared_ai_platform_client: SharedAIPlatformClient = Depends(get_shared_ai_platform_client),
    audit_log_service: AuditLogService = Depends(get_audit_log_service),
    current_user: dict = Depends(require_current_user),
):
    content = await file.read() if file is not None else None
    result = shared_ai_platform_client.retry_document(
        current_user=current_user,
        knowledge_base_id=knowledge_base_id,
        document_id=document_id,
        file_name=file.filename if file is not None else None,
        mime_type=file.content_type if file is not None else None,
        content=content,
    )
    audit_log_service.record_event(
        tenant_id=current_user.get("tenant_id"),
        user_id=current_user.get("id"),
        action="retry_knowledge_document",
        resource_type="knowledge_base_document",
        resource_id=result["id"],
        detail_json={
            "knowledge_base_id": knowledge_base_id,
            "original_document_id": document_id,
            "file_name": result["file_name"],
            "mime_type": result["mime_type"],
            "file_size": result["file_size"],
            "status": result.get("status"),
            "current_stage": result.get("current_stage"),
            "parse_status": result.get("parse_status"),
            "index_status": result.get("index_status"),
            "error_message": result.get("error_message"),
            "replacement_file": file is not None,
            "note": note,
        },
        ip_address=request.client.host if request.client else None,
        user_agent=request.headers.get("user-agent"),
    )
    return success_response(result, message="文档重试已完成" if result["status"] == "indexed" else "文档重试仍需处理")