from model.tools import llm
from model.tools.wikipedia import wikipedia_tool
# from model.tools.sql_tools import sql_tool
from model.tools.predictor import word_problem_tool
from langchain.agents.agent_types import AgentType
from langchain.agents import initialize_agent

class LLMAgent(object):
    def __init__(self) -> None:
        self.agent = initialize_agent(
            tools=[wikipedia_tool, word_problem_tool], #sql_tool],
            llm=llm,
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=False,
            handle_parsing_errors=True,
            max_execution_time=3600,  # Set the maximum execution time (in seconds)
            max_iterations=15  # Set the maximum number of iterations
        )

    def prompt(self, text):
        result = self.agent.invoke(text)

        return result
    
llm_agent = LLMAgent()