""" 知识库路由层 ============ 将检索请求智能路由到本地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 _get_web_services(db_session, enabled_only: bool = True) -> List[Dict]: """获取已启用的 AnySearch 联网检索服务配置。""" from sqlalchemy import text try: query = "SELECT * FROM external_kb_services WHERE service_type='anysearch'" if enabled_only: query += " AND 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 AnySearch services: {e}") return [] async def _search_anysearch(service: Dict, query: str, top_k: int = 5) -> List[Dict]: """调用 AnySearch 联网检索（走 app.anysearch_client）。""" try: from app.anysearch_client import anysearch_search api_key = service.get("api_key", "") api_url = service.get("api_url") or "https://api.anysearch.com/v1/search" cfg = service.get("config") if isinstance(cfg, str): try: import json cfg = json.loads(cfg) except Exception: cfg = None cfg = cfg or {} zone = cfg.get("zone", "cn") freshness = cfg.get("freshness") k = int(cfg.get("top_k", top_k)) return await anysearch_search(api_key, api_url, query, top_k=k, zone=zone, freshness=freshness) except Exception as e: logger.error(f"AnySearch search error: {e}") return [] async def kb_router_search( query: str, kb_id: Optional[str] = None, kb_ids: Optional[List[str]] = None, top_k: int = 5, bm25_results: Optional[List[Dict]] = None, db_session=None, use_vector: bool = True, use_external: bool = True, use_tree: bool = True, use_web: bool = False, ) -> List[Dict]: """ 知识库路由层 — 混合检索 检索策略: 1. 本地 BM25 检索 (MySQL, 已有) 2. 向量语义检索 (PostgreSQL+pgvector, 新增) 3. 结构化分片段检索 (TreeSearch, 新增: 知识库目录树 + chunk 级精准命中) 4. 外部KB检索 (Dify, 按配置) 5. 联网检索 (AnySearch, 可选 use_web) 融合方式: Reciprocal Rank Fusion (RRF) Args: query: 查询文本 kb_id: 知识库ID top_k: 返回结果数 bm25_results: BM25检索结果 db_session: 数据库会话(用于获取外部KB配置与分片段检索) use_vector: 是否使用向量检索 use_external: 是否使用外部KB检索 use_tree: 是否使用 TreeSearch 结构化分片段检索 use_web: 是否使用 AnySearch 联网检索 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 kb_targets = kb_ids if kb_ids else ([kb_id] if kb_id else [None]) vec_results = [] for _kid in kb_targets: try: vec_results.extend(await vector_search(query, _kid, top_k)) except Exception as e: logger.warning(f"Vector search failed for kb={_kid}: {e}") 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}") # 2.5 结构化分片段检索 (TreeSearch, Level 2) if use_tree and db_session is not None: try: from app.tree_search import tree_search kb_targets = kb_ids if kb_ids else ([kb_id] if kb_id else None) tree_hits = await tree_search( db_session, query, kb_ids=kb_targets, top_k=top_k, mode="hybrid", ) for rank, item in enumerate(tree_hits): all_results.append({ "content": item.get("content", ""), "score": item.get("score", 0.5), "source": "tree_chunk", "rank": rank + 1, "metadata": item.get("metadata", {}), }) except Exception as e: logger.warning(f"TreeSearch 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", {}), }) # 4. 联网检索 (AnySearch, 可选) if use_web and db_session is not None: try: web_services = await _get_web_services(db_session) for svc in web_services: web_hits = await _search_anysearch(svc, query, top_k) for rank, item in enumerate(web_hits): all_results.append({ "content": item.get("content", ""), "score": item.get("score", 0.6), "source": item.get("source", "web_anysearch"), "rank": rank + 1, "metadata": item.get("metadata", {}), }) except Exception as e: logger.warning(f"AnySearch web search failed: {e}") 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]