Add some basic tools

gentopia/tools/calculator.py

@@ -0,0 +1,23 @@
+from typing import AnyStr
+
+from langchain import OpenAI, LLMMathChain
+
+from .basetool import *
+
+
+class Calculator(BaseTool):
+    """docstring for Calculator"""
+    name = "calculator"
+    description = "A calculator that can compute arithmetic expressions. Useful when you need to perform " \
+                  "math calculations. Input should be a mathematical expression"
+    args_schema = create_model("CalculatorArgs", expression=(str, ...))
+
+    def _run(self, expression: AnyStr) -> Any:
+        llm = OpenAI(temperature=0)
+        tool = LLMMathChain(llm=llm, verbose=self.verbose)
+        response = tool(input)
+        evidence = response["answer"].replace("Answer:", "").strip()
+        return evidence
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError

gentopia/tools/doc_store_look_up.py

@@ -0,0 +1,25 @@
+from typing import AnyStr
+
+from .basetool import *
+
+
+class DocStoreLookUp(BaseTool):
+    """docstring for DocStoreLookUp"""
+    name = "doc_store_look_up"
+    description = "Worker that search the direct sentence in current Wikipedia result page. Useful when you " \
+                  "need to find information about a specific keyword from a existing Wikipedia search " \
+                  "result. Input should be a search keyword."
+    args_schema = create_model("DocStoreLookUpArgs", query=(str, ...))
+
+    def __init__(self, doc_store=None):
+        super().__init__()
+        self.doc_store = doc_store
+
+    def _run(self, query: AnyStr) -> AnyStr:
+        assert self.doc_store is not None, "doc_store is not set"
+        tool = self.doc_store
+        evidence = tool.search(query)
+        return evidence
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError

gentopia/tools/dummy_tool.py

@@ -0,0 +1 @@
+

gentopia/tools/google_search.py

@@ -0,0 +1,22 @@
+from typing import AnyStr
+
+from langchain import SerpAPIWrapper
+
+from .basetool import *
+
+
+class GoogleSearch(BaseTool):
+    """Tool that adds the capability to query the Google search API."""
+
+    name = "GoogleSearch"
+    description = "Worker that searches results from Google. Useful when you need to find short " \
+                  "and succinct answers about a specific topic. Input should be a search query."
+
+    args_schema = create_model("GoogleSearchArgs", query=(str, ...))
+
+    def _run(self, query: AnyStr) -> AnyStr:
+        tool = SerpAPIWrapper()
+        return tool.run(query)
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError

gentopia/tools/llm.py

@@ -0,0 +1,27 @@
+from typing import AnyStr
+
+from langchain import OpenAI, PromptTemplate, LLMChain
+
+from .basetool import *
+
+
+class LLM(BaseTool):
+    """docstring for LLM"""
+    name = "llm"
+    description = "A pretrained LLM like yourself. Useful when you need to act with general world " \
+                  "knowledge and common sense. Prioritize it when you are confident in solving the problem " \
+                  "yourself. Input can be any instruction."
+
+    args_schema = create_model("LLMArgs", text=(str, ...))
+
+    def _run(self, text: AnyStr) -> AnyStr:
+        llm = OpenAI(temperature=0)
+        prompt = PromptTemplate(template="Respond in short directly with no extra words.\n\n{request}",
+                                input_variables=["request"])
+        tool = LLMChain(prompt=prompt, llm=llm, verbose=False)
+        response = tool(input)
+        evidence = response["text"].strip("\n")
+        return evidence
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError

gentopia/tools/search_doc.py

@@ -0,0 +1,27 @@
+from typing import AnyStr
+
+from langchain.document_loaders import TextLoader
+from langchain.indexes import VectorstoreIndexCreator
+
+from .basetool import *
+
+
+class SearchDoc(BaseTool):
+    name = "SearchDoc"
+    args_schema = create_model("SearchDocArgs", query=(str, ...))
+
+    def __init__(self, doc_name, doc_path):
+        super().__init__()
+        self.doc_path = doc_path
+        self.description = f"A vector store that searches for similar and related content in document: {doc_name}. " \
+                           f"The result is a huge chunk of text related to your search but can also " \
+                           f"contain irrelevant info. Input should be a search query."
+
+    def _run(self, query: AnyStr) -> AnyStr:
+        loader = TextLoader(self.doc_path)
+        vector_store = VectorstoreIndexCreator().from_loaders([loader]).vectorstore
+        evidence = vector_store.similarity_search(query, k=1)[0].page_content
+        return evidence
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError

