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 _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)