tools/web_search.py

from langchain_core.tools.base import BaseTool
from dotenv import load_dotenv
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_community.tools import TavilySearchResults, DuckDuckGoSearchResults
from langchain_tavily import TavilySearch
import os
from pydantic import PrivateAttr
from langchain_community.document_loaders import WebBaseLoader
import json
import requests

load_dotenv(".env", override=True)


class WebSearchTool(BaseTool):
    name: str = "web_search_tool"
    description: str = "Perform a web search and extract concise factual answers. Use for online facts not in GAIA/Wikipedia—e.g. sports stats, Olympic participation, published papers, museum specimen locations, competition winners, and other up-to-date info."
    #_search: BraveSearch = PrivateAttr()
    _search: DuckDuckGoSearchResults = PrivateAttr()

    def __init__(self):
        super().__init__()
        #wrapper = DuckDuckGoSearchAPIWrapper(region="en", max_results=2)
        #self._search = DuckDuckGoSearchResults(api_wrapper=wrapper, output_format="json")

        self._search = TavilySearch(max_results=2)

    def _run_old(self, query: str) -> str:
        json_str = self._search.run(query)          # list[Document]
        docs = json.loads(json_str)
        urls = [doc["link"] for doc in docs]

        pages = [requests.get(url) for url in urls]

        res = "\n\n---\n\n".join(
            page.text for page in pages
        )

        try:
            with open("./web_search.txt", "wt", encoding="utf-8") as f:
                f.write(str(res))
        except Exception as e:
            print(e)

        return res
    
    def _run(self, query: str) -> str:
        # import pdb;pdb.set_trace()
        search_results = []
        search_results.append(self._search.invoke(query))
        # print(f"Search results: {search_results} \n type: {type(search_results)}")
        return str(search_results)