263 lines
9.6 KiB
Python
263 lines
9.6 KiB
Python
"""
|
||
知识库路由层
|
||
============
|
||
将检索请求智能路由到本地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]
|