"""Chat agent with question answering

"""
import os
from utils.giphy import GiphyAPIWrapper
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 (
    ZeroShotAgent,
    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


import langchain
from langchain.cache import InMemoryCache

langchain.llm_cache = InMemoryCache()


news_api_key = os.environ["NEWS_API_KEY"]
tmdb_bearer_token = os.environ["TMDB_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, description="Search wikipedia"),
            Tool(
                name="Lookup",
                func=docstore.lookup,
                description="Lookup a wikipedia page",
            ),
        ]
        docstore_llm = OpenAI(temperature=0, model_name="gpt-3.5-turbo")
        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 = LLMRequestsChain(
                llm_chain=LLMChain(
                    llm=OpenAI(temperature=0), prompt=PROMPT, verbose=True
                )
            ).run(input)
            return out.strip()

        return lambda_func

    def __init__(self, *, conversation_chain: LLMChain = None, history_array):
        date = datetime.today().strftime("%B %d, %Y")

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

        giphy = GiphyAPIWrapper()

        # tool_names = get_all_tool_names()
        tool_names = [
            "wolfram-alpha",
            "llm-math",
            "open-meteo-api",
            "news-api",
            "tmdb-api",
            "wikipedia",
        ]

        requests_tool = self._get_requests_llm_tool()

        tools = load_tools(
            tool_names,
            llm=OpenAI(temperature=0, model_name="gpt-3.5-turbo", verbose=True),
            news_api_key=news_api_key,
            tmdb_bearer_token=tmdb_bearer_token,
        )

        # Tweak some of the tool descriptions
        for tool in tools:
            if tool.name == "Calculator":
                tool.description = (
                    "Use this only to solve numeric math problems and to do arithmetic."
                )

        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="GiphySearch",
                func=giphy.run,
                return_direct=True,
                description="useful for when you need to find a gif or picture, and for adding humor to your replies. Input should be a query, and output will be an html embed code which you MUST include in your Final Answer.",
            ),
            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.",
            ),
        ]

        # set up the google search tool if the env var is set
        if "GOOGLE_API_KEY" in os.environ:
            from langchain.utilities import GoogleSearchAPIWrapper

            tools.append(
                Tool(
                    name="Search",
                    func=GoogleSearchAPIWrapper().run,
                    description="Use this tool for questions relating to current events, or when you can't find an answer using any of the other tools.",
                )
            )

        # set up the serpapi search tool if the env var is set
        if "SERPAPI_API_KEY" in os.environ:
            from langchain import SerpAPIWrapper

            serpapi = SerpAPIWrapper()

            tools.append(
                Tool(
                    name="Search",
                    func=serpapi.run,
                    description="Use this tool for questions relating to current events, or when you can't find an answer using any of the other tools.",
                )
            )

        ai_prefix = "AI"
        human_prefix = "Human"

        prefix = f"""{ai_prefix} is a large language model. {ai_prefix} is represented by a 🤖.

{ai_prefix} uses a light, humorous tone, and very frequently includes emojis its responses. Responses with code examples should be formatted in code blocks using <pre><code></code></pre> tags.

{ai_prefix} is designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, {ai_prefix} is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.

If {ai_prefix} can't provide a good response, it will truthfully answer that it can't help with the user's request.

Overall, Assistant is a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether you need help with a specific question or just want to have a conversation about a particular topic, Assistant is here to assist.

TOOLS:
------

Assistant has access to the following tools:
"""

        suffix = f"""
The current date is {date}. Questions that refer to a specific date or time period will be interpreted relative to this 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=0.5, model_name="gpt-3.5-turbo")
        llm_chain = LLMChain(
            llm=llm,
            prompt=ConversationalAgent.create_prompt(
                tools,
                ai_prefix=ai_prefix,
                human_prefix=human_prefix,
                prefix=prefix,
                suffix=suffix,
            ),
            verbose=True,
        )

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

        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.utilities.GoogleSearchAPIWrapper langchain.Wikipedia langchain.agents.conversational.base.ConversationalAgent langchain.SerpAPIWrapper langchain.agents.conversational.base.ConversationalAgent.create_prompt langchain.agents.Tool langchain.chains.conversation.memory.ConversationBufferMemory langchain.cache.InMemoryCache langchain.prompts.PromptTemplate langchain.OpenAI
[((726, 741), 'langchain.cache.InMemoryCache', 'InMemoryCache', ([], {}), '()\n', (739, 741), False, 'from langchain.cache import InMemoryCache\n'), ((1293, 1342), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'model_name': '"""gpt-3.5-turbo"""'}), "(temperature=0, model_name='gpt-3.5-turbo')\n", (1299, 1342), False, 'from langchain import Wikipedia, OpenAI\n'), ((1368, 1456), '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", (1384, 1456), False, 'from langchain.agents import ZeroShotAgent, Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((1625, 1695), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['requests_result']", 'template': 'template'}), "(input_variables=['requests_result'], template=template)\n", (1639, 1695), False, 'from langchain.prompts import PromptTemplate\n'), ((2261, 2278), 'utils.giphy.GiphyAPIWrapper', 'GiphyAPIWrapper', ([], {}), '()\n', (2276, 2278), False, 'from utils.giphy import GiphyAPIWrapper\n'), ((6817, 6868), 'langchain.chains.conversation.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""'}), "(memory_key='chat_history')\n", (6841, 6868), False, 'from langchain.chains.conversation.memory import ConversationBufferMemory\n'), ((7054, 7105), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0.5)', 'model_name': '"""gpt-3.5-turbo"""'}), "(temperature=0.5, model_name='gpt-3.5-turbo')\n", (7060, 7105), False, 'from langchain import Wikipedia, OpenAI\n'), ((7448, 7523), 'langchain.agents.conversational.base.ConversationalAgent', 'ConversationalAgent', ([], {'llm_chain': 'llm_chain', 'ai_prefix': 'ai_prefix', 'verbose': '(True)'}), '(llm_chain=llm_chain, ai_prefix=ai_prefix, verbose=True)\n', (7467, 7523), False, 'from langchain.agents.conversational.base import ConversationalAgent\n'), ((7577, 7693), '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', (7611, 7693), False, 'from langchain.agents import ZeroShotAgent, Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((976, 987), 'langchain.Wikipedia', 'Wikipedia', ([], {}), '()\n', (985, 987), False, 'from langchain import Wikipedia, OpenAI\n'), ((1028, 1101), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Search"""', 'func': 'docstore.search', 'description': '"""Search wikipedia"""'}), "(name='Search', func=docstore.search, description='Search wikipedia')\n", (1032, 1101), False, 'from langchain.agents import ZeroShotAgent, Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((1115, 1200), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Lookup"""', 'func': 'docstore.lookup', 'description': '"""Lookup a wikipedia page"""'}), "(name='Lookup', func=docstore.lookup, description='Lookup a wikipedia page'\n    )\n", (1119, 1200), False, 'from langchain.agents import ZeroShotAgent, Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((4690, 4706), 'langchain.SerpAPIWrapper', 'SerpAPIWrapper', ([], {}), '()\n', (4704, 4706), False, 'from langchain import SerpAPIWrapper\n'), ((2108, 2124), 'datetime.datetime.today', 'datetime.today', ([], {}), '()\n', (2122, 2124), False, 'from datetime import datetime\n'), ((2637, 2700), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'model_name': '"""gpt-3.5-turbo"""', 'verbose': '(True)'}), "(temperature=0, model_name='gpt-3.5-turbo', verbose=True)\n", (2643, 2700), False, 'from langchain import Wikipedia, OpenAI\n'), ((3099, 3281), '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", (3103, 3281), False, 'from langchain.agents import ZeroShotAgent, Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((3348, 3630), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""GiphySearch"""', 'func': 'giphy.run', 'return_direct': '(True)', 'description': '"""useful for when you need to find a gif or picture, and for adding humor to your replies. Input should be a query, and output will be an html embed code which you MUST include in your Final Answer."""'}), "(name='GiphySearch', func=giphy.run, return_direct=True, description=\n    'useful for when you need to find a gif or picture, and for adding humor to your replies. Input should be a query, and output will be an html embed code which you MUST include in your Final Answer.'\n    )\n", (3352, 3630), False, 'from langchain.agents import ZeroShotAgent, Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((3713, 3947), '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", (3717, 3947), False, 'from langchain.agents import ZeroShotAgent, Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((4750, 4931), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Search"""', 'func': 'serpapi.run', 'description': '"""Use this tool for questions relating to current events, or when you can\'t find an answer using any of the other tools."""'}), '(name=\'Search\', func=serpapi.run, description=\n    "Use this tool for questions relating to current events, or when you can\'t find an answer using any of the other tools."\n    )\n', (4754, 4931), False, 'from langchain.agents import ZeroShotAgent, Tool, AgentExecutor, get_all_tool_names, load_tools, initialize_agent\n'), ((7176, 7299), 'langchain.agents.conversational.base.ConversationalAgent.create_prompt', 'ConversationalAgent.create_prompt', (['tools'], {'ai_prefix': 'ai_prefix', 'human_prefix': 'human_prefix', 'prefix': 'prefix', 'suffix': 'suffix'}), '(tools, ai_prefix=ai_prefix, human_prefix=\n    human_prefix, prefix=prefix, suffix=suffix)\n', (7209, 7299), False, 'from langchain.agents.conversational.base import ConversationalAgent\n'), ((4294, 4318), 'langchain.utilities.GoogleSearchAPIWrapper', 'GoogleSearchAPIWrapper', ([], {}), '()\n', (4316, 4318), False, 'from langchain.utilities import GoogleSearchAPIWrapper\n'), ((1860, 1881), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)'}), '(temperature=0)\n', (1866, 1881), False, 'from langchain import Wikipedia, OpenAI\n')]