gentopia/tools/wikipedia.py

@@ -0,0 +1,31 @@
+from typing import AnyStr
+
+from langchain import Wikipedia
+from langchain.agents.react.base import DocstoreExplorer
+
+from .basetool import *
+
+
+class WikipediaSearch(BaseTool):
+    """Tool that adds the capability to query the Wikipedia API."""
+
+    name = "Wikipedia"
+    description = "Worker that search for similar page contents from Wikipedia. Useful when you need to " \
+                  "get holistic knowledge about people, places, companies, historical events, " \
+                  "or other subjects. The response are long and might contain some irrelevant information. " \
+                  "Input should be a search query."
+    args_schema = create_model("WikipediaArgs", query=(str, ...))
+
+    def __init__(self, doc_store=None):
+        super().__init__()
+        self.doc_store = doc_store
+
+    def _run(self, query: AnyStr) -> AnyStr:
+        if not self.doc_store:
+            self.doc_store = DocstoreExplorer(Wikipedia())
+        tool = self.doc_store
+        evidence = tool.search(query)
+        return evidence
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError

gentopia/tools/wolfram_alpha.py

@@ -0,0 +1,39 @@
+from typing import AnyStr
+
+from langchain import WolframAlphaAPIWrapper
+
+from .basetool import *
+
+
+class CustomWolframAlphaAPITool(WolframAlphaAPIWrapper):
+    def __init__(self):
+        super().__init__()
+
+    def run(self, query: str) -> str:
+        """Run query through WolframAlpha and parse result."""
+        res = self.wolfram_client.query(query)
+
+        try:
+            answer = next(res.results).text
+        except StopIteration:
+            return "Wolfram Alpha wasn't able to answer it"
+
+        if answer is None or answer == "":
+            return "No good Wolfram Alpha Result was found"
+        else:
+            return f"Answer: {answer}"
+
+
+class WolframAlphaWorker(BaseTool):
+    name = "wolfram_alpha"
+    description = "A WolframAlpha search engine. Useful when you need to solve a complicated Mathematical or " \
+                  "Algebraic equation. Input should be an equation or function."
+    args_schema = create_model("WolframAlphaArgs", query=(str, ...))
+
+    def _run(self, query: AnyStr) -> AnyStr:
+        tool = CustomWolframAlphaAPITool()
+        evidence = tool.run(query).replace("Answer:", "").strip()
+        return evidence
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError

gentopia/tools/zip_code.py

@@ -0,0 +1,40 @@
+from typing import AnyStr
+
+import requests
+from geopy import Nominatim
+
+from .basetool import *
+
+
+class ZipCodeRetriever(BaseTool):
+    name = "zip_code_retriever"
+    description = "A zip code retriever. Useful when you need to get users' current zip code. Input can be " \
+                  "left blank."
+    args_schema = create_model("ZipCodeRetrieverArgs")
+
+    def get_ip_address(self):
+        response = requests.get("https://ipinfo.io/json")
+        data = response.json()
+        return data["ip"]
+
+    def get_location_data(sefl, ip_address):
+        url = f"https://ipinfo.io/{ip_address}/json"
+        response = requests.get(url)
+        data = response.json()
+        return data
+
+    def get_zipcode_from_lat_long(self, lat, long):
+        geolocator = Nominatim(user_agent="zipcode_locator")
+        location = geolocator.reverse((lat, long))
+        return location.raw["address"]["postcode"]
+
+    def get_current_zipcode(self):
+        ip_address = self.get_ip_address()
+        location_data = self.get_location_data(ip_address)
+        lat, long = location_data["loc"].split(",")
+        zipcode = self.get_zipcode_from_lat_long(float(lat), float(long))
+        return zipcode
+
+    def _run(self) -> AnyStr:
+        evidence = self.get_current_zipcode()
+        return evidence