"""Chat agent with question answering
"""
from dotenv import load_dotenv
from langchain.cache import InMemoryCache
import langchain
import os
from dataclasses import dataclass

from langchain.chains import LLMChain, LLMRequestsChain
from langchain import Wikipedia, OpenAI
from langchain.agents.react.base import DocstoreExplorer
from langchain.agents import Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent
from langchain.prompts import PromptTemplate
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain.agents.conversational.base import ConversationalAgent
from datetime import datetime

# Load the environment variables
load_dotenv()

langchain.llm_cache = InMemoryCache()
news_api_key = os.getenv("NEWS_API_KEY")


@dataclass
class ChatAgent:
    agent_executor: AgentExecutor = None

    def _get_docstore_agent(self):
        docstore = DocstoreExplorer(Wikipedia())
        docstore_tools = [
            Tool(
                name="Search",
                func=docstore.search
            ),
            Tool(
                name="Lookup",
                func=docstore.lookup
            )
        ]
        docstore_llm = OpenAI(temperature=0, model_name="text-davinci-003")
        docstore_agent = initialize_agent(
            docstore_tools, docstore_llm, agent="react-docstore", verbose=True)
        return docstore_agent

    def _get_requests_llm_tool(self):

        template = """
        Extracted: {requests_result}"""

        PROMPT = PromptTemplate(
            input_variables=["requests_result"],
            template=template,
        )

        def lambda_func(input):
            out = chain = LLMRequestsChain(llm_chain=LLMChain(
                llm=OpenAI(temperature=0),
                prompt=PROMPT)).run(input)
            return out.strip()
        return lambda_func

    def __init__(self, *, conversation_chain: LLMChain = None, history_array):
        date = datetime.today().strftime('%A %d, %B, %Y, %I:%M%p')
        print("DATETIME:", date)

        # set up a Wikipedia docstore agent
        docstore_agent = self._get_docstore_agent()

        tool_names = get_all_tool_names()

        tool_names.remove("pal-math")
        tool_names.remove("requests")  # let's use the llm_requests instead
        # let's use the llm_requests instead
        tool_names.remove("google-search")
        tool_names.remove("pal-colored-objects")
        tool_names.remove("python_repl")
        tool_names.remove("terminal")
        tool_names.remove("serpapi"),
        tool_names.remove("tmdb-api")

        requests_tool = self._get_requests_llm_tool()

        print("ALL TOOLS:", tool_names)
        tools = load_tools(tool_names,
                           llm=OpenAI(temperature=0,
                                      model_name="text-davinci-003"),
                           news_api_key=news_api_key,
                           )

        # Tweak some of the tool descriptions
        for tool in tools:
            if tool.name == "Search":
                tool.description = "Use this tool exclusively for questions relating to current events, or when you can't find an answer using any of the other tools."
            if tool.name == "Calculator":
                tool.description = "Use this to solve numeric math questions and do arithmetic. Don't use it for general or abstract math questions."

        tools = tools + [
            Tool(
                name="WikipediaSearch",
                description="Useful for answering a wide range of factual, scientific, academic, political and historical questions.",
                func=docstore_agent.run
            ),
            Tool(
                name="Requests",
                func=requests_tool,
                description="A portal to the internet. Use this when you need to get specific content from a site. Input should be a specific url, and the output will be all the text on that page."
            )
        ]
        ai_prefix = "FWROG-E"
        human_prefix = "Bence"

        prefix = os.getenv("PROMPT_PREFIX")

        suffix = f"""
        The person's name that you are interacting with is {human_prefix}. Please be entertaining and respectful towards him.
        The current date is {date}. Questions that refer to a specific date or time period will be interpreted relative to this date.
After you answer the question, you MUST to determine which langauge your answer is written in, and append the language code to the end of the Final Answer, within parentheses, like this (en-US).
Begin!
Previous conversation history:
{{chat_history}}
New input: {{input}}
{{agent_scratchpad}}
"""

        memory = ConversationBufferMemory(memory_key="chat_history")
        for item in history_array:
            memory.save_context(
                {f"{ai_prefix}": item["prompt"]}, {f"{human_prefix}": item["response"]})

        llm = OpenAI(temperature=.5, max_tokens=384,
                     model_name="text-davinci-003")
        llm_chain = LLMChain(
            llm=llm,
            prompt=ConversationalAgent.create_prompt(
                tools,
                prefix=prefix,
                ai_prefix=ai_prefix,
                human_prefix=human_prefix,
                suffix=suffix
            ),
        )

        agent_obj = ConversationalAgent(
            llm_chain=llm_chain, ai_prefix=ai_prefix)

        self.agent_executor = AgentExecutor.from_agent_and_tools(
            agent=agent_obj,
            tools=tools,
            verbose=True,
            max_iterations=5,
            memory=memory)

