JOJO
e6d8f194d2
- 模型输出 【cite:src_xxx】/【file:相对路径】 标记，渲染为可交互来源胶囊（悬停延迟开/点击固定/滚动收起） - 后端：modules/citations.py 注册表与落库校验（剥离幻觉 id 与不存在的文件）；web_search/extract_webpage 注册来源并在工具结果带 citation 信息；assistant 消息落库时 finalize 并挂载 message.metadata.citations - 前端：marker 输出瞬间即渲染（工具结果查表 + 文件 token 自解析），task_complete 后权威裁决与富化；弹层支持网页摘要/文件内容片段/图片预览，文件头部点击打开右侧预览面板，宿主机模式可在文件管理器中打开 - prompt：三份系统提示加入引用格式与规则（提到文件名不等于引用；图片等二进制文件同样可引用）
541 lines
19 KiB
Python
541 lines
19 KiB
Python
# modules/search_engine.py - 网络搜索模块
|
||
|
||
import httpx
|
||
import json
|
||
from typing import Dict, Optional, Any, List
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
import re
|
||
from urllib.parse import urlparse
|
||
try:
|
||
from config import TAVILY_API_KEY, SEARCH_MAX_RESULTS, OUTPUT_FORMATS, DATA_DIR
|
||
except ImportError:
|
||
import sys
|
||
project_root = Path(__file__).resolve().parents[1]
|
||
if str(project_root) not in sys.path:
|
||
sys.path.insert(0, str(project_root))
|
||
from config import TAVILY_API_KEY, SEARCH_MAX_RESULTS, OUTPUT_FORMATS, DATA_DIR
|
||
|
||
from modules.i18n import tr
|
||
|
||
|
||
class SearchEngine:
|
||
def __init__(self):
|
||
self.api_key = TAVILY_API_KEY
|
||
self.api_url = "https://api.tavily.com/search"
|
||
|
||
self._valid_topics = {"general", "news", "finance"}
|
||
self._valid_time_ranges = {
|
||
"day": "day",
|
||
"d": "day",
|
||
"week": "week",
|
||
"w": "week",
|
||
"month": "month",
|
||
"m": "month",
|
||
"year": "year",
|
||
"y": "year"
|
||
}
|
||
self._date_pattern = re.compile(r"^\d{4}-\d{2}-\d{2}$")
|
||
|
||
async def search(
|
||
self,
|
||
query: str,
|
||
max_results: Optional[int] = None,
|
||
topic: Optional[str] = None,
|
||
time_range: Optional[str] = None,
|
||
days: Optional[int] = None,
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
country: Optional[str] = None,
|
||
include_domains: Optional[List[str]] = None
|
||
) -> Dict:
|
||
"""
|
||
执行网络搜索
|
||
|
||
Args:
|
||
query: 搜索关键词
|
||
max_results: 最大结果数
|
||
topic: 搜索类型（general/news/finance）
|
||
time_range: 相对时间范围（day/week/month/year 或 d/w/m/y）
|
||
days: 过去N天，仅topic=news可用
|
||
start_date: 起始日期，格式YYYY-MM-DD
|
||
end_date: 结束日期，格式YYYY-MM-DD
|
||
country: 国家过滤，仅topic=general可用
|
||
include_domains: 仅包含这些域名（最多300个）
|
||
|
||
Returns:
|
||
搜索结果字典
|
||
"""
|
||
if not self.api_key or self.api_key == "your-tavily-api-key":
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.api_key_not_configured"),
|
||
"results": []
|
||
}
|
||
|
||
validation = self._build_payload(
|
||
query=query,
|
||
max_results=max_results,
|
||
topic=topic,
|
||
time_range=time_range,
|
||
days=days,
|
||
start_date=start_date,
|
||
end_date=end_date,
|
||
country=country,
|
||
include_domains=include_domains
|
||
)
|
||
|
||
if not validation["success"]:
|
||
return validation
|
||
|
||
payload = validation["payload"]
|
||
applied_filters = validation["filters"]
|
||
|
||
max_results = payload.get("max_results", SEARCH_MAX_RESULTS)
|
||
|
||
print(f"{OUTPUT_FORMATS['search']} 搜索: {query}")
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=30) as client:
|
||
response = await client.post(
|
||
self.api_url,
|
||
json={
|
||
**payload
|
||
},
|
||
headers={
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
)
|
||
|
||
if response.status_code != 200:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.api_request_failed", status_code=response.status_code),
|
||
"results": []
|
||
}
|
||
|
||
data = response.json()
|
||
|
||
# 格式化结果
|
||
formatted_results = self._format_results(data, applied_filters)
|
||
|
||
print(f"{OUTPUT_FORMATS['success']} 搜索完成，找到 {len(formatted_results['results'])} 条结果")
|
||
|
||
return formatted_results
|
||
|
||
except httpx.TimeoutException:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.search_timeout"),
|
||
"results": []
|
||
}
|
||
except Exception as e:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.search_failed", error=str(e)),
|
||
"results": []
|
||
}
|
||
|
||
def _format_results(self, raw_data: Dict, filters: Dict[str, Any]) -> Dict:
|
||
"""格式化搜索结果"""
|
||
formatted = {
|
||
"success": True,
|
||
"query": raw_data.get("query", ""),
|
||
"answer": raw_data.get("answer", ""),
|
||
"results": [],
|
||
"timestamp": datetime.now().isoformat(),
|
||
"filters": filters,
|
||
"total_results": len(raw_data.get("results", []))
|
||
}
|
||
|
||
# 处理每个搜索结果
|
||
for idx, result in enumerate(raw_data.get("results", []), 1):
|
||
url = result.get("url", "")
|
||
formatted_result = {
|
||
"index": idx,
|
||
"title": result.get("title", "无标题"),
|
||
"url": url,
|
||
"domain": urlparse(url).netloc.lower() if url else "",
|
||
"content": result.get("content", ""),
|
||
"score": result.get("score", 0),
|
||
"published_date": result.get("published_date", "")
|
||
}
|
||
formatted["results"].append(formatted_result)
|
||
|
||
return formatted
|
||
|
||
def build_summary_text(
|
||
self,
|
||
query: str,
|
||
results: List[Dict[str, Any]],
|
||
filters: Dict[str, Any],
|
||
timestamp: str
|
||
) -> str:
|
||
"""构建给模型看的搜索摘要文本。
|
||
|
||
若结果项带 citation_id（tools_execution 注册 citation 后回填），
|
||
标题行会带 [src_xxx] 前缀，供模型在行内引用中使用。
|
||
"""
|
||
summary_lines = [
|
||
f"🔍 搜索查询: {query}",
|
||
f"📅 搜索时间: {timestamp}"
|
||
]
|
||
|
||
filter_notes = self._summarize_filters(filters or {})
|
||
if filter_notes:
|
||
summary_lines.append(filter_notes)
|
||
summary_lines.append("")
|
||
|
||
# 添加搜索结果
|
||
if results:
|
||
summary_lines.append("📊 搜索结果:")
|
||
|
||
for result in results:
|
||
cid = result.get("citation_id")
|
||
title_line = f"\n{result['index']}. [{cid}] {result['title']}" if cid else f"\n{result['index']}. {result['title']}"
|
||
summary_lines.extend([
|
||
title_line,
|
||
f" 🔗 {result['url']}",
|
||
f" 📄 {result['content'][:200]}..." if len(result['content']) > 200 else f" 📄 {result['content']}",
|
||
])
|
||
|
||
if result.get("published_date"):
|
||
summary_lines.append(f" 📅 发布时间: {result['published_date']}")
|
||
else:
|
||
summary_lines.append("未找到相关结果")
|
||
|
||
return "\n".join(summary_lines)
|
||
|
||
async def search_with_summary(
|
||
self,
|
||
query: str,
|
||
max_results: Optional[int] = None,
|
||
topic: Optional[str] = None,
|
||
time_range: Optional[str] = None,
|
||
days: Optional[int] = None,
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
country: Optional[str] = None,
|
||
include_domains: Optional[List[str]] = None
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
搜索并返回格式化的摘要
|
||
|
||
Args:
|
||
query: 搜索关键词
|
||
max_results: 最大结果数
|
||
|
||
Returns:
|
||
格式化的搜索摘要字符串
|
||
"""
|
||
results = await self.search(
|
||
query=query,
|
||
max_results=max_results,
|
||
topic=topic,
|
||
time_range=time_range,
|
||
days=days,
|
||
start_date=start_date,
|
||
end_date=end_date,
|
||
country=country,
|
||
include_domains=include_domains
|
||
)
|
||
|
||
if not results["success"]:
|
||
return {
|
||
"success": False,
|
||
"error": results.get("error", tr("search_engine.unknown_error")),
|
||
"summary": ""
|
||
}
|
||
|
||
return {
|
||
"success": True,
|
||
"summary": self.build_summary_text(
|
||
query,
|
||
results["results"],
|
||
results.get("filters", {}),
|
||
results["timestamp"]
|
||
),
|
||
"timestamp": results["timestamp"],
|
||
"filters": results.get("filters", {}),
|
||
"query": results.get("query", query),
|
||
"results": results.get("results", []),
|
||
"total_results": results.get("total_results", len(results.get("results", [])))
|
||
}
|
||
|
||
async def quick_answer(self, query: str) -> str:
|
||
"""
|
||
快速获取答案（返回首个搜索结果的摘要）
|
||
|
||
Args:
|
||
query: 查询问题
|
||
|
||
Returns:
|
||
首个结果摘要或错误信息
|
||
"""
|
||
results = await self.search(query, max_results=5)
|
||
|
||
if not results["success"]:
|
||
return tr("search_engine.search_failed", error=results["error"])
|
||
|
||
# 返回第一个结果的摘要
|
||
if results["results"]:
|
||
first_result = results["results"][0]
|
||
return f"{first_result['title']}\n{first_result['content'][:300]}..."
|
||
|
||
return tr("search_engine.no_relevant_info")
|
||
|
||
def save_results(self, results: Dict, filename: str = None) -> str:
|
||
"""
|
||
保存搜索结果到文件
|
||
|
||
Args:
|
||
results: 搜索结果
|
||
filename: 文件名（可选）
|
||
|
||
Returns:
|
||
保存的文件路径
|
||
"""
|
||
if filename is None:
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
filename = f"search_{timestamp}.json"
|
||
|
||
file_path = Path(DATA_DIR).expanduser().resolve() / "searches" / filename
|
||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 保存结果
|
||
with file_path.open('w', encoding='utf-8') as f:
|
||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||
|
||
print(f"{OUTPUT_FORMATS['file']} 搜索结果已保存到: {file_path}")
|
||
|
||
return str(file_path)
|
||
|
||
def load_results(self, filename: str) -> Optional[Dict]:
|
||
"""
|
||
加载之前的搜索结果
|
||
|
||
Args:
|
||
filename: 文件名
|
||
|
||
Returns:
|
||
搜索结果字典或None
|
||
"""
|
||
file_path = Path(DATA_DIR).expanduser().resolve() / "searches" / filename
|
||
|
||
try:
|
||
with file_path.open('r', encoding='utf-8') as f:
|
||
return json.load(f)
|
||
except FileNotFoundError:
|
||
print(f"{OUTPUT_FORMATS['error']} 文件不存在: {file_path}")
|
||
return None
|
||
except Exception as e:
|
||
print(f"{OUTPUT_FORMATS['error']} 加载失败: {e}")
|
||
return None
|
||
|
||
def _build_payload(
|
||
self,
|
||
query: str,
|
||
max_results: Optional[int],
|
||
topic: Optional[str],
|
||
time_range: Optional[str],
|
||
days: Optional[int],
|
||
start_date: Optional[str],
|
||
end_date: Optional[str],
|
||
country: Optional[str],
|
||
include_domains: Optional[List[str]]
|
||
) -> Dict[str, Any]:
|
||
"""验证并构建 Tavily 请求参数"""
|
||
payload: Dict[str, Any] = {
|
||
"query": query,
|
||
"search_depth": "advanced",
|
||
"include_answer": False,
|
||
"include_images": False,
|
||
"include_raw_content": False
|
||
}
|
||
|
||
filters: Dict[str, Any] = {}
|
||
|
||
if max_results:
|
||
payload["max_results"] = max_results
|
||
else:
|
||
payload["max_results"] = SEARCH_MAX_RESULTS
|
||
|
||
normalized_topic = (topic or "general").strip().lower()
|
||
if not normalized_topic:
|
||
normalized_topic = "general"
|
||
if normalized_topic not in self._valid_topics:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.invalid_topic", topic=topic, valid=", ".join(self._valid_topics)),
|
||
"results": []
|
||
}
|
||
payload["topic"] = normalized_topic
|
||
filters["topic"] = normalized_topic
|
||
|
||
# 时间参数互斥检查
|
||
has_time_range = bool(time_range)
|
||
has_days = days is not None
|
||
has_date_range = bool(start_date or end_date)
|
||
selected_filters = sum([has_time_range, has_days, has_date_range])
|
||
if selected_filters > 1:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.time_params_mutually_exclusive"),
|
||
"results": []
|
||
}
|
||
|
||
# 验证 days
|
||
if has_days:
|
||
try:
|
||
days_value = int(days) # type: ignore[arg-type]
|
||
except (TypeError, ValueError):
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.days_must_be_positive_int", days=days),
|
||
"results": []
|
||
}
|
||
if days_value <= 0:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.days_must_be_greater_than_zero", days=days_value),
|
||
"results": []
|
||
}
|
||
if normalized_topic != "news":
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.days_only_for_news"),
|
||
"results": []
|
||
}
|
||
payload["days"] = days_value
|
||
filters["days"] = days_value
|
||
|
||
# 验证 time_range
|
||
if has_time_range:
|
||
normalized_range = time_range.strip().lower() # type: ignore[union-attr]
|
||
normalized_range = self._valid_time_ranges.get(normalized_range, "")
|
||
if not normalized_range:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.invalid_time_range", time_range=time_range),
|
||
"results": []
|
||
}
|
||
payload["time_range"] = normalized_range
|
||
filters["time_range"] = normalized_range
|
||
|
||
# 验证日期范围
|
||
if has_date_range:
|
||
if not start_date or not end_date:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.date_range_requires_both"),
|
||
"results": []
|
||
}
|
||
if not self._date_pattern.match(start_date):
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.start_date_invalid_format", start_date=start_date),
|
||
"results": []
|
||
}
|
||
if not self._date_pattern.match(end_date):
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.end_date_invalid_format", end_date=end_date),
|
||
"results": []
|
||
}
|
||
try:
|
||
start_dt = datetime.fromisoformat(start_date)
|
||
end_dt = datetime.fromisoformat(end_date)
|
||
except ValueError:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.invalid_calendar_date"),
|
||
"results": []
|
||
}
|
||
if start_dt > end_dt:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.start_date_after_end_date", start_date=start_date, end_date=end_date),
|
||
"results": []
|
||
}
|
||
payload["start_date"] = start_date
|
||
payload["end_date"] = end_date
|
||
filters["start_date"] = start_date
|
||
filters["end_date"] = end_date
|
||
|
||
# 国家过滤
|
||
if country:
|
||
normalized_country = country.strip().lower()
|
||
if normalized_country:
|
||
if normalized_topic != "general":
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.country_only_for_general"),
|
||
"results": []
|
||
}
|
||
payload["country"] = normalized_country
|
||
filters["country"] = normalized_country
|
||
|
||
# 域名白名单
|
||
if include_domains is not None:
|
||
if not isinstance(include_domains, list):
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.include_domains_must_be_array"),
|
||
"results": []
|
||
}
|
||
cleaned_domains = []
|
||
for item in include_domains:
|
||
if not isinstance(item, str):
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.include_domains_item_must_be_string"),
|
||
"results": []
|
||
}
|
||
domain = item.strip().lower()
|
||
if domain:
|
||
cleaned_domains.append(domain)
|
||
if len(cleaned_domains) > 300:
|
||
return {
|
||
"success": False,
|
||
"error": tr("search_engine.include_domains_too_many", count=len(cleaned_domains)),
|
||
"results": []
|
||
}
|
||
if cleaned_domains:
|
||
payload["include_domains"] = cleaned_domains
|
||
filters["include_domains"] = cleaned_domains
|
||
|
||
return {
|
||
"success": True,
|
||
"payload": payload,
|
||
"filters": filters,
|
||
"results": []
|
||
}
|
||
|
||
def _summarize_filters(self, filters: Dict[str, Any]) -> str:
|
||
"""构建过滤条件摘要"""
|
||
if not filters:
|
||
return ""
|
||
|
||
parts = []
|
||
topic = filters.get("topic")
|
||
if topic:
|
||
parts.append(f"Topic: {topic}")
|
||
|
||
if "time_range" in filters:
|
||
parts.append(f"Time Range: {filters['time_range']}")
|
||
elif "days" in filters:
|
||
parts.append(f"最近 {filters['days']} 天")
|
||
elif "start_date" in filters and "end_date" in filters:
|
||
parts.append(f"{filters['start_date']} 至 {filters['end_date']}")
|
||
|
||
if "country" in filters:
|
||
parts.append(f"Country: {filters['country']}")
|
||
if "include_domains" in filters:
|
||
parts.append(f"Domains: {len(filters['include_domains'])}")
|
||
|
||
if not parts:
|
||
return ""
|
||
|
||
return "🎯 过滤条件: " + " | ".join(parts)
|