|
from model.tools import llm |
|
from model.tools.wikipedia import wikipedia_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], |
|
llm=llm, |
|
agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, |
|
verbose=False, |
|
handle_parsing_errors=True, |
|
max_execution_time=3600, |
|
max_iterations=15 |
|
) |
|
|
|
def prompt(self, text): |
|
result = self.agent.invoke(text) |
|
|
|
return result |
|
|
|
llm_agent = LLMAgent() |