chat-ai-api/app/core/viking_client.py

"""
Viking 知识库上传工具模块
职责：
  1. 将本地临时文件上传到火山引擎 TOS 对象存储
  2. 调用 Viking /api/knowledge/doc/v2/add 接口将 TOS 文件导入知识库
"""
import os
import uuid
import asyncio
from concurrent.futures import ThreadPoolExecutor

import httpx

from ..config.config import Config

cfg = Config()

# 线程池用于执行 TOS 同步 SDK 调用
_executor = ThreadPoolExecutor(max_workers=4)


def _tos_upload_sync(local_path: str, object_key: str) -> str:
    """
    同步上传本地文件到 TOS，返回 tos:// URI。
    在线程池中调用，避免阻塞事件循环。
    """
    try:
        import tos as tos_sdk
    except ImportError:
        raise RuntimeError("缺少 tos 依赖，请执行 pip install tos")

    client = tos_sdk.TosClientV2(
        ak=cfg.VIKING_AK,
        sk=cfg.VIKING_SK,
        endpoint=cfg.VIKING_TOS_ENDPOINT,
        region=cfg.VIKING_TOS_REGION,
    )
    client.put_object_from_file(
        bucket=cfg.VIKING_TOS_BUCKET,
        key=object_key,
        file_path=local_path,
    )
    return f"tos://{cfg.VIKING_TOS_BUCKET}/{object_key}"


async def upload_to_tos(local_path: str, app_name: str, original_filename: str) -> tuple[str, str]:
    """
    异步上传本地文件到 TOS。
    返回 (object_key, tos_uri)。
    object_key 格式：{app_name}/{uuid32}_{原始文件名}
    """
    # 清理文件名中的特殊字符，只保留字母数字下划线中文和常见标点
    safe_name = os.path.basename(original_filename)
    uid = uuid.uuid4().hex
    object_key = f"{app_name}/{uid}_{safe_name}"

    loop = asyncio.get_event_loop()
    tos_uri = await loop.run_in_executor(_executor, _tos_upload_sync, local_path, object_key)
    return object_key, tos_uri


async def add_doc_to_knowledge(
    resource_id: str,
    doc_id: str,
    tos_uri: str,
    doc_name: str,
    doc_type: str,
    tag_list: list | None = None,
    path_segments: list | None = None,
) -> dict:
    """
    调用 Viking /api/knowledge/doc/v2/add 接口，将 TOS 文件导入知识库。
    使用 resource_id 指定知识库（对应 ai_config.knowledgeId）。
    返回响应体 data 字段。
    """
    url = f"https://{cfg.VIKING_HOST}/api/knowledge/doc/v2/add"
    payload: dict = {
        "resource_id": resource_id,
        "doc_id": doc_id,
        "doc_name": doc_name,
        "doc_type": doc_type,
        "uri": tos_uri,
    }
    if tag_list:
        payload["tag_list"] = tag_list
    if path_segments:
        payload["path_segments"] = path_segments

    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.post(
            url,
            json=payload,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {cfg.VIKING_API_KEY}",
            },
        )
        body = resp.json()
        code = body.get("code", -1)
        if code != 0:
            raise ValueError(
                f"Viking API 返回错误: {body.get('message')} "
                f"(code={code}, request_id={body.get('request_id')})"
            )
        return body["data"]


def _tos_delete_sync(object_key: str) -> None:
    """
    同步删除 TOS 对象。
    在线程池中调用，避免阻塞事件循环。
    """
    try:
        import tos as tos_sdk
    except ImportError:
        raise RuntimeError("缺少 tos 依赖，请执行 pip install tos")

    client = tos_sdk.TosClientV2(
        ak=cfg.VIKING_AK,
        sk=cfg.VIKING_SK,
        endpoint=cfg.VIKING_TOS_ENDPOINT,
        region=cfg.VIKING_TOS_REGION,
    )
    client.delete_object(bucket=cfg.VIKING_TOS_BUCKET, key=object_key)


async def delete_from_tos(object_key: str) -> None:
    """异步删除 TOS 对象。"""
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(_executor, _tos_delete_sync, object_key)


def _search_knowledge_sync(resource_id: str, query: str, limit: int = 5) -> list:
    """
    同步检索 Viking 知识库，返回 PointInfo 列表。
    在线程池中调用，避免阻塞事件循环。
    配置：开启重排序（base-multilingual-rerank）、邻近切片扩散1、文本聚合。
    """
    from vikingdb.knowledge import VikingKnowledge
    from vikingdb.auth import APIKey
    from vikingdb.knowledge.models.search import SearchKnowledgeRequest

    viking_client = VikingKnowledge(
        host=cfg.VIKING_HOST,
        region="cn-beijing",
        auth=APIKey(api_key=cfg.VIKING_API_KEY),
        scheme="https",
    )
    collection = viking_client.collection(resource_id=resource_id)
    resp = collection.search_knowledge(SearchKnowledgeRequest(
        query=query,
        limit=limit,
        post_processing={
            "rerank_switch": True,
            "retrieve_count": max(limit * 3, 15),   # 进入重排的候选数，须 >= limit
            "chunk_diffusion_count": 1,              # 命中切片上下各扩散1片
            "chunk_group": True,                     # 按文档原始顺序重排，适合叙述性文档
        },
    ))
    if resp and resp.result and resp.result.result_list:
        return resp.result.result_list
    return []


async def search_knowledge(resource_id: str, query: str, limit: int = 5) -> list:
    """
    异步检索 Viking 知识库，返回 PointInfo 列表。
    result_list 非空表示知识库有命中，为空则降级到联网搜索。
    """
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(_executor, _search_knowledge_sync, resource_id, query, limit)


async def delete_doc_from_knowledge(resource_id: str, doc_id: str) -> dict:
    """
    调用 Viking /api/knowledge/doc/delete 接口删除知识库文档。
    返回完整响应体。
    """
    url = f"https://{cfg.VIKING_HOST}/api/knowledge/doc/delete"
    payload = {
        "resource_id": resource_id,
        "doc_id": doc_id,
    }

    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.post(
            url,
            json=payload,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {cfg.VIKING_API_KEY}",
            },
        )
        body = resp.json()
        code = body.get("code", -1)
        if code != 0:
            raise ValueError(
                f"Viking API 返回错误: {body.get('message')} "
                f"(code={code}, request_id={body.get('request_id')})"
            )
        return body