import os
from smolagents import Tool
from tavily import TavilyClient

class TavilyResearchTool(Tool):
    """
    Tavily deep-research search tool.
    Use this when the question needs academic papers,
    scientific background or research-level accuracy.
    """

    name = "tavily_research"
    description = (
        "Use Tavily deep research mode to find academic-level content "
        "including papers, research summaries, and high-quality sources."
    )

    inputs = {
        "query": {"type": "string", "description": "Research topic to search"}
    }

    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.")
        self.client = TavilyClient(api_key=api_key)

    def forward(self, query: str) -> str:
        try:
            response = self.client.search(
                query=query,
                search_depth="advanced",     # <-- academic/longform mode
                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 research error: {e}"