Yuzhiran Dev
3fab87ee11
- search_utils.py: 新增 _call_360/_call_sogou/_call_wechat HTML爬取函数 - initial_data.py: 种子数据新增三个搜索提供商（priority 3/4/5） - models.py: provider_type 注释补充新类型 - admin.html: 搜索提供商类型下拉框新增三个选项 - AGENTS.md/PROGRESS.md/README.md: PostgreSQL 15→16 - README.md: 移除硬编码数据库密码
396 lines
15 KiB
Python
396 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
统一搜索工具：DB驱动多提供商自动降级
|
||
"""
|
||
import json, logging, os, sys, time
|
||
from pathlib import Path
|
||
from typing import List, Dict, Optional
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||
sys.path.insert(0, str(PROJECT_ROOT))
|
||
sys.path.insert(0, str(PROJECT_ROOT / "platform" / "backend"))
|
||
|
||
try:
|
||
from app.database import SessionLocal
|
||
from app.models import SearchProvider
|
||
HAVE_DB = True
|
||
except ImportError:
|
||
HAVE_DB = False
|
||
|
||
SEARCH_CACHE_FILE = PROJECT_ROOT / "automation" / "data" / "search_cache.json"
|
||
|
||
|
||
def _get_providers() -> List[Dict]:
|
||
"""从 DB 加载启用的搜索提供商，按优先级排序（自动跨日重置用量）"""
|
||
if not HAVE_DB:
|
||
return []
|
||
try:
|
||
import datetime as _dt
|
||
db = SessionLocal()
|
||
today = _dt.date.today()
|
||
|
||
rows = db.query(SearchProvider).filter(
|
||
SearchProvider.enabled == True
|
||
).order_by(SearchProvider.priority).all()
|
||
|
||
needs_commit = False
|
||
for r in rows:
|
||
if (r.usage_today or 0) > 0 and r.last_used_at:
|
||
last_date = r.last_used_at
|
||
if hasattr(last_date, 'date'):
|
||
last_date = last_date.date()
|
||
elif isinstance(last_date, _dt.datetime):
|
||
last_date = last_date.date()
|
||
if last_date < today:
|
||
r.usage_today = 0
|
||
needs_commit = True
|
||
if needs_commit:
|
||
db.commit()
|
||
|
||
db.close()
|
||
return [r.to_dict() if hasattr(r, 'to_dict') else {
|
||
"id": r.id, "name": r.name, "provider_type": r.provider_type,
|
||
"api_key": r.api_key or "", "api_url": r.api_url or "",
|
||
"priority": r.priority, "daily_limit": r.daily_limit,
|
||
"usage_today": getattr(r, 'usage_today', 0),
|
||
} for r in rows]
|
||
except Exception as e:
|
||
logger.warning("加载搜索提供商失败: %s", e)
|
||
return []
|
||
|
||
|
||
def reset_all_usage():
|
||
"""手动重置所有提供商当日用量（供 API/定时任务调用）"""
|
||
if not HAVE_DB:
|
||
return
|
||
try:
|
||
db = SessionLocal()
|
||
db.query(SearchProvider).update({SearchProvider.usage_today: 0})
|
||
db.commit()
|
||
db.close()
|
||
logger.info("所有搜索提供商用量已重置")
|
||
except Exception as e:
|
||
logger.warning("重置用量失败: %s", e)
|
||
|
||
|
||
def _increment_usage(provider_id: int):
|
||
"""增加提供商当日用量"""
|
||
if not HAVE_DB:
|
||
return
|
||
try:
|
||
db = SessionLocal()
|
||
p = db.query(SearchProvider).filter(SearchProvider.id == provider_id).first()
|
||
if p:
|
||
p.usage_today = (p.usage_today or 0) + 1
|
||
p.last_used_at = __import__('datetime').datetime.now(__import__('datetime').timezone.utc)
|
||
db.commit()
|
||
db.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _call_baidu(api_key: str, api_url: str, query: str, max_results: int) -> List[Dict]:
|
||
import requests
|
||
resp = requests.post(
|
||
api_url or "https://qianfan.baidubce.com/v2/ai_search/web_search",
|
||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||
json={
|
||
"messages": [{"role": "user", "content": query}],
|
||
"search_source": "baidu_search_v2",
|
||
"resource_type_filter": [{"type": "web", "top_k": max_results}],
|
||
},
|
||
timeout=15
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.warning("百度搜索返回 %s: %s", resp.status_code, resp.text[:100])
|
||
return []
|
||
data = resp.json()
|
||
results = data.get("results", []) or data.get("webPages", {}).get("value", [])
|
||
return [{
|
||
"title": r.get("title", "")[:120],
|
||
"url": r.get("url", "") or r.get("link", ""),
|
||
"content": r.get("snippet", "") or r.get("content", "") or r.get("summary", "")[:300],
|
||
"source": "baidu",
|
||
} for r in results[:max_results]]
|
||
|
||
|
||
def _call_qiniu(api_key: str, api_url: str, query: str, max_results: int) -> List[Dict]:
|
||
import requests
|
||
resp = requests.post(
|
||
api_url or "https://api.qnaigc.com/v1/search/web",
|
||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||
json={"query": query, "max_results": max_results, "search_type": "web"},
|
||
timeout=15
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.warning("七牛搜索返回 %s: %s", resp.status_code, resp.text[:100])
|
||
return []
|
||
data = resp.json()
|
||
items = data.get("results", data.get("data", []))
|
||
return [{
|
||
"title": r.get("title", "")[:120],
|
||
"url": r.get("url", ""),
|
||
"content": r.get("content", "") or r.get("snippet", "")[:300],
|
||
"source": "qiniu",
|
||
} for r in items[:max_results]]
|
||
|
||
|
||
def _call_tinyfish(api_key: str, api_url: str, query: str, max_results: int) -> List[Dict]:
|
||
import requests
|
||
resp = requests.get(
|
||
api_url or "https://api.search.tinyfish.ai",
|
||
params={"query": query, "location": "CN", "language": "zh"},
|
||
headers={"X-API-Key": api_key},
|
||
timeout=15
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.warning("TinyFish搜索返回 %s: %s", resp.status_code, resp.text[:100])
|
||
return []
|
||
data = resp.json()
|
||
items = data.get("results", [])
|
||
return [{
|
||
"title": r.get("title", "")[:120],
|
||
"url": r.get("url", ""),
|
||
"content": r.get("snippet", "")[:300],
|
||
"source": "tinyfish",
|
||
} for r in items[:max_results]]
|
||
|
||
|
||
def _call_bing(api_key: str, api_url: str, query: str, max_results: int) -> List[Dict]:
|
||
import requests
|
||
resp = requests.get(
|
||
api_url or "https://api.bing.microsoft.com/v7.0/search",
|
||
params={"q": query, "count": max_results, "mkt": "zh-CN"},
|
||
headers={"Ocp-Apim-Subscription-Key": api_key},
|
||
timeout=15
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.warning("Bing搜索返回 %s: %s", resp.status_code, resp.text[:100])
|
||
return []
|
||
data = resp.json()
|
||
items = data.get("webPages", {}).get("value", [])
|
||
return [{
|
||
"title": r.get("name", "")[:120],
|
||
"url": r.get("url", ""),
|
||
"content": r.get("snippet", "")[:300],
|
||
"source": "bing",
|
||
} for r in items[:max_results]]
|
||
|
||
|
||
def _call_mcp(api_key: str, api_url: str, query: str, max_results: int) -> List[Dict]:
|
||
"""Call the MCP search server directly (no API key needed)."""
|
||
import subprocess
|
||
try:
|
||
r = subprocess.run(
|
||
[sys.executable, str(PROJECT_ROOT / "scripts" / "mcp_search_server.py"),
|
||
"--query", query],
|
||
capture_output=True, text=True, timeout=90,
|
||
)
|
||
if r.returncode != 0:
|
||
logger.warning("MCP搜索子进程返回非零: %s", r.stderr[:100])
|
||
return []
|
||
results = json.loads(r.stdout)
|
||
if isinstance(results, list):
|
||
for res in results:
|
||
res["source"] = "opencode"
|
||
return results[:max_results]
|
||
except json.JSONDecodeError as e:
|
||
logger.warning("MCP搜索JSON解析失败: %s", e)
|
||
except subprocess.TimeoutExpired:
|
||
logger.warning("MCP搜索超时 (90s)")
|
||
except Exception as e:
|
||
logger.warning("MCP搜索失败: %s", e)
|
||
return []
|
||
|
||
|
||
def _call_360(api_key: str, api_url: str, query: str, max_results: int) -> List[Dict]:
|
||
"""360搜索（HTML爬取，无需 API Key）"""
|
||
from bs4 import BeautifulSoup
|
||
import requests
|
||
try:
|
||
resp = requests.get(
|
||
api_url or "https://www.so.com/s",
|
||
params={"q": query},
|
||
headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"},
|
||
timeout=15,
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.warning("360搜索返回 %s", resp.status_code)
|
||
return []
|
||
soup = BeautifulSoup(resp.text, "html.parser")
|
||
results = []
|
||
for item in soup.select("li.res-list, li[class*=result], .rb"):
|
||
title_el = item.select_one("h3.res-title a, h3[class*=title] a, .res-title a")
|
||
if not title_el:
|
||
continue
|
||
title = title_el.get_text(strip=True)[:120]
|
||
url = title_el.get("href", "")
|
||
snippet_el = item.select_one("p.res-desc, p[class*=desc], .res-desc")
|
||
snippet = snippet_el.get_text(strip=True)[:300] if snippet_el else ""
|
||
results.append({"title": title, "url": url, "content": snippet, "source": "360"})
|
||
if len(results) >= max_results:
|
||
break
|
||
return results
|
||
except Exception as e:
|
||
logger.warning("360搜索失败: %s", e)
|
||
return []
|
||
|
||
|
||
def _call_sogou(api_key: str, api_url: str, query: str, max_results: int) -> List[Dict]:
|
||
"""搜狗搜索（HTML爬取，无需 API Key）"""
|
||
from bs4 import BeautifulSoup
|
||
import requests
|
||
try:
|
||
resp = requests.get(
|
||
api_url or "https://sogou.com/web",
|
||
params={"query": query},
|
||
headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"},
|
||
timeout=15,
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.warning("搜狗搜索返回 %s", resp.status_code)
|
||
return []
|
||
soup = BeautifulSoup(resp.text, "html.parser")
|
||
results = []
|
||
for item in soup.select("div.vrwrap, div[class*=vr], .rb"):
|
||
title_el = item.select_one("h3.vr-title a, h3[class*=title] a, .vr-title a")
|
||
if not title_el:
|
||
continue
|
||
title = title_el.get_text(strip=True)[:120]
|
||
url = title_el.get("href", "")
|
||
snippet_el = item.select_one("p.str-text, div.str-text, p[class*=str]")
|
||
snippet = snippet_el.get_text(strip=True)[:300] if snippet_el else ""
|
||
results.append({"title": title, "url": url, "content": snippet, "source": "sogou"})
|
||
if len(results) >= max_results:
|
||
break
|
||
return results
|
||
except Exception as e:
|
||
logger.warning("搜狗搜索失败: %s", e)
|
||
return []
|
||
|
||
|
||
def _call_wechat(api_key: str, api_url: str, query: str, max_results: int) -> List[Dict]:
|
||
"""微信搜一搜（通过搜狗抓取，无需 API Key）"""
|
||
from bs4 import BeautifulSoup
|
||
import requests
|
||
try:
|
||
resp = requests.get(
|
||
api_url or "https://wx.sogou.com/weixin",
|
||
params={"type": 2, "query": query},
|
||
headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"},
|
||
timeout=15,
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.warning("微信搜索返回 %s", resp.status_code)
|
||
return []
|
||
soup = BeautifulSoup(resp.text, "html.parser")
|
||
results = []
|
||
for item in soup.select("div.news-box, li.news-list, div[class*=news]"):
|
||
title_el = item.select_one("h3 a, .txt-box h3 a")
|
||
if not title_el:
|
||
continue
|
||
title = title_el.get_text(strip=True)[:120]
|
||
url = title_el.get("href", "")
|
||
if url and not url.startswith("http"):
|
||
url = "https://wx.sogou.com" + url
|
||
snippet_el = item.select_one("p.txt-info, div.txt-info, .txt-info")
|
||
snippet = snippet_el.get_text(strip=True)[:300] if snippet_el else ""
|
||
results.append({"title": title, "url": url, "content": snippet, "source": "wechat"})
|
||
if len(results) >= max_results:
|
||
break
|
||
return results
|
||
except Exception as e:
|
||
logger.warning("微信搜索失败: %s", e)
|
||
return []
|
||
|
||
|
||
_PROVIDER_CALLS = {
|
||
"baidu": _call_baidu,
|
||
"qiniu": _call_qiniu,
|
||
"tinyfish": _call_tinyfish,
|
||
"bing": _call_bing,
|
||
"mcp": _call_mcp,
|
||
"360": _call_360,
|
||
"sogou": _call_sogou,
|
||
"wechat": _call_wechat,
|
||
}
|
||
|
||
|
||
def search(query: str, max_results: int = 5) -> List[Dict]:
|
||
"""统一搜索：DB提供商 → 本地缓存 → 空结果"""
|
||
providers = _get_providers()
|
||
for p in providers:
|
||
if (p.get("usage_today") or 0) >= (p.get("daily_limit") or 99999):
|
||
logger.info("提供商 %s 已达日限 %s，跳过", p.get("name"), p.get("daily_limit"))
|
||
continue
|
||
no_key_types = {"mcp", "360", "sogou", "wechat"}
|
||
if not p.get("api_key") and p.get("provider_type") not in no_key_types:
|
||
logger.info("提供商 %s 未配置 API Key，跳过", p.get("name"))
|
||
continue
|
||
call_fn = _PROVIDER_CALLS.get(p.get("provider_type"))
|
||
if not call_fn:
|
||
continue
|
||
try:
|
||
results = call_fn(p["api_key"], p.get("api_url", ""), query, max_results)
|
||
if results:
|
||
_increment_usage(p["id"])
|
||
logger.info("搜索 '%s' 通过 %s 获得 %d 条结果", query[:20], p.get("name"), len(results))
|
||
return results
|
||
logger.warning("提供商 %s 返回空结果", p.get("name"))
|
||
except Exception as e:
|
||
logger.warning("提供商 %s 失败: %s", p.get("name"), e)
|
||
continue
|
||
|
||
logger.info("搜索 '%s' 无结果（所有提供商均不可用）", query[:20])
|
||
return []
|
||
|
||
|
||
def search_from_cache(query: str, max_results: int = 5) -> List[Dict]:
|
||
"""从本地缓存读取搜索结果"""
|
||
if not SEARCH_CACHE_FILE.exists():
|
||
return []
|
||
try:
|
||
cache = json.loads(SEARCH_CACHE_FILE.read_text(encoding="utf-8"))
|
||
meta = cache.get("_metadata", {})
|
||
updated = meta.get("updated_at", "")
|
||
if updated:
|
||
import datetime
|
||
age = (datetime.datetime.now() - datetime.datetime.fromisoformat(updated)).total_seconds()
|
||
if age > 129600:
|
||
logger.warning("搜索缓存过时(%dh)，跳过", int(age // 3600))
|
||
return []
|
||
results = cache.get(query, [])
|
||
return results[:max_results]
|
||
except Exception:
|
||
return []
|
||
|
||
|
||
def enrich_topic_research(topic: dict, max_results: int = 5) -> str:
|
||
"""对选题进行网络搜索，返回格式化的研究发现文本"""
|
||
title = topic.get('title', '')
|
||
field = topic.get('field', '')
|
||
queries = [title]
|
||
if field and field not in title:
|
||
queries.append(f"{field} {title[:40]}")
|
||
seen_urls = set()
|
||
results = []
|
||
for q in queries:
|
||
for r in search(q, max_results):
|
||
url = r.get('url', '')
|
||
if url and url not in seen_urls:
|
||
seen_urls.add(url)
|
||
results.append(r)
|
||
if not results:
|
||
return ""
|
||
lines = ["\n## 网络搜索参考", ""]
|
||
for r in results[:max_results]:
|
||
snippet = r.get('snippet', r.get('content', ''))
|
||
lines.append(f"- **{r.get('title', '无标题')}**")
|
||
lines.append(f" {snippet[:200]}")
|
||
if r.get('url'):
|
||
lines.append(f" [{r['url']}]")
|
||
lines.append("")
|
||
return "\n".join(lines)
|