chat-ai-api/app/routers/knowledge_base.py

"""
知识库管理 API
提供文档上传、查询、删除、智能客服配置等接口
"""
from fastapi import APIRouter, HTTPException, Depends, UploadFile, File, BackgroundTasks
from typing import List, Optional
from pydantic import BaseModel
import os

from ..db.knowledge_base import (
    get_tenant_config, upsert_tenant_config, get_agent_config, update_agent_config,
    create_document, get_document_by_id, list_documents, count_documents,
    update_document_status, delete_document, create_chunks, get_chunks_by_doc,
    delete_chunks_by_doc, get_tenant_stats
)
from ..core.oss_client import get_oss_client
from ..core.qdrant_client import get_qdrant_manager
from ..utils.document_processor import get_document_processor
from ..utils.embedding_service import get_embedding_service
from ..dependencies.auth import resolve_user_id, UserContext

router = APIRouter()

# 允许的文件类型
ALLOWED_FILE_TYPES = ['.xlsx', '.xls', '.docx', '.doc', '.pdf']
MAX_FILE_SIZE = 20 * 1024 * 1024  # 20MB


# ==================== Pydantic 模型 ====================

class AgentConfigUpdate(BaseModel):
    enabled: bool
    name: str
    greeting: str


class DocumentInfo(BaseModel):
    id: str
    filename: str
    file_type: str
    file_size: int
    status: str
    chunk_count: int
    created_at: str


class ChunkInfo(BaseModel):
    chunk_index: int
    content: str
    vector_id: Optional[str]


# ==================== 后台任务：文档处理 ====================

async def process_document_task(doc_id: str, tenant_id: str, file_content: bytes, file_type: str):
    """
    后台任务：处理文档（提取文本、切片、向量化、存储）
    """
    try:
        # 更新状态为处理中
        update_document_status(doc_id, "processing")

        # 1. 提取文本
        processor = get_document_processor()
        text = processor.extract_text(file_content, file_type)

        if not text:
            update_document_status(doc_id, "failed", error_msg="文本提取失败")
            return

        # 清理文本
        text = processor.clean_text(text)

        # 2. 智能切片
        chunks = processor.split_text_into_chunks(text)

        if not chunks:
            update_document_status(doc_id, "failed", error_msg="文档切片失败")
            return

        # 3. 向量化（批量）
        embedding_service = get_embedding_service()
        chunk_contents = [chunk["content"] for chunk in chunks]
        vectors = embedding_service.generate_embeddings(chunk_contents, tenant_id)

        if len(vectors) != len(chunks):
            update_document_status(doc_id, "failed", error_msg="向量化失败")
            return

        # 4. 存储到 Qdrant
        qdrant = get_qdrant_manager()
        chunks_with_vectors = [
            {"content": chunk["content"], "vector": vector}
            for chunk, vector in zip(chunks, vectors)
        ]
        vector_ids = qdrant.add_vectors(tenant_id, doc_id, chunks_with_vectors)

        # 5. 存储切片元数据到 MongoDB
        chunks_data = [
            {
                "content": chunk["content"],
                "vector_id": vector_id,
                "metadata": chunk.get("metadata", {})
            }
            for chunk, vector_id in zip(chunks, vector_ids)
        ]
        create_chunks(doc_id, tenant_id, chunks_data)

        # 6. 更新文档状态为完成
        update_document_status(doc_id, "completed", chunk_count=len(chunks))

        print(f"文档处理完成: {doc_id}, 切片数: {len(chunks)}")

    except Exception as e:
        print(f"文档处理异常: {str(e)}")
        update_document_status(doc_id, "failed", error_msg=str(e))


# ==================== 文档管理接口 ====================

@router.post("/documents/upload")
async def upload_document(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    user_context: UserContext = Depends(resolve_user_id),
):
    """
    上传文档到知识库
    1. 上传原始文件到阿里云OSS
    2. 后台任务：提取文本、切片、向量化、存储
    """
    tenant_id = user_context.app_name

    # 验证文件类型
    file_ext = os.path.splitext(file.filename)[1].lower()
    if file_ext not in ALLOWED_FILE_TYPES:
        raise HTTPException(status_code=400, detail=f"不支持的文件类型: {file_ext}")

    # 读取文件内容
    file_content = await file.read()
    file_size = len(file_content)

    # 验证文件大小
    if file_size > MAX_FILE_SIZE:
        raise HTTPException(status_code=400, detail=f"文件大小超出限制 ({MAX_FILE_SIZE // 1024 // 1024}MB)")

    # 上传到 OSS
    oss_client = get_oss_client()
    success, oss_path = oss_client.upload_file(tenant_id, file_content, file.filename)

    if not success:
        raise HTTPException(status_code=500, detail=f"文件上传失败: {oss_path}")

    # 创建文档记录
    doc_id = create_document(
        tenant_id=tenant_id,
        filename=file.filename,
        file_type=file_ext,
        file_size=file_size,
        oss_path=oss_path
    )

    # 后台处理文档
    background_tasks.add_task(process_document_task, doc_id, tenant_id, file_content, file_ext)

    return {
        "success": True,
        "doc_id": doc_id,
        "message": "文档上传成功，正在后台处理..."
    }


