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,
]