EVAL/core/agents/builder.py

72 lines
2.3 KiB
Python
Raw Normal View History

2023-03-18 12:26:19 +00:00
from langchain.chat_models.base import BaseChatModel
2023-03-20 08:27:20 +00:00
from langchain.output_parsers.base import BaseOutputParser
2023-03-18 06:05:02 +00:00
from core.prompts.input import EVAL_PREFIX, EVAL_SUFFIX
from core.tools.base import BaseToolSet
from core.tools.factory import ToolsFactory
2023-04-03 07:43:34 +00:00
from env import settings
from .chat_agent import ConversationalChatAgent
2023-04-03 07:43:34 +00:00
from .llm import ChatOpenAI
from .parser import EvalOutputParser
2023-03-18 06:05:02 +00:00
2023-03-18 12:42:14 +00:00
class AgentBuilder:
def __init__(self, toolsets: list[BaseToolSet] = []):
2023-03-18 12:26:19 +00:00
self.llm: BaseChatModel = None
2023-03-20 08:27:20 +00:00
self.parser: BaseOutputParser = None
self.global_tools: list = None
self.toolsets = toolsets
2023-03-18 12:26:19 +00:00
2023-03-18 12:42:14 +00:00
def build_llm(self):
2023-03-18 12:26:19 +00:00
self.llm = ChatOpenAI(temperature=0)
2023-03-20 08:27:20 +00:00
def build_parser(self):
self.parser = EvalOutputParser()
def build_global_tools(self):
2023-03-18 12:26:19 +00:00
if self.llm is None:
raise ValueError("LLM must be initialized before tools")
toolnames = ["wikipedia"]
2023-03-18 12:26:19 +00:00
if settings["SERPAPI_API_KEY"]:
toolnames.append("serpapi")
if settings["BING_SEARCH_URL"] and settings["BING_SUBSCRIPTION_KEY"]:
toolnames.append("bing-search")
self.global_tools = [
*ToolsFactory.create_global_tools_from_names(toolnames, llm=self.llm),
*ToolsFactory.create_global_tools(self.toolsets),
2023-03-18 12:26:19 +00:00
]
def get_global_tools(self):
if self.global_tools is None:
raise ValueError("Global tools are not initialized yet")
2023-03-18 12:42:14 +00:00
return self.global_tools
2023-03-20 08:27:20 +00:00
def get_agent(self):
2023-03-18 12:42:14 +00:00
if self.llm is None:
raise ValueError("LLM must be initialized before agent")
2023-03-20 08:27:20 +00:00
if self.parser is None:
raise ValueError("Parser must be initialized before agent")
if self.global_tools is None:
raise ValueError("Global tools must be initialized before agent")
2023-03-20 08:27:20 +00:00
return ConversationalChatAgent.from_llm_and_tools(
llm=self.llm,
tools=[
*self.global_tools,
*ToolsFactory.create_per_session_tools(
self.toolsets
), # for names and descriptions
],
system_message=EVAL_PREFIX.format(bot_name=settings["BOT_NAME"]),
human_message=EVAL_SUFFIX.format(bot_name=settings["BOT_NAME"]),
output_parser=self.parser,
max_iterations=30,
2023-03-18 12:26:19 +00:00
)