戒酒的李白
251 lines
10 KiB
Python
251 lines
10 KiB
Python
"""
|
|
专为 AI Agent 设计的舆情搜索工具集 (Tavily)
|
|
|
|
版本: 1.5
|
|
最后更新: 2025-08-22
|
|
|
|
此脚本将复杂的Tavily搜索功能分解为一系列目标明确、参数极少的独立工具，
|
|
专为AI Agent调用而设计。Agent只需根据任务意图选择合适的工具，
|
|
无需理解复杂的参数组合。所有工具默认搜索“新闻”(topic='news')。
|
|
|
|
新特性:
|
|
- 新增 `basic_search_news` 工具，用于执行标准、通用的新闻搜索。
|
|
- 每个搜索结果现在都包含 `published_date` (新闻发布日期)。
|
|
|
|
主要工具:
|
|
- basic_search_news: (新增) 执行标准、快速的通用新闻搜索。
|
|
- deep_search_news: 对主题进行最全面的深度分析。
|
|
- search_news_last_24_hours: 获取24小时内的最新动态。
|
|
- search_news_last_week: 获取过去一周的主要报道。
|
|
- search_images_for_news: 查找与新闻主题相关的图片。
|
|
- search_news_by_date: 在指定的历史日期范围内搜索。
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from typing import List, Dict, Any, Optional
|
|
|
|
# 添加utils目录到Python路径
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
root_dir = os.path.dirname(os.path.dirname(current_dir))
|
|
utils_dir = os.path.join(root_dir, 'utils')
|
|
if utils_dir not in sys.path:
|
|
sys.path.append(utils_dir)
|
|
|
|
from retry_helper import with_graceful_retry, SEARCH_API_RETRY_CONFIG
|
|
from dataclasses import dataclass, field
|
|
|
|
# 运行前请确保已安装Tavily库: pip install tavily-python
|
|
try:
|
|
from tavily import TavilyClient
|
|
except ImportError:
|
|
raise ImportError("Tavily库未安装，请运行 `pip install tavily-python` 进行安装。")
|
|
|
|
# --- 1. 数据结构定义 ---
|
|
|
|
@dataclass
|
|
class SearchResult:
|
|
"""
|
|
网页搜索结果数据类
|
|
包含 published_date 属性来存储新闻发布日期
|
|
"""
|
|
title: str
|
|
url: str
|
|
content: str
|
|
score: Optional[float] = None
|
|
raw_content: Optional[str] = None
|
|
published_date: Optional[str] = None
|
|
|
|
@dataclass
|
|
class ImageResult:
|
|
"""图片搜索结果数据类"""
|
|
url: str
|
|
description: Optional[str] = None
|
|
|
|
@dataclass
|
|
class TavilyResponse:
|
|
"""封装Tavily API的完整返回结果，以便在工具间传递"""
|
|
query: str
|
|
answer: Optional[str] = None
|
|
results: List[SearchResult] = field(default_factory=list)
|
|
images: List[ImageResult] = field(default_factory=list)
|
|
response_time: Optional[float] = None
|
|
|
|
|
|
# --- 2. 核心客户端与专用工具集 ---
|
|
|
|
class TavilyNewsAgency:
|
|
"""
|
|
一个包含多种专用新闻舆情搜索工具的客户端。
|
|
每个公共方法都设计为供 AI Agent 独立调用的工具。
|
|
"""
|
|
|
|
def __init__(self, api_key: Optional[str] = None):
|
|
"""
|
|
初始化客户端。
|
|
Args:
|
|
api_key: Tavily API密钥，若不提供则从环境变量 TAVILY_API_KEY 读取。
|
|
"""
|
|
if api_key is None:
|
|
api_key = os.getenv("TAVILY_API_KEY")
|
|
if not api_key:
|
|
raise ValueError("Tavily API Key未找到！请设置TAVILY_API_KEY环境变量或在初始化时提供")
|
|
self._client = TavilyClient(api_key=api_key)
|
|
|
|
@with_graceful_retry(SEARCH_API_RETRY_CONFIG, default_return=TavilyResponse(query="搜索失败"))
|
|
def _search_internal(self, **kwargs) -> TavilyResponse:
|
|
"""内部通用的搜索执行器，所有工具最终都调用此方法"""
|
|
try:
|
|
kwargs['topic'] = 'general'
|
|
api_params = {k: v for k, v in kwargs.items() if v is not None}
|
|
response_dict = self._client.search(**api_params)
|
|
|
|
search_results = [
|
|
SearchResult(
|
|
title=item.get('title'),
|
|
url=item.get('url'),
|
|
content=item.get('content'),
|
|
score=item.get('score'),
|
|
raw_content=item.get('raw_content'),
|
|
published_date=item.get('published_date')
|
|
) for item in response_dict.get('results', [])
|
|
]
|
|
|
|
image_results = [ImageResult(url=item.get('url'), description=item.get('description')) for item in response_dict.get('images', [])]
|
|
|
|
return TavilyResponse(
|
|
query=response_dict.get('query'), answer=response_dict.get('answer'),
|
|
results=search_results, images=image_results,
|
|
response_time=response_dict.get('response_time')
|
|
)
|
|
except Exception as e:
|
|
print(f"搜索时发生错误: {str(e)}")
|
|
raise e # 让重试机制捕获并处理
|
|
|
|
# --- Agent 可用的工具方法 ---
|
|
|
|
def basic_search_news(self, query: str, max_results: int = 7) -> TavilyResponse:
|
|
"""
|
|
【工具】基础新闻搜索: 执行一次标准、快速的新闻搜索。
|
|
这是最常用的通用搜索工具，适用于不确定需要何种特定搜索时。
|
|
Agent可提供搜索查询(query)和可选的最大结果数(max_results)。
|
|
"""
|
|
print(f"--- TOOL: 基础新闻搜索 (query: {query}) ---")
|
|
return self._search_internal(
|
|
query=query,
|
|
max_results=max_results,
|
|
search_depth="basic",
|
|
include_answer=False
|
|
)
|
|
|
|
def deep_search_news(self, query: str) -> TavilyResponse:
|
|
"""
|
|
【工具】深度新闻分析: 对一个主题进行最全面、最深入的搜索。
|
|
返回AI生成的“高级”详细摘要答案和最多20条最相关的新闻结果。适用于需要全面了解某个事件背景的场景。
|
|
Agent只需提供搜索查询(query)。
|
|
"""
|
|
print(f"--- TOOL: 深度新闻分析 (query: {query}) ---")
|
|
return self._search_internal(
|
|
query=query, search_depth="advanced", max_results=20, include_answer="advanced"
|
|
)
|
|
|
|
def search_news_last_24_hours(self, query: str) -> TavilyResponse:
|
|
"""
|
|
【工具】搜索24小时内新闻: 获取关于某个主题的最新动态。
|
|
此工具专门查找过去24小时内发布的新闻。适用于追踪突发事件或最新进展。
|
|
Agent只需提供搜索查询(query)。
|
|
"""
|
|
print(f"--- TOOL: 搜索24小时内新闻 (query: {query}) ---")
|
|
return self._search_internal(query=query, time_range='d', max_results=10)
|
|
|
|
def search_news_last_week(self, query: str) -> TavilyResponse:
|
|
"""
|
|
【工具】搜索本周新闻: 获取关于某个主题过去一周内的主要新闻报道。
|
|
适用于进行周度舆情总结或回顾。
|
|
Agent只需提供搜索查询(query)。
|
|
"""
|
|
print(f"--- TOOL: 搜索本周新闻 (query: {query}) ---")
|
|
return self._search_internal(query=query, time_range='w', max_results=10)
|
|
|
|
def search_images_for_news(self, query: str) -> TavilyResponse:
|
|
"""
|
|
【工具】查找新闻图片: 搜索与某个新闻主题相关的图片。
|
|
此工具会返回图片链接及描述，适用于需要为报告或文章配图的场景。
|
|
Agent只需提供搜索查询(query)。
|
|
"""
|
|
print(f"--- TOOL: 查找新闻图片 (query: {query}) ---")
|
|
return self._search_internal(
|
|
query=query, include_images=True, include_image_descriptions=True, max_results=5
|
|
)
|
|
|
|
def search_news_by_date(self, query: str, start_date: str, end_date: str) -> TavilyResponse:
|
|
"""
|
|
【工具】按指定日期范围搜索新闻: 在一个明确的历史时间段内搜索新闻。
|
|
这是唯一需要Agent提供详细时间参数的工具。适用于需要对特定历史事件进行分析的场景。
|
|
Agent需要提供查询(query)、开始日期(start_date)和结束日期(end_date)，格式均为 'YYYY-MM-DD'。
|
|
"""
|
|
print(f"--- TOOL: 按指定日期范围搜索新闻 (query: {query}, from: {start_date}, to: {end_date}) ---")
|
|
return self._search_internal(
|
|
query=query, start_date=start_date, end_date=end_date, max_results=15
|
|
)
|
|
|
|
|
|
# --- 3. 测试与使用示例 ---
|
|
|
|
def print_response_summary(response: TavilyResponse):
|
|
"""简化的打印函数，用于展示测试结果，现在会显示发布日期"""
|
|
if not response or not response.query:
|
|
print("未能获取有效响应。")
|
|
return
|
|
|
|
print(f"\n查询: '{response.query}' | 耗时: {response.response_time}s")
|
|
if response.answer:
|
|
print(f"AI摘要: {response.answer[:120]}...")
|
|
print(f"找到 {len(response.results)} 条网页, {len(response.images)} 张图片。")
|
|
if response.results:
|
|
first_result = response.results[0]
|
|
date_info = f"(发布于: {first_result.published_date})" if first_result.published_date else ""
|
|
print(f"第一条结果: {first_result.title} {date_info}")
|
|
print("-" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# 在运行前，请确保您已设置 TAVILY_API_KEY 环境变量
|
|
|
|
try:
|
|
# 初始化“新闻社”客户端，它内部包含了所有工具
|
|
agency = TavilyNewsAgency()
|
|
|
|
# 场景1: Agent 进行一次常规、快速的搜索
|
|
response1 = agency.basic_search_news(query="奥运会最新赛况", max_results=5)
|
|
print_response_summary(response1)
|
|
|
|
# 场景2: Agent 需要全面了解“全球芯片技术竞争”的背景
|
|
response2 = agency.deep_search_news(query="全球芯片技术竞争")
|
|
print_response_summary(response2)
|
|
|
|
# 场景3: Agent 需要追踪“GTC大会”的最新消息
|
|
response3 = agency.search_news_last_24_hours(query="Nvidia GTC大会 最新发布")
|
|
print_response_summary(response3)
|
|
|
|
# 场景4: Agent 需要为一篇关于“自动驾驶”的周报查找素材
|
|
response4 = agency.search_news_last_week(query="自动驾驶商业化落地")
|
|
print_response_summary(response4)
|
|
|
|
# 场景5: Agent 需要查找“韦伯太空望远镜”的新闻图片
|
|
response5 = agency.search_images_for_news(query="韦伯太空望远镜最新发现")
|
|
print_response_summary(response5)
|
|
|
|
# 场景6: Agent 需要研究2025年第一季度关于“人工智能法规”的新闻
|
|
response6 = agency.search_news_by_date(
|
|
query="人工智能法规",
|
|
start_date="2025-01-01",
|
|
end_date="2025-03-31"
|
|
)
|
|
print_response_summary(response6)
|
|
|
|
except ValueError as e:
|
|
print(f"初始化失败: {e}")
|
|
print("请确保 TAVILY_API_KEY 环境变量已正确设置。")
|
|
except Exception as e:
|
|
print(f"测试过程中发生未知错误: {e}") |