| from bs4 import BeautifulSoup |
|
|
| from langchain_community.document_loaders import WikipediaLoader |
| from langchain_community.document_loaders import ArxivLoader |
| from langchain_community.tools.tavily_search import TavilySearchResults |
|
|
| from langchain_core.tools import tool |
| |
|
|
| from playwright.sync_api import sync_playwright |
|
|
|
|
| @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 = "\n\n---\n\n".join( |
| [ |
| f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>' |
| for doc in 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 = "\n\n---\n\n".join( |
| [ |
| f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>' |
| for doc in 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() |
| formatted_search_docs = "\n\n---\n\n".join( |
| [ |
| f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>' |
| for doc in search_docs |
| ] |
| ) |
| return {"arxiv_results": formatted_search_docs} |
|
|
|
|
| @tool |
| def website_scrape(url: str, question: str) -> str: |
| """Scrapes a website and returns the text. |
| Args: |
| url: the URL to the website to scrape. |
| question: the question to answer when searching for the answer in the website. |
| Returns: |
| str: The text of the website. |
| """ |
|
|
| with sync_playwright() as p: |
| browser = p.chromium.launch(headless=True) |
| page = browser.new_page() |
| page.goto(url) |
| html_content = page.content() |
| browser.close() |
|
|
| soup = BeautifulSoup(html_content, "html.parser") |
|
|
| |
| text = soup.get_text() |
|
|
| return text |
|
|