@router.get("/documents")
async def get_documents(
    skip: int = 0,
    limit: int = 100,
    user_context: UserContext = Depends(resolve_user_id),
):
    """获取租户的所有文档列表"""
    tenant_id = user_context.app_name
    docs = list_documents(tenant_id, skip, limit)
    total = count_documents(tenant_id)

    return {
        "success": True,
        "total": total,
        "documents": docs
    }


@router.get("/documents/{doc_id}")
async def get_document_detail(
    doc_id: str,
    user_context: UserContext = Depends(resolve_user_id),
):
    """获取文档详情（包含切片）"""
    tenant_id = user_context.app_name

    # 获取文档信息
    doc = get_document_by_id(doc_id)
    if not doc:
        raise HTTPException(status_code=404, detail="文档不存在")

    # 验证租户权限
    if doc["tenant_id"] != tenant_id:
        raise HTTPException(status_code=403, detail="无权访问该文档")

    # 获取切片
    chunks = get_chunks_by_doc(doc_id)

    # 获取原始文件URL（临时访问链接）
    oss_client = get_oss_client()
    file_url = oss_client.get_file_url(doc["oss_path"])

    return {
        "success": True,
        "document": doc,
        "chunks": chunks,
        "file_url": file_url
    }


@router.delete("/documents/{doc_id}")
async def delete_document_endpoint(
    doc_id: str,
    user_context: UserContext = Depends(resolve_user_id),
):
    """删除文档（包括OSS文件、向量、切片）"""
    tenant_id = user_context.app_name

    # 获取文档信息
    doc = get_document_by_id(doc_id)
    if not doc:
        raise HTTPException(status_code=404, detail="文档不存在")

    # 验证租户权限
    if doc["tenant_id"] != tenant_id:
        raise HTTPException(status_code=403, detail="无权删除该文档")

    # 1. 删除 OSS 文件
    oss_client = get_oss_client()
    oss_client.delete_file(doc["oss_path"])

    # 2. 删除 Qdrant 向量
    qdrant = get_qdrant_manager()
    qdrant.delete_vectors_by_doc(doc_id)

    # 3. 删除 MongoDB 切片
    delete_chunks_by_doc(doc_id)

    # 4. 软删除文档记录
    delete_document(doc_id)

    return {
        "success": True,
        "message": "文档删除成功"
    }


# ==================== 智能客服配置接口 ====================

@router.get("/agent/config")
async def get_agent_config_endpoint(
    user_context: UserContext = Depends(resolve_user_id),
):
    """获取智能客服配置"""
    tenant_id = user_context.app_name
    config = get_agent_config(tenant_id)

    if not config:
        # 返回默认配置
        return {
            "success": True,
            "config": {
                "enabled": False,
                "name": "智能小助手",
                "greeting": "您好！我是智能客服小助手，有什么可以帮您的吗？",
                "h5_url": f"https://aichat.zhongsou.com/?tenant={tenant_id}"
            }
        }

    return {
        "success": True,
        "config": config
    }


@router.post("/agent/config")
async def update_agent_config_endpoint(
    config: AgentConfigUpdate,
    user_context: UserContext = Depends(resolve_user_id),
):
    """更新智能客服配置"""
    tenant_id = user_context.app_name

    success = update_agent_config(
        tenant_id=tenant_id,
        enabled=config.enabled,
        name=config.name,
        greeting=config.greeting
    )

    if not success:
        raise HTTPException(status_code=500, detail="配置更新失败")

    return {
        "success": True,
        "message": "配置更新成功"
    }


# ==================== 统计接口 ====================

@router.get("/stats")
async def get_stats(
    user_context: UserContext = Depends(resolve_user_id),
):
    """获取租户的知识库统计数据"""
    tenant_id = user_context.app_name
    stats = get_tenant_stats(tenant_id)

    return {
        "success": True,
        "stats": stats
    }