# modules/search_engine.py - 网络搜索模块 import httpx import json from typing import Dict, Optional, Any from datetime import datetime import re try: from config import TAVILY_API_KEY, SEARCH_MAX_RESULTS, OUTPUT_FORMATS except ImportError: import sys from pathlib import Path 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 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 ) -> 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可用 Returns: 搜索结果字典 """ if not self.api_key or self.api_key == "your-tavily-api-key": return { "success": False, "error": "Tavily API密钥未配置", "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 ) 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": f"API请求失败: {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": "搜索超时", "results": [] } except Exception as e: return { "success": False, "error": f"搜索失败: {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): formatted_result = { "index": idx, "title": result.get("title", "无标题"), "url": result.get("url", ""), "content": result.get("content", ""), "score": result.get("score", 0), "published_date": result.get("published_date", "") } formatted["results"].append(formatted_result) return formatted 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 ) -> 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 ) if not results["success"]: return { "success": False, "error": results.get("error", "未知错误"), "summary": "" } # 构建摘要 summary_lines = [ f"🔍 搜索查询: {query}", f"📅 搜索时间: {results['timestamp']}" ] filter_notes = self._summarize_filters(results.get("filters", {})) if filter_notes: summary_lines.append(filter_notes) summary_lines.append("") # 添加AI答案（如果有） if results.get("answer"): summary_lines.extend([ "📝 AI摘要:", results["answer"], "", "---", "" ]) # 添加搜索结果 if results["results"]: summary_lines.append("📊 搜索结果:") for result in results["results"]: summary_lines.extend([ f"\n{result['index']}. {result['title']}", 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 { "success": True, "summary": "\n".join(summary_lines), "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: """ 快速获取答案（只返回AI摘要） Args: query: 查询问题 Returns: AI答案或错误信息 """ results = await self.search(query, max_results=5) if not results["success"]: return f"搜索失败: {results['error']}" if results.get("answer"): return results["answer"] # 如果没有AI答案，返回第一个结果的摘要 if results["results"]: first_result = results["results"][0] return f"{first_result['title']}\n{first_result['content'][:300]}..." return "未找到相关信息" 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 = f"./data/searches/{filename}" # 确保目录存在 import os os.makedirs(os.path.dirname(file_path), exist_ok=True) # 保存结果 with open(file_path, 'w', encoding='utf-8') as f: json.dump(results, f, ensure_ascii=False, indent=2) print(f"{OUTPUT_FORMATS['file']} 搜索结果已保存到: {file_path}") return file_path def load_results(self, filename: str) -> Optional[Dict]: """ 加载之前的搜索结果 Args: filename: 文件名 Returns: 搜索结果字典或None """ file_path = f"./data/searches/{filename}" try: with open(file_path, '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] ) -> Dict[str, Any]: """验证并构建 Tavily 请求参数""" payload: Dict[str, Any] = { "query": query, "search_depth": "advanced", "include_answer": True, "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": f"无效的topic: {topic}. 可选值: {', '.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": "时间参数只能三选一：time_range、days、start_date+end_date 不能同时使用", "results": [] } # 验证 days if has_days: try: days_value = int(days) # type: ignore[arg-type] except (TypeError, ValueError): return { "success": False, "error": f"days 必须是正整数，当前值: {days}", "results": [] } if days_value <= 0: return { "success": False, "error": f"days 必须大于0，当前值: {days_value}", "results": [] } if normalized_topic != "news": return { "success": False, "error": "days 参数仅在 topic=\"news\" 时可用，请调整 topic 或改用其他时间参数", "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": f"无效的time_range: {time_range}. 可选值: day/week/month/year 或缩写 d/w/m/y", "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": "start_date 与 end_date 必须同时提供且格式为 YYYY-MM-DD", "results": [] } if not self._date_pattern.match(start_date): return { "success": False, "error": f"start_date 格式无效: {start_date}，请使用 YYYY-MM-DD", "results": [] } if not self._date_pattern.match(end_date): return { "success": False, "error": f"end_date 格式无效: {end_date}，请使用 YYYY-MM-DD", "results": [] } try: start_dt = datetime.fromisoformat(start_date) end_dt = datetime.fromisoformat(end_date) except ValueError: return { "success": False, "error": "start_date 或 end_date 含无效日期，请检查是否为有效的公历日期", "results": [] } if start_dt > end_dt: return { "success": False, "error": f"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": "country 参数仅在 topic=\"general\" 时可用，请调整 topic 或移除 country", "results": [] } payload["country"] = normalized_country filters["country"] = normalized_country 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 not parts: return "" return "🎯 过滤条件: " + " | ".join(parts)