""" 知识库路由层 ============ 将检索请求智能路由到本地BM25、向量检索、或外部KB服务(FastGPT/Dify/RAGFlow)。 根据知识库配置自动选择最优检索策略。 使用方式: from app.kb_router import kb_router_search results = await kb_router_search(query, kb_id, top_k=5) """ import os import json import logging import httpx 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服务""" 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_fastgpt(api_url: str, api_key: str, query: str, top_k: int = 5) -> List[Dict]: """搜索 FastGPT 知识库""" try: async with httpx.AsyncClient(timeout=30) as client: resp = await client.post( f"{api_url.rstrip('/')}/api/v1/dataset/search", headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", }, json={ "query": query, "similarity": 0.5, "limit": top_k, }, ) if resp.status_code == 200: data = resp.json() results = [] for item in data.get("data", {}).get("list", []): results.append({ "content": item.get("q", item.get("content", "")), "score": float(item.get("score", 0.5)), "source": f"FastGPT: {item.get('collectionId', '')}", "metadata": item, }) return results else: logger.warning(f"FastGPT search returned {resp.status_code}") return [] except Exception as e: logger.error(f"FastGPT search error: {e}") return [] async def _search_dify(api_url: str, api_key: str, query: str, top_k: int = 5) -> List[Dict]: """搜索 Dify 知识库""" try: async with httpx.AsyncClient(timeout=30) as client: resp = await client.post( f"{api_url.rstrip('/')}/v1/datasets/retrieval", headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", }, json={ "query": query, "retrieval_mode": "semantic", "top_k": top_k, }, ) if resp.status_code == 200: data = resp.json() results = [] for item in data.get("records", []): segment = item.get("segment", {}) results.append({ "content": segment.get("content", ""), "score": float(item.get("score", 0.5)), "source": f"Dify: {segment.get('id', '')}", "metadata": item, }) return results else: logger.warning(f"Dify search returned {resp.status_code}") return [] except Exception as e: logger.error(f"Dify search error: {e}") return [] async def _search_ragflow(api_url: str, api_key: str, query: str, top_k: int = 5) -> List[Dict]: """搜索 RAGFlow 知识库""" try: async with httpx.AsyncClient(timeout=30) as client: resp = await client.post( f"{api_url.rstrip('/')}/api/v1/retrieval", headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", }, json={ "question": query, "top_k": top_k, }, ) if resp.status_code == 200: data = resp.json() results = [] for item in data.get("data", {}).get("chunks", []): results.append({ "content": item.get("content", ""), "score": float(item.get("similarity", 0.5)), "source": f"RAGFlow: {item.get('document_id', '')}", "metadata": item, }) return results else: logger.warning(f"RAGFlow search returned {resp.status_code}") return [] except Exception as e: logger.error(f"RAGFlow search error: {e}") return [] async def _search_external(service: Dict, query: str, top_k: int = 5) -> List[Dict]: """路由到对应的外部KB服务""" service_type = service.get("service_type", "") api_url = service.get("api_url", "") api_key = service.get("api_key", "") if service_type == "fastgpt": return await _search_fastgpt(api_url, api_key, query, top_k) elif service_type == "dify": return await _search_dify(api_url, api_key, query, top_k) elif service_type == "ragflow": return await _search_ragflow(api_url, api_key, 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检索 (FastGPT/Dify/RAGFlow, 按配置) 融合方式: 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检索 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", f"external_{service['service_type']}"), "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]