langchain

[Proteusiq]

def query_wikipedia(title, first_paragraph_only=True):
  base_url = "https://en.wikipedia.org"
  url = f"{base_url}/w/api.php?format=json&action=query&prop=extracts&explaintext=1&titles={title}"
  if first_paragraph_only:
    url += "&exintro=1"
  data = requests.get(url).json()
  return Document(
    metadata={"source": f"{base_url}/wiki/{title}"},
    page_content=list(data["query"]["pages"].values())[0]["extract"],
  )

# example
from langchain.agents import Tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser
from langchain.prompts import StringPromptTemplate
from langchain import OpenAI, LLMChain
from langchain.utilities import WikipediaAPIWrapper
from typing import List, Union
from langchain.schema import AgentAction, AgentFinish
import re
from termcolor import colored
import os

# os.environ["OPENAI_API_KEY"] = 
# os.environ["SERPAPI_API_KEY"] = 

search = WikipediaAPIWrapper()

def search_wikipedia(input):

    result = search.run(input)

    if type(result) == str:
        return result[:5000]
    else:
        return "Agent could not find a result."

tools = [
    Tool(
        name="Wikipedia",
        description="Useful for finding information about a specific topic. You cannot use this tool to ask questions, only to find information about a specific topic.",
        func=search_wikipedia,
    )
]

template = """I want you to be FritzAgent. An agent that use tools to get answers. You are reliable and trustworthy. You follow the rules:

Rule 1: Answer the following questions as best as you can with the Observations presented to you.
Rule 2: Never use information outside of the Observations presented to you.
Rule 3: Never jump to conclusions unless the information for the final answer is explicitly presented to you in Observation.

You have access to the following tools:

{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
Thought: you should always think about what to do next. Use the Observation to gather extra information, but never use information outside of the Observation.
Action: the action to take, should be one of [{tool_names}]
Action_input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer.
Final Answer: the final answer to the original input question

Begin!

Question: {input}
{agent_scratchpad}
"""

class CustomPromptTemplate(StringPromptTemplate):

    template: str
    tools: List[Tool]

    def format(self, **kwargs) -> str:
        
        intermediate_steps = kwargs.pop("intermediate_steps")
        thoughts = ""

        for action, observation in intermediate_steps:
            thoughts += action.log
            thoughts += f"\nObservation: {observation}\nThought: "
        
        kwargs["agent_scratchpad"] = thoughts
        kwargs["tools"] = "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])
        kwargs["tool_names"] = ", ".join([tool.name for tool in self.tools])
        return self.template.format(**kwargs)

prompt = CustomPromptTemplate(template=template, tools=tools, input_variables=["input", "intermediate_steps"])

class CustomOutputParser(AgentOutputParser):

    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:

        if "Final Answer:" in llm_output:
            return AgentFinish(return_values={"output": llm_output.split("Final Answer:")[1].strip()},
                               log=llm_output
            )
        
        regex = r"Action: (.*?)[\n]*Action Input:[\s]*(.*)"
        match = re.search(regex, llm_output, re.DOTALL)
        if not match:
            raise ValueError(f"Could not parse output: {llm_output}")
        action = match.group(1).strip()
        action_input = match.group(2).strip()
        return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)
    
output_parser = CustomOutputParser()

llm = OpenAI(temperature=0)

llm_chain = LLMChain(llm=llm, prompt=prompt)
tool_names = [tool.name for tool in tools]
agent = LLMSingleActionAgent(llm_chain=llm_chain, output_parser=output_parser, stop=["\nObservation:"], allowed_tools=tool_names)

agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)

while True:
    user_input = input(colored("> ", "green", attrs=["bold"]))
    if user_input == "exit":
        break
    output = agent_executor.run(input=user_input)
    print(colored("FritzGPT:\n", "red"))
    print(output)