Create my_tools.py

my_tools.py

@@ -0,0 +1,178 @@
+import cmath
+import json
+
+from langchain_community.document_loaders import ArxivLoader
+from langchain_community.document_loaders import WebBaseLoader
+from langchain_community.tools import tool
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 page links.
+    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
+    #     ]
+    # )
+    formatted_search_docs = [
+        {"source": doc.metadata.get("source", ""), "page": doc.metadata.get("page", "")}
+        for doc in search_docs
+    ]
+    return {"wiki_page_links": json.dumps(formatted_search_docs)}
+
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 page links.
+    Args:
+        query: The search query."""
+    search_docs = TavilySearchResults(
+        max_results=3,
+        tavily_api_key="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",
+    ).invoke(query)
+    # formatted_search_docs = "\n\n---\n\n".join(
+    #     [
+    #         f'<Document source="{doc.get("url", "")}" title="{doc.get("title", "")}"/>\n{doc.get("content", "")}\n</Document>'
+    #         for doc in search_docs
+    #     ]
+    # )
+    formatted_search_docs = [
+        {"source": doc.get("url", ""), "page": doc.get("title", "")}
+        for doc in search_docs
+    ]
+    return {"web_page_links": json.dumps(formatted_search_docs)}
+
+@tool
+def visit_webpage(url: str) -> str:
+    """Retrieve content from a webpage using page links. A good option when you need detailed information from a specific webpage.
+    Args:
+        url: The URL to retrieve content from."""
+    search_docs = WebBaseLoader(url).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata.get("source", "")}" title="{doc.metadata.get("title", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+    return {"web_page_content": 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["Title"]}" page="{doc.metadata.get("Summary", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+    return {"arxiv_results": formatted_search_docs}
+
+@tool
+def translate_to_english(text: str) -> str:
+    """Translate the given text to English.
+    Args:
+        text: The text to translate to English."""
+    from langchain_ollama import ChatOllama
+
+    llm = ChatOllama(
+        model="tinyllama:1.1b",
+        api_base="http://localhost:11434",
+    )
+
+    prompt = f"Translate the following text to English:\n\n{text}"
+    response = llm.invoke(
+        [
+            {"role": "system", "content": "You are a helpful assistant that translates text to English."},
+            {"role": "user", "content": prompt},
+        ]
+    )
+    return {"translated_text": response.content}
+
+@tool
+def multiply(a: float, b: float) -> float:
+    """
+    Multiplies two numbers.
+    Args:
+        a (float): the first number
+        b (float): the second number
+    """
+    return a * b
+
+
+@tool
+def add(a: float, b: float) -> float:
+    """
+    Adds two numbers.
+    Args:
+        a (float): the first number
+        b (float): the second number
+    """
+    return a + b
+
+
+@tool
+def subtract(a: float, b: float) -> int:
+    """
+    Subtracts two numbers.
+    Args:
+        a (float): the first number
+        b (float): the second number
+    """
+    return a - b
+
+
+@tool
+def divide(a: float, b: float) -> float:
+    """
+    Divides two numbers.
+    Args:
+        a (float): the first float number
+        b (float): the second float number
+    """
+    if b == 0:
+        raise ValueError("Cannot divided by zero.")
+    return a / b
+
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """
+    Get the modulus of two numbers.
+    Args:
+        a (int): the first number
+        b (int): the second number
+    """
+    return a % b
+
+
+@tool
+def power(a: float, b: float) -> float:
+    """
+    Get the power of two numbers.
+    Args:
+        a (float): the first number
+        b (float): the second number
+    """
+    return a**b
+
+
+@tool
+def square_root(a: float) -> float | complex:
+    """
+    Get the square root of a number.
+    Args:
+        a (float): the number to get the square root of
+    """
+    if a >= 0:
+        return a**0.5
+    return cmath.sqrt(a)