69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
import os
|
|
from typing import List, Dict
|
|
from tavily import TavilyClient
|
|
from src.utils.logging_config import setup_logger
|
|
|
|
logger = setup_logger("web-search")
|
|
|
|
class WebSearcher:
|
|
def __init__(self):
|
|
api_key = os.getenv("TAVILY_API_KEY")
|
|
if not api_key:
|
|
raise ValueError("TAVILY_API_KEY environment variable is not set")
|
|
self.client = TavilyClient(api_key)
|
|
logger.info("WebSearcher initialized with Tavily client")
|
|
|
|
def search(self, query: str, max_results: int = 1) -> List[Dict]:
|
|
"""
|
|
使用 Tavily 搜索相关内容
|
|
|
|
Args:
|
|
query: 搜索查询
|
|
max_results: 最大返回结果数
|
|
|
|
Returns:
|
|
搜索结果列表
|
|
"""
|
|
try:
|
|
search_results = self.client.search(
|
|
query=query,
|
|
search_depth="basic",
|
|
max_results=max_results
|
|
)
|
|
|
|
# 提取需要的信息
|
|
formatted_results = []
|
|
for result in search_results['results'][:max_results]:
|
|
formatted_results.append({
|
|
'title': result.get('title', ''),
|
|
'content': result.get('content', ''),
|
|
'url': result.get('url', ''),
|
|
'score': result.get('score', 0)
|
|
})
|
|
|
|
return formatted_results
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during web search: {str(e)}")
|
|
return []
|
|
|
|
def format_search_results(self, results: List[Dict]) -> str:
|
|
"""
|
|
将搜索结果格式化为文本
|
|
|
|
Args:
|
|
results: 搜索结果列表
|
|
|
|
Returns:
|
|
格式化后的文本
|
|
"""
|
|
if not results:
|
|
return "没有找到相关的网络搜索结果。"
|
|
|
|
formatted_text = "以下是相关的网络搜索结果：\n\n"
|
|
for i, result in enumerate(results, 1):
|
|
formatted_text += f"{i}. {result['title']}\n"
|
|
formatted_text += f" {result['content']}\n"
|
|
formatted_text += f" 来源: {result['url']}\n\n"
|
|
|
|
return formatted_text |