from langchain.agents import OpenAIFunctionsAgent
from langchain.schema import SystemMessage
from langchain.agents import AgentExecutor
from langchain.tools.base import ToolException


class Agent:
    def __init__(self, llm, tools, max_iterations=15):
        for tool in tools:
            def _handle_error(error: ToolException) -> str:
                return (
                        f"The following errors occurred during tool '{tool.name}' execution:"
                        + error.args[0]
                        + "Please try another tool."
                )

            tool.handle_tool_error = _handle_error
        self._tools = tools
        system_message = SystemMessage(
            content="You are a web researcher who uses search engines to look up information and list references to sources.")
        prompt = OpenAIFunctionsAgent.create_prompt(system_message=system_message)
        agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)
        self.agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, max_iterations=max_iterations)

    def run(self, query):
        try:
            query = query + [" Along with your answer, return a list of markdown formated references to the sources used in following format: <example>[^1]: [name](link)</example>"]
            return self.agent_executor.run(query)
        except Exception as e:
            msg = f"Agent encounter an error.\n\n Error: {str(e)}"
            return msg