baLLseM / model /llm_agent.py
hqms's picture
initial commit
8578816
raw
history blame contribute delete
863 Bytes
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()