langchain.agents.initialize_agent langchain.agents.AgentExecutor.from_agent_and_tools langchain.cache.InMemoryCache langchain.Wikipedia langchain.agents.conversational.base.ConversationalAgent langchain.agents.conversational.base.ConversationalAgent.create_prompt langchain.agents.Tool langchain.chains.conversation.memory.ConversationBufferMemory langchain.agents.get_all_tool_names langchain.prompts.PromptTemplate langchain.OpenAI
[((681, 694), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (692, 694), False, 'from dotenv import load_dotenv\n'), ((718, 733), 'langchain.cache.InMemoryCache', 'InMemoryCache', ([], {}), '()\n', (731, 733), False, 'from langchain.cache import InMemoryCache\n'), ((749, 774), 'os.getenv', 'os.getenv', (['"""NEWS_API_KEY"""'], {}), "('NEWS_API_KEY')\n", (758, 774), False, 'import os\n'), ((1192, 1244), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'model_name': '"""text-davinci-003"""'}), "(temperature=0, model_name='text-davinci-003')\n", (1198, 1244), False, 'from langchain import Wikipedia, OpenAI\n'), ((1270, 1358), 'langchain.agents.initialize_agent', 'initialize_agent', (['docstore_tools', 'docstore_llm'], {'agent': '"""react-docstore"""', 'verbose': '(True)'}), "(docstore_tools, docstore_llm, agent='react-docstore',\n    verbose=True)\n", (1286, 1358), False, 'from langchain.agents import Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((1519, 1589), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['requests_result']", 'template': 'template'}), "(input_variables=['requests_result'], template=template)\n", (1533, 1589), False, 'from langchain.prompts import PromptTemplate\n'), ((2164, 2184), 'langchain.agents.get_all_tool_names', 'get_all_tool_names', ([], {}), '()\n', (2182, 2184), False, 'from langchain.agents import Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((4068, 4094), 'os.getenv', 'os.getenv', (['"""PROMPT_PREFIX"""'], {}), "('PROMPT_PREFIX')\n", (4077, 4094), False, 'import os\n'), ((4692, 4743), 'langchain.chains.conversation.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""'}), "(memory_key='chat_history')\n", (4716, 4743), False, 'from langchain.chains.conversation.memory import ConversationBufferMemory\n'), ((4916, 4986), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0.5)', 'max_tokens': '(384)', 'model_name': '"""text-davinci-003"""'}), "(temperature=0.5, max_tokens=384, model_name='text-davinci-003')\n", (4922, 4986), False, 'from langchain import Wikipedia, OpenAI\n'), ((5322, 5383), 'langchain.agents.conversational.base.ConversationalAgent', 'ConversationalAgent', ([], {'llm_chain': 'llm_chain', 'ai_prefix': 'ai_prefix'}), '(llm_chain=llm_chain, ai_prefix=ai_prefix)\n', (5341, 5383), False, 'from langchain.agents.conversational.base import ConversationalAgent\n'), ((5428, 5544), 'langchain.agents.AgentExecutor.from_agent_and_tools', 'AgentExecutor.from_agent_and_tools', ([], {'agent': 'agent_obj', 'tools': 'tools', 'verbose': '(True)', 'max_iterations': '(5)', 'memory': 'memory'}), '(agent=agent_obj, tools=tools, verbose=\n    True, max_iterations=5, memory=memory)\n', (5462, 5544), False, 'from langchain.agents import Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((918, 929), 'langchain.Wikipedia', 'Wikipedia', ([], {}), '()\n', (927, 929), False, 'from langchain import Wikipedia, OpenAI\n'), ((970, 1011), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Search"""', 'func': 'docstore.search'}), "(name='Search', func=docstore.search)\n", (974, 1011), False, 'from langchain.agents import Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((1071, 1112), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Lookup"""', 'func': 'docstore.lookup'}), "(name='Lookup', func=docstore.lookup)\n", (1075, 1112), False, 'from langchain.agents import Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((1960, 1976), 'datetime.datetime.today', 'datetime.today', ([], {}), '()\n', (1974, 1976), False, 'from datetime import datetime\n'), ((2758, 2810), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'model_name': '"""text-davinci-003"""'}), "(temperature=0, model_name='text-davinci-003')\n", (2764, 2810), False, 'from langchain import Wikipedia, OpenAI\n'), ((3444, 3626), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""WikipediaSearch"""', 'description': '"""Useful for answering a wide range of factual, scientific, academic, political and historical questions."""', 'func': 'docstore_agent.run'}), "(name='WikipediaSearch', description=\n    'Useful for answering a wide range of factual, scientific, academic, political and historical questions.'\n    , func=docstore_agent.run)\n", (3448, 3626), False, 'from langchain.agents import Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((3692, 3926), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Requests"""', 'func': 'requests_tool', 'description': '"""A portal to the internet. Use this when you need to get specific content from a site. Input should be a specific url, and the output will be all the text on that page."""'}), "(name='Requests', func=requests_tool, description=\n    'A portal to the internet. Use this when you need to get specific content from a site. Input should be a specific url, and the output will be all the text on that page.'\n    )\n", (3696, 3926), False, 'from langchain.agents import Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((5077, 5199), 'langchain.agents.conversational.base.ConversationalAgent.create_prompt', 'ConversationalAgent.create_prompt', (['tools'], {'prefix': 'prefix', 'ai_prefix': 'ai_prefix', 'human_prefix': 'human_prefix', 'suffix': 'suffix'}), '(tools, prefix=prefix, ai_prefix=ai_prefix,\n    human_prefix=human_prefix, suffix=suffix)\n', (5110, 5199), False, 'from langchain.agents.conversational.base import ConversationalAgent\n'), ((1741, 1762), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)'}), '(temperature=0)\n', (1747, 1762), False, 'from langchain import Wikipedia, OpenAI\n')]