Merge pull request #41 from lambrou/hotfix/fix-continuity

Hotfix/fix continuity

Author: lambroulabs

semterm/agent/SemanticTerminalAgent.py

@@ -1,7 +1,7 @@
 import os
 
 from langchain.agents import AgentExecutor
-from langchain.memory import ConversationEntityMemory
+from langchain.memory import ConversationTokenBufferMemory
 from langchain_openai import ChatOpenAI
 
 from semterm.agent.TerminalAgentPrompt import PREFIX
@@ -40,9 +40,11 @@ def load_tools(self):
         return tools
 
     def initialize_memory(self):
-        return ConversationEntityMemory(
+        return ConversationTokenBufferMemory(
             llm=self.llm,
             return_messages=True,
+            input_key="input",
+            output_key="output",
             chat_history_key="chat_history",
         )
 
@@ -62,8 +64,9 @@ def initialize_executor(self):
             self.tools,
             memory=self.memory,
             max_iterations=self.max_iterations,
+            return_intermediate_steps=True,
             verbose=self.verbose,
         )
 
     def invoke(self, user_input):
-        return self.terminal_agent_executor.invoke(input=user_input)["output"]
+        return self.terminal_agent_executor.invoke({"input": user_input, "chat_history": []})["output"]

semterm/agent/TerminalAgentPrompt.py

@@ -1,5 +1,5 @@
 # flake8: noqa
-PREFIX = """You are a Semantic Terminal. Users will ask for you to perform tasks, expecting you to use the Terminal 
+PREFIX = """You are a Semantic Terminal, an intelligent AI that can only respond in thoughts followed by valid JSON. Users will ask for you to perform tasks, expecting you to use the Terminal 
 tool. Use it often and go above and beyond in completing the users request. Remember that you don't have to run all your commands in one go. You can run a command, look at the 
 output, and then run a new command based on that output. You can also use the Terminal tool to run multiple commands.
 If you need to install a program to use it, use the Human tool to get permission from the user and then install it.
@@ -39,11 +39,15 @@
 
 {format_instructions}
 
+CHAT HISTORY ----------
+
+{{{{chat_history}}}}
+
 USER'S INPUT ---------
 
 {{{{input}}}}
 
-Plan out what you will do to complete the task and then respond with an action.
+Plan out what you will do to complete the task and then respond with an action, all in one response. (Thoughts and actions need to be in one response)
 """
 
 TEMPLATE_TOOL_RESPONSE = """Observation: 

semterm/terminal/TerminalOutputParser.py

@@ -3,6 +3,8 @@
 from typing import Union
 from langchain.agents.conversational_chat.output_parser import ConvoOutputParser
 from langchain.schema import AgentAction, AgentFinish
+from langchain_core.agents import AgentFinish, AgentAction
+
 from semterm.agent.TerminalAgentPrompt import FORMAT_INSTRUCTIONS
 from semterm.langchain_extensions.schema import AgentMistake
 
@@ -11,14 +13,14 @@ class TerminalOutputParser(ConvoOutputParser, ABC):
     def get_format_instructions(self) -> str:
         return FORMAT_INSTRUCTIONS
 
-    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
+    def parse(self, text: str) -> AgentFinish | AgentAction | AgentMistake:
         text = text.strip().replace("\xa0", " ")
         start_positions = [i for i, c in enumerate(text) if c == "{"]
         end_positions = [i for i, c in enumerate(text) if c == "}"]
 
         for start in start_positions:
             for end in end_positions:
-                if start < end:  # ensure the end position is after the start
+                if start < end:
                     try:
                         cleaned_output = text[start : end + 1]
                         response = json.loads(cleaned_output)

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="semterm",
-    version="0.6.0",
+    version="0.6.1",
     description="The Semantic Terminal",
     long_description="The Semantic Terminal",
     author="Lambrou",

tests/agent/test_SemanticTerminalAgent.py

@@ -119,4 +119,4 @@ def test_invoke(self, mrkl_agent):
 
         mrkl_agent.terminal_agent_executor.invoke = MagicMock()
         mrkl_agent.invoke(user_input)
-        mrkl_agent.terminal_agent_executor.invoke.assert_called_with(input=user_input)
+        mrkl_agent.terminal_agent_executor.invoke.assert_called_with({"input": user_input, "chat_history": []})

tests/agent/test_TerminalAgent.py

@@ -7,6 +7,8 @@
     AIMessage,
     SystemMessage,
 )
+from langchain_core.messages import HumanMessage
+
 from semterm.agent.TerminalAgent import TerminalAgent
 from semterm.agent.TerminalAgentPrompt import PREFIX, SUFFIX
 
@@ -30,7 +32,9 @@ def test_create_prompt(self, terminal_agent, mock_tools):
         # Extract properties from the returned ChatPromptTemplate
         system_message_from_prompt = prompt.messages[0].format_messages()[0].content
         human_message_from_prompt = (
-            prompt.messages[2].format_messages(input="test input")[0].content
+            prompt.messages[2].format_messages(input="test input", chat_history=[HumanMessage(
+                content="test input"
+            )])[0].content
         )
 
         # Assert that the properties have the expected values