Omnitopia's picture
添加媒体分析工具,tavily_search
9ddc718 verified
raw
history blame
4.43 kB
from smolagents import DuckDuckGoSearchTool, Tool, tool
import wikipediaapi
import os
import requests
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)}"
# 添加媒体分析工具
@tool
def analyze_media_file(file_path: str) -> str:
"""
分析媒体文件并返回基本信息
Args:
file_path: 文件路径
Returns:
文件信息和分析结果
"""
if not os.path.exists(file_path):
return f"File not found: {file_path}"
try:
file_size = os.path.getsize(file_path)
file_ext = os.path.splitext(file_path)[1].lower()
result = f"File: {os.path.basename(file_path)}\n"
result += f"Size: {file_size} bytes\n"
result += f"Extension: {file_ext}\n"
# 根据文件类型进行不同的分析
if file_ext in ['.png', '.jpg', '.jpeg', '.gif', '.bmp']:
try:
from PIL import Image
with Image.open(file_path) as img:
result += f"Image dimensions: {img.size[0]}x{img.size[1]}\n"
result += f"Image mode: {img.mode}\n"
result += f"Image format: {img.format}\n"
# 检测是否可能是棋盘
if abs(img.size[0] - img.size[1]) < 50:
result += "Note: Square aspect ratio - possibly a chess board\n"
except Exception as e:
result += f"Image analysis error: {str(e)}\n"
elif file_ext in ['.mp4', '.avi', '.mov', '.mkv']:
result += "Video file detected\n"
elif file_ext in ['.mp3', '.wav', '.m4a', '.flac']:
result += "Audio file detected\n"
elif file_ext in ['.pdf']:
result += "PDF file detected\n"
elif file_ext in ['.txt', '.csv', '.json']:
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
result += f"Text length: {len(content)} characters\n"
result += f"Line count: {len(content.splitlines())}\n"
preview = content[:200].replace('\n', ' ')
result += f"Content preview: {preview}...\n"
except Exception as e:
result += f"Text analysis error: {str(e)}\n"
return result
except Exception as e:
return f"Error analyzing file: {str(e)}"
my_tool_list = [
WikipediaSearchTool(),
DuckDuckGoSearchTool(),
tavily_search,
analyze_media_file,
]