tools/search.py

from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.utilities import GoogleSerperAPIWrapper
from dotenv import load_dotenv
from typing import Annotated
from typing_extensions import Annotated
from langchain_core.tools.base import InjectedToolCallId
from langchain_core.runnables import RunnableConfig
from langgraph.types import Command
from langchain_core.messages import ToolMessage

load_dotenv()

@tool
def taivily_search(query: str, tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig):
    """Tavily is a search engine optimized for LLMs, aimed at efficient, quick and persistent search results. Tavily take care of all the burden of searching, scraping, filtering and extracting the most relevant information from online sources.
    Args:
        query: The search query."""
    
    search_docs = TavilySearchResults(max_results=1).invoke(input=query)
    
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f"""# Taivily \n source="{doc["url"]}"" \n ## Content \n {doc["content"]} \n---"""
            for doc in search_docs
        ]
    )
    
    return Command(
        update={
            "external_information": f"{config.get('external_information', '')}\n\n---\n\n{formatted_search_docs}",
            "messages": [ToolMessage(content=formatted_search_docs, tool_call_id=tool_call_id)]
        }
    )

@tool
def serper_search(query: str, tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig) -> Command:
    """
    lightning-fast Google search results in 1-2 seconds.
    Args: 
        query: The search query.
    """
    search_docs = GoogleSerperAPIWrapper(k=8)
    result = search_docs.run(query)
    
    formatted_result = f"""# Google Search Results \n\n## Results \n\n{result}"""
    
    return Command(
        update={
            "external_information": f"{config.get('external_information', '')}\n\n---\n\n{formatted_result}",
            "messages": [ToolMessage(content=formatted_result, tool_call_id=tool_call_id)]
        }
    )