100d1a06eb
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
153 lines
4.5 KiB
Python
153 lines
4.5 KiB
Python
"""
|
||
Agentic RAG - 检索 Mixin
|
||
|
||
包含知识库检索、网络搜索等方法
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
import requests
|
||
|
||
from .agentic_base import (
|
||
logger, HAS_SERPER, SERPER_API_KEY,
|
||
SOURCE_KB, SOURCE_WEB
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class SearchMixin:
|
||
"""检索功能方法"""
|
||
|
||
def _web_search(self, query: str, top_k: int = 5) -> list:
|
||
"""网络搜索（使用Serper API）"""
|
||
if not HAS_SERPER:
|
||
return []
|
||
|
||
try:
|
||
url = "https://google.serper.dev/search"
|
||
payload = json.dumps({
|
||
"q": query,
|
||
"gl": "cn",
|
||
"hl": "zh-cn",
|
||
"num": top_k
|
||
})
|
||
headers = {
|
||
'X-API-KEY': SERPER_API_KEY,
|
||
'Content-Type': 'application/json'
|
||
}
|
||
|
||
response = requests.post(url, headers=headers, data=payload, timeout=10)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
results = []
|
||
for item in data.get('organic', [])[:top_k]:
|
||
results.append({
|
||
'title': item.get('title', ''),
|
||
'link': item.get('link', ''),
|
||
'snippet': item.get('snippet', ''),
|
||
'date': item.get('date', '')
|
||
})
|
||
|
||
return results
|
||
|
||
except Exception as e:
|
||
logger.warning(f"网络搜索失败: {e}")
|
||
return []
|
||
|
||
def _should_web_search(self, query: str) -> bool:
|
||
"""判断是否需要网络搜索"""
|
||
realtime_keywords = [
|
||
"今天", "最新", "今日", "当前", "现在",
|
||
"天气", "新闻", "股价", "行情", "汇率",
|
||
"最近", "近期", "这周", "本月", "今年",
|
||
"实时", "动态", "热点", "发生"
|
||
]
|
||
|
||
query_lower = query.lower()
|
||
return any(kw in query_lower for kw in realtime_keywords)
|
||
|
||
def _web_search_flow(self, query: str, log_trace: list, emit_log, verbose: bool,
|
||
allowed_levels: list = None) -> list:
|
||
"""
|
||
网络搜索流程
|
||
|
||
Args:
|
||
query: 查询
|
||
log_trace: 日志追踪列表
|
||
emit_log: 日志发射函数
|
||
verbose: 是否详细输出
|
||
allowed_levels: 允许的安全级别
|
||
|
||
Returns:
|
||
网络搜索结果列表
|
||
"""
|
||
if not self.enable_web_search or not HAS_SERPER:
|
||
return []
|
||
|
||
if emit_log:
|
||
emit_log("🌐 触发网络搜索...")
|
||
|
||
web_results = self._web_search(query, top_k=5)
|
||
|
||
if not web_results:
|
||
if emit_log:
|
||
emit_log("⚠️ 网络搜索未返回结果")
|
||
return []
|
||
|
||
# 转换为统一上下文格式
|
||
web_contexts = []
|
||
for item in web_results:
|
||
web_contexts.append({
|
||
'doc': f"{item.get('title', '')}\n{item.get('snippet', '')}",
|
||
'meta': {
|
||
'source': self.SOURCE_WEB,
|
||
'link': item.get('link', ''),
|
||
'date': item.get('date', '')
|
||
},
|
||
'source_type': self.SOURCE_WEB,
|
||
'query': query
|
||
})
|
||
|
||
log_trace.append({
|
||
'phase': 'web_search',
|
||
'query': query,
|
||
'results_count': len(web_contexts)
|
||
})
|
||
|
||
if emit_log:
|
||
emit_log(f"✅ 网络搜索返回 {len(web_contexts)} 条结果")
|
||
|
||
return web_contexts
|
||
|
||
def _is_kb_result_sufficient(self, query: str, docs: list) -> bool:
|
||
"""判断知识库检索结果是否充分"""
|
||
if not docs:
|
||
return False
|
||
|
||
# 结果数量检查
|
||
if len(docs) >= 3:
|
||
# 至少3条结果，检查相关性
|
||
high_rel_count = 0
|
||
for doc in docs:
|
||
score = doc.get('score', 0) or doc.get('distance', 1)
|
||
# cosine 距离转相似度
|
||
if isinstance(score, (int, float)):
|
||
sim = 1 - score if score <= 1 else score
|
||
if sim >= 0.6:
|
||
high_rel_count += 1
|
||
|
||
if high_rel_count >= 2:
|
||
return True
|
||
|
||
# 有高质量结果
|
||
for doc in docs[:2]:
|
||
score = doc.get('score', 0) or doc.get('distance', 1)
|
||
if isinstance(score, (int, float)):
|
||
sim = 1 - score if score <= 1 else score
|
||
if sim >= 0.8:
|
||
return True
|
||
|
||
return False
|