"""Tools to search the web and extract information."""
import os

from tavily import AsyncTavilyClient
from llama_index.core.tools import FunctionTool

async def tavily_web_tool(query: str) -> str:
    """
    Search the web with a given query and return the first two results.
    If results do not return what you are looking for, rephrasing the query can help.
    """
    client = AsyncTavilyClient(api_key=os.getenv("TAVILY"))
    res = await client.search(query)
    urls = [r.get("url") for r in res.get("results", []) if "url" in r][:2]
    page_contents = {
        url: await client.extract(url)
        for url in urls
    }
    return "\n\n".join(
        f"RESULT {i}:\n{content}"
        for i, content in enumerate(page_contents.values())
    )

def get_search_web_tool():
    return FunctionTool.from_defaults(
        fn=tavily_web_tool,
        description="Search the web with a given query and return the first two results."
    )