chat-ai-api/app/utils/kb_search.py

"""
知识库检索工具
用于在聊天中检索自建知识库内容
"""
from typing import List, Dict, Optional
from ..core.qdrant_client import get_qdrant_manager
from ..utils.embedding_service import get_embedding_service
from ..db.knowledge_base import search_chunks_by_ids


def search_knowledge_base(query: str, tenant_id: str, top_k: int = 5) -> List[Dict]:
    """
    在知识库中搜索相关内容
    :param query: 用户查询文本
    :param tenant_id: 租户ID（appName）
    :param top_k: 返回Top K个结果
    :return: 检索结果列表 [{"content": "...", "score": 0.95, "doc_id": "..."}, ...]
    """
    try:
        # 1. 将查询文本向量化
        embedding_service = get_embedding_service()
        query_vector = embedding_service.generate_single_embedding(query, tenant_id)

        if not query_vector:
            return []

        # 2. 在 Qdrant 中进行向量相似度搜索
        qdrant = get_qdrant_manager()
        results = qdrant.search_vectors(tenant_id, query_vector, top_k)

        # 3. 从 MongoDB 获取完整的切片内容
        if results:
            vector_ids = [r["vector_id"] for r in results]
            chunks = search_chunks_by_ids(vector_ids)

            # 合并 Qdrant 的相似度分数和 MongoDB 的完整内容
            vector_id_to_chunk = {chunk["vector_id"]: chunk for chunk in chunks}

            enriched_results = []
            for result in results:
                chunk = vector_id_to_chunk.get(result["vector_id"])
                if chunk:
                    enriched_results.append({
                        "content": chunk["content"],
                        "score": result["score"],
                        "doc_id": chunk["doc_id"],
                        "chunk_index": chunk["chunk_index"],
                    })

            return enriched_results

        return []

    except Exception as e:
        print(f"知识库检索失败: {str(e)}")
        return []


def format_knowledge_context(search_results: List[Dict]) -> str:
    """
    将检索结果格式化为上下文文本，用于注入到 prompt
    :param search_results: 检索结果列表
    :return: 格式化的上下文文本
    """
    if not search_results:
        return ""

    context_parts = ["<知识库相关信息>"]

    for idx, result in enumerate(search_results, 1):
        context_parts.append(f"\n[文档片段 {idx}]")
        context_parts.append(result["content"])

    context_parts.append("\n</知识库相关信息>")

    return "\n".join(context_parts)


def get_custom_knowledge_tools(tenant_id: str) -> List[Dict]:
    """
    构造自定义知识库检索工具配置
    用于替代豆包原生的知识库工具
    :param tenant_id: 租户ID
    :return: 工具配置列表
    """
    # 注意：这里返回空列表，因为我们将知识库检索前置到prompt中
    # 而不是作为工具让模型调用
    return []