qihuaiflow/backend/python/app/kb_router.py

"""
知识库路由层
============
将检索请求智能路由到本地BM25、向量检索、或外部KB服务(Dify)。
根据知识库配置自动选择最优检索策略。

使用方式:
    from app.kb_router import kb_router_search
    results = await kb_router_search(query, kb_id, top_k=5)
"""

import logging
from typing import List, Dict, Optional

logger = logging.getLogger("aiflow.kb_router")


async def _get_external_services(db_session, enabled_only: bool = True) -> List[Dict]:
    """获取已配置的外部KB服务（仅 dify）"""
    from sqlalchemy import text
    try:
        query = "SELECT * FROM external_kb_services"
        if enabled_only:
            query += " WHERE is_enabled = 1"
        query += " ORDER BY created_at DESC"
        result = db_session.execute(text(query))
        columns = result.keys()
        return [dict(zip(columns, row)) for row in result.fetchall()]
    except Exception as e:
        logger.error(f"Failed to get external services: {e}")
        return []


async def _search_dify(service: Dict, query: str, top_k: int = 5) -> List[Dict]:
    """检索 Dify 知识库（走 app.dify_kb 客户端）"""
    try:
        from app.dify_kb import dify_retrieve
        api_url = service.get("api_url", "")
        api_key = service.get("api_key", "")
        dataset_name = None
        cfg = service.get("config")
        if isinstance(cfg, str):
            try:
                import json
                cfg = json.loads(cfg)
            except Exception:
                cfg = None
        if isinstance(cfg, dict):
            dataset_name = cfg.get("dataset_name")
        hits = dify_retrieve(api_url, api_key, query, top_k=top_k, dataset_name=dataset_name)
        return hits
    except Exception as e:
        logger.error(f"Dify search error: {e}")
        return []


async def _search_external(service: Dict, query: str, top_k: int = 5) -> List[Dict]:
    """路由到对应的外部KB服务（当前仅支持 dify）"""
    service_type = service.get("service_type", "")
    if service_type == "dify":
        return await _search_dify(service, query, top_k)
    else:
        logger.warning(f"Unknown external KB type: {service_type}")
        return []


async def kb_router_search(
    query: str,
    kb_id: Optional[str] = None,
    top_k: int = 5,
    bm25_results: Optional[List[Dict]] = None,
    db_session=None,
    use_vector: bool = True,
    use_external: bool = True,
) -> List[Dict]:
    """
    知识库路由层 — 混合检索

    检索策略:
    1. 本地 BM25 检索 (MySQL, 已有)
    2. 向量语义检索 (PostgreSQL+pgvector, 新增)
    3. 外部KB检索 (Dify, 按配置)

    融合方式: Reciprocal Rank Fusion (RRF)

    Args:
        query: 查询文本
        kb_id: 知识库ID
        top_k: 返回结果数
        bm25_results: BM25检索结果
        db_session: 数据库会话(用于获取外部KB配置)
        use_vector: 是否使用向量检索
        use_external: 是否使用外部KB检索

    Returns:
        融合后的检索结果列表
    """
    all_results = []

    # 1. BM25 结果
    if bm25_results:
        for rank, item in enumerate(bm25_results):
            all_results.append({
                "content": item.get("content", item.get("text", "")),
                "score": item.get("score", 1.0 / (60 + rank + 1)),
                "source": "local_bm25",
                "rank": rank + 1,
                "metadata": item.get("metadata", {}),
            })

    # 2. 向量检索
    if use_vector:
        try:
            from app.pgvector_search import vector_search
            vec_results = await vector_search(query, kb_id, top_k)
            for rank, item in enumerate(vec_results):
                all_results.append({
                    "content": item.get("content", ""),
                    "score": item.get("score", 0.5),
                    "source": "vector_pg",
                    "rank": rank + 1,
                    "metadata": item.get("metadata", {}),
                })
        except ImportError:
            logger.debug("pgvector_search not available, skipping vector search")
        except Exception as e:
            logger.warning(f"Vector search failed: {e}")

    # 3. 外部KB检索 (Dify)
    if use_external and db_session:
        services = await _get_external_services(db_session)
        for service in services:
            if service.get("status") == "connected":
                ext_results = await _search_external(service, query, top_k)
                for rank, item in enumerate(ext_results):
                    all_results.append({
                        "content": item.get("content", ""),
                        "score": item.get("score", 0.5),
                        "source": item.get("source", "external_dify"),
                        "rank": rank + 1,
                        "metadata": item.get("metadata", {}),
                    })

    if not all_results:
        return []

    # RRF 融合
    rrf_k = 60
    score_map = {}
    for item in all_results:
        content_key = item["content"][:200]  # 用内容前200字符作为去重键
        rrf_score = 1.0 / (rrf_k + item.get("rank", 1))

        if content_key in score_map:
            score_map[content_key]["score"] += rrf_score
            score_map[content_key]["sources"].append(item["source"])
        else:
            score_map[content_key] = {
                "content": item["content"],
                "score": rrf_score,
                "sources": [item["source"]],
                "metadata": item.get("metadata", {}),
            }

    # 按融合分数排序
    merged = sorted(score_map.values(), key=lambda x: x["score"], reverse=True)

    # 标准化分数到 0-1
    if merged:
        max_score = merged[0]["score"]
        for item in merged:
            item["score"] = item["score"] / max_score if max_score > 0 else 0
            item["sources"] = list(set(item["sources"]))

    return merged[:top_k]