import os
from smolagents import Tool
from tavily import TavilyClient
class TavilySearchTool(Tool):
    """
    General-purpose Tavily web search.
    Use this FIRST before Wikipedia or other tools.
    Returns summarized results with sources.
    """
    
    name = "tavily_search"
    description = (
        "Use Tavily to search the web for up-to-date information. "
        "Returns summarized content and source links."
    )

    inputs = {
        "query": {"type": "string", "description": "Search query."}
    }

    output_type = "string"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        api_key = os.getenv("TAVILY_API_KEY")
        if not api_key:
            raise ValueError("Missing TAVILY_API_KEY in environment variables.")
        self.client = TavilyClient(api_key=api_key)

    def forward(self, query: str) -> str:
        try:
            response = self.client.search(query=query, max_results=5)
            out = []
            for r in response.get("results", []):
                out.append(
                    f"TITLE: {r.get('title')}\n"
                    f"CONTENT: {r.get('content')}\n"
                    f"SOURCE: {r.get('url')}"
                )
            return "\n\n---SEPARATOR---\n\n".join(out)
        except Exception as e:
            return f"Tavily search error: {e}"