import copy from langchain_community.tools.tavily_search import TavilySearchResults from langchain_community.document_loaders import WikipediaLoader from langchain_community.document_loaders import ArxivLoader from langchain_core.tools import tool def format_search_docs(search_docs): """Format search documents into a consistent string format. Args: search_docs: List of document objects with metadata and page_content. Returns: Formatted string with document sources and content. """ return "\n\n---\n\n".join( [ f'\n{doc.page_content}\n' for doc in search_docs ] ) @tool def wiki_search(query: str) -> str: """Search Wikipedia for a query and return maximum 2 results. Args: query: The search query.""" search_docs = WikipediaLoader(query=query, load_max_docs=2).load() formatted_search_docs = format_search_docs(search_docs) return {"wiki_results": formatted_search_docs} @tool def web_search(query: str) -> str: """Search Tavily for a query and return maximum 3 results. Args: query: The search query.""" search_docs = TavilySearchResults(max_results=3).invoke(query=query) formatted_search_docs = format_search_docs(search_docs) return {"web_results": formatted_search_docs} @tool def arxiv_search(query: str) -> str: """Search Arxiv for a query and return maximum 3 result. Args: query: The search query.""" search_docs = ArxivLoader(query=query, load_max_docs=3).load() truncated_docs = [] for doc in search_docs: doc_copy = copy.copy(doc) doc_copy.page_content = doc.page_content[:1000] truncated_docs.append(doc_copy) formatted_search_docs = format_search_docs(truncated_docs) return {"arxiv_results": formatted_search_docs}