EVAL/core/agents/manager.py
ChungHwan Han 3ec0cc786c
Feature/better usability (#6)
* refactor: api, core

* feat: static uploader

* doc: update readme

* fix: mkdir static file

* doc: typo
2023-03-23 16:33:45 +09:00

68 lines
2.1 KiB
Python

from typing import Dict
from langchain.agents.tools import BaseTool
from langchain.agents.agent import Agent, AgentExecutor
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain.memory.chat_memory import BaseChatMemory
from core.tools.base import BaseToolSet
from core.tools.factory import ToolsFactory
from .builder import AgentBuilder
class AgentManager:
def __init__(
self,
agent: Agent,
global_tools: list[BaseTool],
toolsets: list[BaseToolSet] = [],
):
self.agent: Agent = agent
self.global_tools: list[BaseTool] = global_tools
self.toolsets: list[BaseToolSet] = toolsets
self.executors: Dict[str, AgentExecutor] = {}
def create_memory(self) -> BaseChatMemory:
return ConversationBufferMemory(memory_key="chat_history", return_messages=True)
def create_executor(self, session: str) -> AgentExecutor:
memory: BaseChatMemory = self.create_memory()
return AgentExecutor.from_agent_and_tools(
agent=self.agent,
tools=[
*self.global_tools,
*ToolsFactory.create_per_session_tools(
self.toolsets,
get_session=lambda: (session, self.executors[session]),
),
],
memory=memory,
verbose=True,
)
def remove_executor(self, session: str) -> None:
if session in self.executors:
del self.executors[session]
def get_or_create_executor(self, session: str) -> AgentExecutor:
if not (session in self.executors):
self.executors[session] = self.create_executor(session=session)
return self.executors[session]
@staticmethod
def create(toolsets: list[BaseToolSet]) -> "AgentManager":
builder = AgentBuilder(toolsets)
builder.build_llm()
builder.build_parser()
builder.build_global_tools()
agent = builder.get_agent()
global_tools = builder.get_global_tools()
return AgentManager(
agent=agent,
global_tools=global_tools,
toolsets=toolsets,
)