|
|
from smolagents import DuckDuckGoSearchTool, Tool,tool |
|
|
import wikipediaapi |
|
|
|
|
|
class WikipediaSearchTool(Tool): |
|
|
name = "wikipedia_search" |
|
|
description = "查找英文维基百科的页面简介，输入应为一个词或短词组" |
|
|
inputs = { |
|
|
"query": {"type": "string", "description": "维基百科搜索关键词，例如人名/专名"} |
|
|
} |
|
|
output_type = "string" |
|
|
def __init__(self, lang="en"): |
|
|
super().__init__() |
|
|
self.wiki = wikipediaapi.Wikipedia(language=lang, user_agent="celum") |
|
|
def forward(self, query: str): |
|
|
page = self.wiki.page(query) |
|
|
if not page.exists(): |
|
|
return "No Wikipedia page found." |
|
|
return page.summary[:1000] |
|
|
|
|
|
@tool |
|
|
def tavily_search(query: str) -> str: |
|
|
""" |
|
|
Search the web using Tavily API |
|
|
Args: |
|
|
query: The search query |
|
|
Returns: |
|
|
Search results as formatted text |
|
|
""" |
|
|
api_key = os.getenv("TAVILY_API_KEY") |
|
|
if not api_key: |
|
|
return "Tavily API key not found" |
|
|
|
|
|
url = "https://api.tavily.com/search" |
|
|
payload = { |
|
|
"api_key": api_key, |
|
|
"query": query, |
|
|
"search_depth": "basic", |
|
|
"include_answer": True, |
|
|
"include_domains": [], |
|
|
"exclude_domains": [], |
|
|
"max_results": 5 |
|
|
} |
|
|
|
|
|
try: |
|
|
response = requests.post(url, json=payload, timeout=10) |
|
|
response.raise_for_status() |
|
|
data = response.json() |
|
|
|
|
|
results = [] |
|
|
if data.get("answer"): |
|
|
results.append(f"Quick Answer: {data['answer']}") |
|
|
|
|
|
for result in data.get("results", [])[:3]: |
|
|
results.append(f"Title: {result.get('title', 'N/A')}") |
|
|
results.append(f"URL: {result.get('url', 'N/A')}") |
|
|
results.append(f"Content: {result.get('content', 'N/A')[:200]}...") |
|
|
results.append("---") |
|
|
|
|
|
return "\n".join(results) |
|
|
except Exception as e: |
|
|
return f"Tavily search error: {str(e)}" |
|
|
|
|
|
my_tool_list = [ |
|
|
WikipediaSearchTool(), |
|
|
DuckDuckGoSearchTool() |
|
|
tavily_search, |
|
|
] |
