import requests from geopy.geocoders import Nominatim from langchain import OpenAI, LLMMathChain, LLMChain, PromptTemplate, Wikipedia from langchain.agents import Tool from langchain.agents.react.base import DocstoreExplorer from langchain.document_loaders import TextLoader from langchain.indexes import VectorstoreIndexCreator from langchain.utilities import SerpAPIWrapper from langchain.utilities.wolfram_alpha import WolframAlphaAPIWrapper from nodes.Node import Node class GoogleWorker(Node): def __init__(self, name="Google"): super().__init__(name, input_type=str, output_type=str) self.isLLMBased = False self.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." def run(self, input, log=False): assert isinstance(input, self.input_type) tool = SerpAPIWrapper() evidence = tool.run(input) assert isinstance(evidence, self.output_type) if log: print(f"Running {self.name} with input {input}\nOutput: {evidence}\n") return evidence class WikipediaWorker(Node): def __init__(self, name="Wikipedia", docstore=None): super().__init__(name, input_type=str, output_type=str) self.isLLMBased = False self.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." self.docstore = docstore def run(self, input, log=False): if not self.docstore: self.docstore = DocstoreExplorer(Wikipedia()) assert isinstance(input, self.input_type) tool = Tool( name="Search", func=self.docstore.search, description="useful for when you need to ask with search" ) evidence = tool.run(input) assert isinstance(evidence, self.output_type) if log: print(f"Running {self.name} with input {input}\nOutput: {evidence}\n") return evidence class DocStoreLookUpWorker(Node): def __init__(self, name="LookUp", docstore=None): super().__init__(name, input_type=str, output_type=str) self.isLLMBased = False self.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." self.docstore = docstore def run(self, input, log=False): if not self.docstore: raise ValueError("Docstore must be provided for lookup") assert isinstance(input, self.input_type) tool = Tool( name="Lookup", func=self.docstore.lookup, description="useful for when you need to ask with lookup" ) evidence = tool.run(input) assert isinstance(evidence, self.output_type) if log: print(f"Running {self.name} with input {input}\nOutput: {evidence}\n") return evidence 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(Node): def __init__(self, name="WolframAlpha"): super().__init__(name, input_type=str, output_type=str) self.isLLMBased = False self.description = "A WolframAlpha search engine. Useful when you need to solve a complicated Mathematical or " \ "Algebraic equation. Input should be an equation or function." def run(self, input, log=False): assert isinstance(input, self.input_type) tool = CustomWolframAlphaAPITool() evidence = tool.run(input).replace("Answer:", "").strip() assert isinstance(evidence, self.output_type) if log: print(f"Running {self.name} with input {input}\nOutput: {evidence}\n") return evidence class CalculatorWorker(Node): def __init__(self, name="Calculator"): super().__init__(name, input_type=str, output_type=str) self.isLLMBased = True self.description = "A calculator that can compute arithmetic expressions. Useful when you need to perform " \ "math calculations. Input should be a mathematical expression" def run(self, input, log=False): assert isinstance(input, self.input_type) llm = OpenAI(temperature=0) tool = LLMMathChain(llm=llm, verbose=False) response = tool(input) evidence = response["answer"].replace("Answer:", "").strip() assert isinstance(evidence, self.output_type) if log: return {"input": response["question"], "output": response["answer"]} return evidence class LLMWorker(Node): def __init__(self, name="LLM"): super().__init__(name, input_type=str, output_type=str) self.isLLMBased = True self.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." def run(self, input, log=False): assert isinstance(input, self.input_type) 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") assert isinstance(evidence, self.output_type) if log: return {"input": response["request"], "output": response["text"]} return evidence class ZipCodeRetriever(Node): def __init__(self, name="ZipCodeRetriever"): super().__init__(name, input_type=str, output_type=str) self.isLLMBased = False self.description = "A zip code retriever. Useful when you need to get users' current zip code. Input can be " \ "left blank." 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, input): assert isinstance(input, self.input_type) evidence = self.get_current_zipcode() assert isinstance(evidence, self.output_type) class SearchDocWorker(Node): def __init__(self, doc_name, doc_path, name="SearchDoc"): super().__init__(name, input_type=str, output_type=str) self.isLLMBased = True 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, input, log=False): assert isinstance(input, self.input_type) loader = TextLoader(self.doc_path) vectorstore = VectorstoreIndexCreator().from_loaders([loader]).vectorstore evidence = vectorstore.similarity_search(input, k=1)[0].page_content assert isinstance(evidence, self.output_type) if log: print(f"Running {self.name} with input {input}\nOutput: {evidence}\n") return evidence class SearchSOTUWorker(SearchDocWorker): def __init__(self, name="SearchSOTU"): super().__init__(name=name, doc_name="state_of_the_union", doc_path="data/docs/state_of_the_union.txt") WORKER_REGISTRY = {"Google": GoogleWorker(), "Wikipedia": WikipediaWorker(), "LookUp": DocStoreLookUpWorker(), "WolframAlpha": WolframAlphaWorker(), "Calculator": CalculatorWorker(), "LLM": LLMWorker(), "SearchSOTU": SearchSOTUWorker()}