refactor: exit conversation

agents/manager.py

@@ -5,16 +5,23 @@ from langchain.chains.conversation.memory import ConversationBufferMemory
 from langchain.memory.chat_memory import BaseChatMemory
 
 from tools.base import BaseToolSet
+from tools.cpu import ExitConversation
 
 from .builder import AgentBuilder
 
 
 class AgentManager:
-    def __init__(self, agent: Agent, tools: list[BaseToolSet]):
-        self.agent: Agent = agent
-        self.tools: list[BaseToolSet] = tools
+    def __init__(self):
+        self.agent: Agent = None
+        self.tools: list[BaseToolSet] = []
         self.executors: Dict[str, AgentExecutor] = {}
 
+    def set_agent(self, agent: Agent) -> None:
+        self.agent = agent
+
+    def set_tools(self, tools: list[BaseToolSet]) -> None:
+        self.tools = tools
+
     def create_memory(self) -> BaseChatMemory:
         return ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
@@ -37,11 +44,14 @@ class AgentManager:
 
     @staticmethod
     def create(toolsets: list[BaseToolSet]) -> "AgentManager":
+        manager = AgentManager()
+
         builder = AgentBuilder()
         builder.build_llm()
         builder.build_parser()
-        builder.build_tools(toolsets)
-        agent = builder.get_agent()
-        tools = builder.get_tools()
+        builder.build_tools([*toolsets, ExitConversation(manager.executors)])
+
+        manager.set_agent(builder.get_agent())
+        manager.set_tools(builder.get_tools())
 
-        return AgentManager(agent, tools)
+        return manager

main.py

@@ -29,13 +29,11 @@ from handlers.dataframe import CsvToDataframe
 
 app = FastAPI()
 
-agent_manager: AgentManager = None
 
 toolsets: List[BaseToolSet] = [
     Terminal(),
     CodeEditor(),
     RequestsGet(),
-    ExitConversation(agent_manager),
     Text2Image("cuda"),
     ImageEditing("cuda"),
     InstructPix2Pix("cuda"),
@@ -78,12 +76,11 @@ async def command(request: Request) -> Response:
 
     print("=============== Running =============")
     print("Inputs:", query, files)
-    # TODO - add state to memory (use key)
-
     executor = agent_manager.get_or_create_executor(key)
 
     print("======>Previous memory:\n %s" % executor.memory)
 
+    # TODO: exit conversation
     promptedQuery = "\n".join([file_handler.handle(file) for file in files])
     promptedQuery += query
     print("======>Prompted Text:\n %s" % promptedQuery)

tools/cpu.py

@@ -1,16 +1,16 @@
 from env import settings
 
+from typing import Dict
 import requests
 
 from llama_index.readers.database import DatabaseReader
 from llama_index import GPTSimpleVectorIndex
 
 from bs4 import BeautifulSoup
-from langchain.memory.chat_memory import BaseChatMemory
+from langchain.agents.agent import AgentExecutor
 
 import subprocess
 
-from agents.manager import AgentManager
 from .base import tool, BaseToolSet
 
 
@@ -147,20 +147,20 @@ class WineDB(BaseToolSet):
 
 
 class ExitConversation(BaseToolSet):
-    def __init__(self, agent_manager: AgentManager):
-        self.agent_manager = agent_manager
+    def __init__(self, executors: Dict[str, AgentExecutor]):
+        self.executors = executors
 
     @tool(
-        name="exit_conversation",
+        name="Exit Conversation",
         description="A tool to exit the conversation. "
-        "Use this when you want to end the conversation. "
+        "Use this when you want to exit the conversation. "
         "Input should be a user's key."
         "The output will be a message that the conversation is over.",
     )
     def exit(self, key: str) -> str:
         """Run the tool."""
-        self.agent_manager.remove_executor(key)
+        self.executors.pop(key)
 
         print(f"\nProcessed ExitConversation.")
 
-        return f"End conversation."
+        return f"Exit conversation."