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

"""
知识库工具路由
提供本地文档上传到 Viking 知识库的接口，无需 JWT 鉴权。
接口：POST /main/knowledgeTool/upload
"""
import os
import json
import uuid
import asyncio
import tempfile
from typing import Optional

from fastapi import APIRouter, HTTPException, UploadFile, File, Form
from pydantic import BaseModel

from ..core.viking_client import upload_to_tos, add_doc_to_knowledge, delete_from_tos, delete_doc_from_knowledge
from ..db.ai_config import get_config_by_app_name
from ..db.mongo import save_knowledge_doc, get_knowledge_docs, get_knowledge_doc_by_doc_id, delete_knowledge_doc

router = APIRouter()


class KnowledgeDocsQuery(BaseModel):
    app_name: str
    page: int = 1
    page_size: int = 20


class KnowledgeDeleteRequest(BaseModel):
    doc_id: str
    app_name: str

# 支持的文档类型：后缀 → Viking doc_type 值
ALLOWED_DOC_EXTENSIONS: dict[str, str] = {
    ".txt":      "txt",
    ".doc":      "doc",
    ".docx":     "docx",
    ".pdf":      "pdf",
    ".md":       "markdown",
    ".markdown": "markdown",
    ".pptx":     "pptx",
    ".ppt":      "ppt",
}

# 单文件大小上限：50 MB
MAX_DOC_SIZE = 50 * 1024 * 1024


def _generate_doc_id() -> str:
    """生成符合 Viking 命名规范的唯一 doc_id（字母开头，仅含字母/数字/下划线）"""
    return f"doc_{uuid.uuid4().hex}"


async def _cleanup_tos(object_key: str) -> None:
    """上传成功后异步清理 TOS 临时文件，失败仅打印日志不影响主流程。"""
    try:
        await delete_from_tos(object_key)
    except Exception as e:
        print(f"TOS 临时文件清理失败（key={object_key}）: {e}")


