import langchain.llms
from langchain import GoogleSearchAPIWrapper, LLMChain
from langchain.agents import initialize_agent, AgentType, Tool, ZeroShotAgent, AgentExecutor
from langchain.schema import BaseMemory


def setup_agent(llm: langchain.llms.BaseLLM, memory: BaseMemory):
    search = GoogleSearchAPIWrapper()

    tools = [
        Tool(
            name="Google Search",
            func=search.run,
            description="Useful for when you need to answer questions about current events, the current state of the world or what you don't know."
        ),
    ]

    prefix = """Answer the following questions as best you can. You have access to the following tools:"""
    suffix = """Begin! Use lots of tools, and please answer finally in Japanese.

        Question: {input}
        {agent_scratchpad}"""

    prompt = ZeroShotAgent.create_prompt(
        tools,
        prefix=prefix,
        suffix=suffix,
        input_variables=["input", "agent_scratchpad"]
    )

    llm_chain = LLMChain(
        llm=llm,
        prompt=prompt)
    agent = ZeroShotAgent(
        llm_chain=llm_chain,
        tools=tools,
        memory=memory)
    agent_executor = AgentExecutor.from_agent_and_tools(
        agent=agent,
        tools=tools,
        verbose=True)

    return agent_executor

langchain.agents.AgentExecutor.from_agent_and_tools langchain.LLMChain langchain.agents.ZeroShotAgent.create_prompt langchain.agents.ZeroShotAgent langchain.GoogleSearchAPIWrapper langchain.agents.Tool
[((291, 315), 'langchain.GoogleSearchAPIWrapper', 'GoogleSearchAPIWrapper', ([], {}), '()\n', (313, 315), False, 'from langchain import GoogleSearchAPIWrapper, LLMChain\n'), ((833, 948), 'langchain.agents.ZeroShotAgent.create_prompt', 'ZeroShotAgent.create_prompt', (['tools'], {'prefix': 'prefix', 'suffix': 'suffix', 'input_variables': "['input', 'agent_scratchpad']"}), "(tools, prefix=prefix, suffix=suffix,\n    input_variables=['input', 'agent_scratchpad'])\n", (860, 948), False, 'from langchain.agents import initialize_agent, AgentType, Tool, ZeroShotAgent, AgentExecutor\n'), ((1000, 1032), 'langchain.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt'}), '(llm=llm, prompt=prompt)\n', (1008, 1032), False, 'from langchain import GoogleSearchAPIWrapper, LLMChain\n'), ((1062, 1124), 'langchain.agents.ZeroShotAgent', 'ZeroShotAgent', ([], {'llm_chain': 'llm_chain', 'tools': 'tools', 'memory': 'memory'}), '(llm_chain=llm_chain, tools=tools, memory=memory)\n', (1075, 1124), False, 'from langchain.agents import initialize_agent, AgentType, Tool, ZeroShotAgent, AgentExecutor\n'), ((1171, 1245), 'langchain.agents.AgentExecutor.from_agent_and_tools', 'AgentExecutor.from_agent_and_tools', ([], {'agent': 'agent', 'tools': 'tools', 'verbose': '(True)'}), '(agent=agent, tools=tools, verbose=True)\n', (1205, 1245), False, 'from langchain.agents import initialize_agent, AgentType, Tool, ZeroShotAgent, AgentExecutor\n'), ((339, 529), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Google Search"""', 'func': 'search.run', 'description': '"""Useful for when you need to answer questions about current events, the current state of the world or what you don\'t know."""'}), '(name=\'Google Search\', func=search.run, description=\n    "Useful for when you need to answer questions about current events, the current state of the world or what you don\'t know."\n    )\n', (343, 529), False, 'from langchain.agents import initialize_agent, AgentType, Tool, ZeroShotAgent, AgentExecutor\n')]