T4: AnySearch 改为全局固定配置 + 修复外部服务弹窗 UI
This commit is contained in:
QihuAI Design committed
1 parent
3c9d9518b1
commit
97125d7b65
6 files changed
+634
-441
No files matched your search
+271
-262
@@ -1,262 +1,271 @@
|
||||
"""
|
||||
知识库路由层
|
||||
============
|
||||
将检索请求智能路由到本地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]
|
||||
"""
|
||||
知识库路由层
|
||||
============
|
||||
将检索请求智能路由到本地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 联网检索配置（从 site_settings 读取）。"""
|
||||
from app.anysearch_client import load_anysearch_config
|
||||
cfg = load_anysearch_config(db_session)
|
||||
if not cfg:
|
||||
return []
|
||||
if enabled_only and not cfg.get("is_enabled"):
|
||||
return []
|
||||
api_key = cfg.get("api_key", "")
|
||||
if not api_key:
|
||||
return []
|
||||
return [{
|
||||
"id": "anysearch_global",
|
||||
"service_type": "anysearch",
|
||||
"api_key": api_key,
|
||||
"api_url": cfg.get("api_url") or "https://api.anysearch.com/v1/search",
|
||||
"config": {
|
||||
"zone": cfg.get("zone", "cn"),
|
||||
"freshness": cfg.get("freshness"),
|
||||
"top_k": cfg.get("top_k", 5),
|
||||
},
|
||||
"is_enabled": cfg.get("is_enabled", False),
|
||||
}]
|
||||
|
||||
|
||||
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]
|
||||
Reference in new issue
Block a user