@router.post("/upload", summary="上传本地文档到知识库")
async def upload_knowledge_doc(
    file: UploadFile = File(..., description="待上传文档，支持 txt/doc/docx/pdf/markdown/pptx/ppt"),
    app_name: str = Form(..., description="应用包名，用于查询对应知识库配置"),
    doc_id: Optional[str] = Form(None, description="文档唯一标识，不填则自动生成"),
    doc_name: Optional[str] = Form(None, description="文档显示名称，不填则使用原始文件名"),
    path_segments: Optional[str] = Form(None, description="目录路径 JSON 数组，如 [\"产品文档\",\"V2.0\"]"),
    tag_list: Optional[str] = Form(None, description="标签列表 JSON 数组，格式与 Viking API 一致"),
):
    """
    上传本地文档到火山引擎 Viking 知识库。

    流程：
    1. 校验文件类型和大小
    2. 从 ai_config 查询 knowledgeId（作为 resource_id）
    3. 上传文件到 TOS 对象存储
    4. 调用 Viking /api/knowledge/doc/v2/add 导入知识库
    5. 返回 doc_id 和知识库信息
    """
    # ── 1. 校验文件类型 ──────────────────────────────────────────
    original_filename = file.filename or ""
    _, ext = os.path.splitext(original_filename.lower())
    if ext not in ALLOWED_DOC_EXTENSIONS:
        allowed = "、".join(ALLOWED_DOC_EXTENSIONS.keys())
        raise HTTPException(
            status_code=400,
            detail=f"不支持的文件类型 '{ext}'，仅支持：{allowed}",
        )
    doc_type = ALLOWED_DOC_EXTENSIONS[ext]

    # ── 2. 读取并校验文件大小 ────────────────────────────────────
    file_bytes = await file.read()
    if len(file_bytes) > MAX_DOC_SIZE:
        raise HTTPException(
            status_code=400,
            detail=f"文件大小超过限制（最大 50 MB，当前 {len(file_bytes) / 1024 / 1024:.1f} MB）",
        )
    if len(file_bytes) == 0:
        raise HTTPException(status_code=400, detail="文件内容为空")

    # ── 3. 查询知识库配置 ────────────────────────────────────────
    ai_config = get_config_by_app_name(app_name)
    if not ai_config:
        raise HTTPException(status_code=404, detail=f"未找到 appName '{app_name}' 的配置")

    resource_id = ai_config.get("knowledgeId", "")
    if not resource_id:
        raise HTTPException(
            status_code=400,
            detail=f"appName '{app_name}' 的配置中未设置 knowledgeId",
        )

    # ── 4. 解析可选 JSON 参数 ────────────────────────────────────
    # 强制携带 app_name 标签，前端额外传入的标签追加在后面
    parsed_tag_list: list = [
        {"field_name": "app_name", "field_type": "string", "field_value": app_name}
    ]
    if tag_list:
        try:
            extra_tags = json.loads(tag_list)
            parsed_tag_list.extend(extra_tags)
        except json.JSONDecodeError:
            raise HTTPException(status_code=400, detail="tag_list 格式错误，需为合法的 JSON 数组")

    parsed_path_segments: list | None = None
    if path_segments:
        try:
            parsed_path_segments = json.loads(path_segments)
        except json.JSONDecodeError:
            raise HTTPException(status_code=400, detail="path_segments 格式错误，需为合法的 JSON 数组")

    # ── 5. 确定 doc_id 和 doc_name ──────────────────────────────
    final_doc_id = doc_id.strip() if doc_id else _generate_doc_id()
    final_doc_name = doc_name.strip() if doc_name else original_filename

    # ── 6. 写临时文件 ────────────────────────────────────────────
    tmp_path = None
    try:
        suffix = ext  # 保留原后缀，TOS SDK 可能依赖文件名做 Content-Type 推断
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(file_bytes)
            tmp_path = tmp.name

        # ── 7. 上传到 TOS ────────────────────────────────────────
        try:
            object_key, tos_uri = await upload_to_tos(tmp_path, app_name, original_filename)
        except RuntimeError as e:
            raise HTTPException(status_code=500, detail=str(e))
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"TOS 上传失败: {str(e)}")

        # ── 8. 调用 Viking API 导入知识库 ───────────────────────
        try:
            kb_data = await add_doc_to_knowledge(
                resource_id=resource_id,
                doc_id=final_doc_id,
                tos_uri=tos_uri,
                doc_name=final_doc_name,
                doc_type=doc_type,
                tag_list=parsed_tag_list,
                path_segments=parsed_path_segments,
            )
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"导入知识库失败: {str(e)}")

        # ── 9. 异步删除 TOS 临时文件（不阻塞响应）─────────────────
        asyncio.create_task(_cleanup_tos(object_key))

    finally:
        # ── 10. 清理本地临时文件 ─────────────────────────────────
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)

    # ── 11. 构造返回结果并存储到 MongoDB ─────────────────────────
    result = {
        "doc_id": kb_data.get("doc_id", final_doc_id),
        "resource_id": kb_data.get("resource_id", resource_id),
        "collection_name": kb_data.get("collection_name", ""),
        "project": kb_data.get("project", "default"),
        "doc_name": final_doc_name,
        "doc_type": doc_type,
        "app_name": app_name,
    }
    inserted_id = save_knowledge_doc(result.copy())
    result["id"] = inserted_id or ""

    return result


@router.post("/docs", summary="分页查询知识库文档上传记录")
async def list_knowledge_docs(request: KnowledgeDocsQuery):
    """
    按 app_name 分页查询已上传的知识库文档记录，按上传时间倒序。
    请求体使用 JSON，方便后续扩展查询条件。
    """
    page = max(1, request.page)
    page_size = request.page_size if 1 <= request.page_size <= 100 else 20
    return get_knowledge_docs(app_name=request.app_name, page=page, page_size=page_size)


@router.post("/delete", summary="删除知识库文档")
async def delete_knowledge_doc_api(request: KnowledgeDeleteRequest):
    """
    删除指定文档，依次执行：
    1. 按 doc_id 查询 MongoDB，获取 resource_id
    2. 调用 Viking /api/knowledge/doc/delete 删除知识库文档
    3. 删除 MongoDB 中的上传记录
    注：TOS 临时文件由桶生命周期规则自动过期清理，此处不处理。
    """
    # ── 1. 查询 MongoDB 记录 ────────────────────────────────────
    doc = get_knowledge_doc_by_doc_id(request.doc_id)
    if not doc:
        raise HTTPException(status_code=404, detail=f"未找到 doc_id '{request.doc_id}' 的文档记录")

    if doc.get("app_name") != request.app_name:
        raise HTTPException(status_code=403, detail="doc_id 与 app_name 不匹配")

    resource_id = doc.get("resource_id", "")

    # ── 2. 删除 Viking 知识库文档 ───────────────────────────────
    try:
        await delete_doc_from_knowledge(resource_id=resource_id, doc_id=request.doc_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"删除 Viking 知识库文档失败: {str(e)}")

    # ── 3. 删除 MongoDB 记录 ────────────────────────────────────
    delete_knowledge_doc(request.doc_id)

    return {"doc_id": request.doc_id, "message": "删除成功"}