Memory doc and UI Refinement.

configs/alice.yaml

@@ -5,7 +5,7 @@ version: 0.0.1
 description: example
 prompt_template: !prompt VanillaPrompt
 llm:
-  model_name: guanaco-7b
+  model_name: vicuna-7b
   param:
     temperature: 0.6
     top_p: 0.8

configs/memory.yaml

@@ -1,13 +1,13 @@
 # Agent Config
-name: main(memory)
+name: main
 type: openai_memory
 version: 0.0.1
 description: main agent leveraging OpenAI function call API.
 prompt_template: !prompt VanillaPrompt
 memory:
   memory_type: chroma   # chroma or pinecone
-  threshold_1: 1        # first-level memory
-  threshold_2: 1        # second-level memory
+  conversation_threshold: 1        # first-level memory
+  reasoning_threshold: 1          # second-level memory
   params:
     index: main
     top_k: 2
@@ -24,6 +24,6 @@ target_tasks:
 plugins:
   - name: google_search
   - name: web_page
-#  - !include mathria.yaml
+  # - !include mathria.yaml
 
 

gentopia/agent/openai_memory/agent.py

@@ -129,7 +129,7 @@ def run(self, instruction: str, output: Optional[BaseOutput] = None) -> AgentOut
         if output is None:
             output = BaseOutput()
         self.memory.clear_memory_II()
-        message_scratchpad = self.__add_system_prompt(self.memory.lastest_context(instruction))
+        message_scratchpad = self.__add_system_prompt(self.memory.lastest_context(instruction, output))
         total_cost = 0
         total_token = 0
 
@@ -145,10 +145,10 @@ def run(self, instruction: str, output: Optional[BaseOutput] = None) -> AgentOut
         if response.state == "success":
             output.done(self.name)
             output.panel_print(response.content)
-            self.memory.save_memory_I({"role": "user", "content": instruction}, response.message_scratchpad[-1], self.llm)
+            self.memory.save_memory_I({"role": "user", "content": instruction}, response.message_scratchpad[-1], output)
 
             if len(response.message_scratchpad) != len(message_scratchpad) + 1: # normal case
-                self.memory.save_memory_II(response.message_scratchpad[-3], response.message_scratchpad[-2],  self.llm)       
+                self.memory.save_memory_II(response.message_scratchpad[-3], response.message_scratchpad[-2], output, self.llm)       
 
             total_cost += calculate_cost(self.llm.model_name, response.prompt_token,
                                          response.completion_token) + response.plugin_cost
@@ -174,7 +174,7 @@ def stream(self, instruction: Optional[str] = None, output: Optional[BaseOutput]
         output.thinking(self.name)
         if is_start:
             self.memory.clear_memory_II()
-        message_scratchpad = self.__add_system_prompt(self.memory.lastest_context(instruction))
+        message_scratchpad = self.__add_system_prompt(self.memory.lastest_context(instruction, output))
         output.debug(message_scratchpad)
         assert len(message_scratchpad) > 1
 
@@ -221,10 +221,10 @@ def stream(self, instruction: Optional[str] = None, output: Optional[BaseOutput]
             self.memory.save_memory_II(dict(role='assistant', content=None, function_call={i: str(j) for i, j in result.items()}),
                                             {"role": "function",
                                             "name": function_name,
-                                            "content": function_response}, self.llm)
+                                            "content": function_response}, output, self.llm)
             self.stream(instruction, output=output, is_start=False)
         else:
-            self.memory.save_memory_I({"role": "user", "content": instruction}, {"role": _role, "content": result}, self.llm)
+            self.memory.save_memory_I({"role": "user", "content": instruction}, {"role": _role, "content": result}, output)
         # else:
         #     self.message_scratchpad.append({"role": "user", "content": "Summarize what you have done and continue if you have not finished."})
         #     self.stream(output=output)

gentopia/agent/react/agent.py

@@ -76,7 +76,6 @@ def _construct_scratchpad(
     ) -> str:
         """Construct the scratchpad that lets the agent continue its thought process."""
         thoughts = ""
-        print(intermediate_steps)
         for action, observation in intermediate_steps:
             thoughts += action.log
             thoughts += f"\nObservation: {observation}\nThought:"

gentopia/assembler/agent_assembler.py

@@ -96,8 +96,7 @@ def _parse_memory(self, obj) -> MemoryWrapper:
         if obj == []:
             return None
         memory_type = obj["memory_type"] # memory_type: ["pinecone"]
-        return create_memory(memory_type, obj['threshold_1'], obj['threshold_2'], **obj["params"]) # params of memory.
-        # Different memories may have different params
+        return create_memory(memory_type, obj['conversation_threshold'], obj['reasoning_threshold'], **obj["params"]) # params of memory. Different memories may have different params
 
 
     def _get_llm(self, obj) -> Union[BaseLLM, Dict[str, BaseLLM]]:

gentopia/memory/api.py

@@ -5,6 +5,7 @@
 from gentopia.memory.embeddings import OpenAIEmbeddings
 from gentopia.llm.base_llm import BaseLLM
 from gentopia import PromptTemplate
+from gentopia.output.base_output import BaseOutput
 import pydantic
 import os
 import queue
@@ -16,12 +17,12 @@ class Config:
     input_variables=["rank", "input", "output"],
     template=
 """
-You are a helpful assistant who are expected to summarize some sentences.
-Another AI assistant are interacting with user and mutiple tools. 
-Here is part of their conversations, you need to summarize them and provide a brief summary, which will help other assisant to recall their thoughts and actions.
-Note that you need to use some words like \"In the fourth step\" according to the rank in to start your summary. For example, you need to use \"In the fifth step\" if it is step 5, or use \"First\" if it is step 1 or \"Second\" in step 2.
+You are a helpful AI assistant to summarize some sentences.
+Another assistant has interacted with the user for multiple rounds.
+Here is part of their conversation for a certain step. You need to provide a brief summary that helps the other assistant recall previous thoughts and actions.
+Note that you need to start with "In the xxx step" depending on the step number. For example, "In the second step" if the step number is 2.
 
-In step {rank}, the part of conversation is:
+In step {rank}, the part of the conversation is:
 Input: {input}
 Output: {output}
 Your summary:
@@ -32,7 +33,7 @@ class Config:
     input_variables=['summary'],
     template=
 """
-The following summaries of context may help you to recall the former conversation.
+The following summaries of context may help you recall the former conversation.
 {summary}.
 End of the summaries.
 """
@@ -42,7 +43,7 @@ class Config:
     input_variables=["summary"],
     template=
 """
-The following summaries of context may help you to recall your memory, which assists you to take your next step.
+The following summaries of context may help you recall your prior steps, which assist you in taking your next step.
 {summary}
 End of the summaries.
 """
@@ -52,51 +53,111 @@ class Config:
     input_variables=["related_history"],
     template=
 """
-Here are some related conversations which may help you to answer the question:
+Here are some related conversations that may help you answer the next question:
 {related_history}
 End of the related history.
 """
 )
 
 @pydantic.dataclasses.dataclass(config=Config)
 class MemoryWrapper:
+    """
+    Wrapper class for memory management.
+    """
+
     memory: BaseMemory
-    threshold_I: int
-    threshold_II: int
+    conversation_threshold: int
+    reasoning_threshold: int
 
-    def __init__(self, memory: VectorStoreRetrieverMemory, threshold1: int, threshold2: int):
+    def __init__(self, memory: VectorStoreRetrieverMemory, conversation_threshold: int, reasoning_threshold: int):
+        """
+        Initialize the MemoryWrapper.
+
+        :param memory: The vector store retriever memory.
+        :type memory: VectorStoreRetrieverMemory
+        :param conversation_threshold: The conversation threshold.
+        :type conversation_threshold: int
+        :param reasoning_threshold: The reasoning threshold.
+        :type reasoning_threshold: int
+        """
         self.memory = memory
-        self.threshold_I = threshold1
-        self.threshold_II = threshold2
-        assert self.threshold_I >= 0
-        assert self.threshold_II >= 0
+        self.conversation_threshold = conversation_threshold
+        self.reasoning_threshold = reasoning_threshold
+        assert self.conversation_threshold >= 0
+        assert self.reasoning_threshold >= 0
         self.history_queue_I = queue.Queue()
         self.history_queue_II = queue.Queue()
-        self.summary_I =  ""     # memory I  - level
-        self.summary_II = ""     # memory II - level
+        self.summary_I = ""     # memory I - level
+        self.summary_II = ""    # memory II - level
         self.rank_I = 0
         self.rank_II = 0
 
     def __save_to_memory(self, io_obj):
+        """
+        Save the input-output pair to memory.
+
+        :param io_obj: The input-output pair.
+        :type io_obj: Any
+        """
         self.memory.save_context(io_obj[0], io_obj[1]) # (input, output)
 
-    def save_memory_I(self, input, output, llm: BaseLLM):
+    def save_memory_I(self, query, response, output: BaseOutput):
+        """
+        Save the conversation to memory (level I).
+
+        :param query: The query.
+        :type query: Any
+        :param response: The response.
+        :type response: Any
+        :param output: The output object.
+        :type output: BaseOutput
+        """
+
+        output.update_status("Conversation Memorizing...")
         self.rank_I += 1
-        self.history_queue_I.put((input, output, self.rank_I))
-        while self.history_queue_I.qsize() > self.threshold_I:
+        self.history_queue_I.put((query, response, self.rank_I))
+        while self.history_queue_I.qsize() > self.conversation_threshold:
             top_context = self.history_queue_I.get()
             self.__save_to_memory(top_context)
             # self.summary_I += llm.completion(prompt=SummaryPrompt.format(rank=top_context[2], input=top_context[0], output=top_context[1])).content + "\n"
+        output.done()
 
-    def save_memory_II(self, input, output, llm: BaseLLM):
+    def save_memory_II(self, query, response, output: BaseOutput, llm: BaseLLM):
+        """
+        Save the conversation to memory (level II).
+
+        :param query: The query.
+        :type query: Any
+        :param response: The response.
+        :type response: Any
+        :param output: The output object.
+        :type output: BaseOutput
+        :param llm: The BaseLLM object.
+        :type llm: BaseLLM
+        """
+        output.update_status("Reasoning Memorizing...")
         self.rank_II += 1
-        self.history_queue_II.put((input, output, self.rank_II))
-        while self.history_queue_II.qsize() > self.threshold_II:
+        self.history_queue_II.put((query, response, self.rank_II))
+        while self.history_queue_II.qsize() > self.reasoning_threshold:
             top_context = self.history_queue_II.get()
             self.__save_to_memory(top_context)
+            output.done()
+            output.update_status("Summarizing...")
             self.summary_II += llm.completion(prompt=SummaryPrompt.format(rank=top_context[2], input=top_context[0], output=top_context[1])).content + "\n"
+        output.done()
+
+    def lastest_context(self, instruction, output: BaseOutput):
+        """
+        Get the latest context history.
 
-    def lastest_context(self, instruction):
+        :param instruction: The instruction.
+        :type instruction: Any
+        :param output: The output object.
+        :type output: BaseOutput
+
+        :return: The context history.
+        :rtype: List[Dict[str, Any]]
+        """
         context_history = []
         # TODO this context_history can only be used in openai agent. This function should be more universal
         if self.summary_I != "":
@@ -105,28 +166,59 @@ def lastest_context(self, instruction):
             context_history.append(i[0])
             context_history.append(i[1])
         related_history = self.load_history(instruction)
+
         if related_history != "":
-            instruction += "\n" + RelatedContextPrompt.format(related_history=related_history)
-        if self.summary_II != "":
-            instruction += "\n" + RecallPrompt.format(summary = self.summary_II)
+            output.panel_print(related_history, f"[green] Related Conversation Memory: ")
+            context_history.append({"role": "user", "content": RelatedContextPrompt.format(related_history=related_history)})
+
         context_history.append({"role": "user", "content": instruction})
+
+        if self.summary_II != "":
+            output.panel_print(self.summary_II, f"[green] Summary of Prior Steps: ")
+            context_history.append({"role": "user", "content": RecallPrompt.format(summary = self.summary_II)})
+
         for i in list(self.history_queue_II.queue):
             context_history.append(i[0])
             context_history.append(i[1])
         return context_history
 
     def clear_memory_II(self):
+        """
+        Clear memory (level II).
+        """
         self.summary_II = ""
         self.history_queue_II = queue.Queue()
         self.rank_II = 0
-
+ 
 
     def load_history(self, input):
+        """
+        Load history from memory.
+
+        :param input: The input.
+        :type input: Any
+
+        :return: The loaded history.
+        :rtype: str
+        """
         return self.memory.load_memory_variables({"query": input})['history']
 
 
+def create_memory(memory_type, conversation_threshold, reasoning_threshold, **kwargs) -> MemoryWrapper:
+    """
+    Create a memory object.
+
+    :param memory_type: The type of memory.
+    :type memory_type: str
+    :param conversation_threshold: The conversation threshold.
+    :type conversation_threshold: int
+    :param reasoning_threshold: The reasoning threshold.
+    :type reasoning_threshold: int
+    :param **kwargs: Additional keyword arguments.
 
-def create_memory(memory_type, threshold1, threshold2, **kwargs) -> MemoryWrapper:
+    :return: The created MemoryWrapper object.
+    :rtype: MemoryWrapper
+    """
     # choose desirable memory you need!
     memory: BaseMemory = None
     if memory_type == "pinecone":
@@ -144,4 +236,4 @@ def create_memory(memory_type, threshold1, threshold2, **kwargs) -> MemoryWrappe
         memory = VectorStoreRetrieverMemory(retriever=retriever)
     else:
         raise ValueError(f"Memory {memory_type} is not supported currently.")   
-    return MemoryWrapper(memory, threshold1, threshold2)
+    return MemoryWrapper(memory, conversation_threshold, reasoning_threshold)

gentopia/memory/base_memory.py

@@ -1,5 +1,3 @@
-"""Class for a VectorStore-backed memory object."""
-
 from typing import Any, Dict, List
 from abc import ABC, abstractmethod
 from gentopia.memory.serializable import Serializable
@@ -19,14 +17,28 @@ def memory_variables(self) -> List[str]:
 
     @abstractmethod
     def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
-        """Return key-value pairs given the text input to the chain.
+        """
+        Return key-value pairs given the text input to the chain.
+
+        If None, return all memories.
+
+        :param inputs: The text inputs to the chain.
+        :type inputs: Dict[str, Any]
 
-        If None, return all memories
+        :return: The key-value pairs representing the memories.
+        :rtype: Dict[str, Any]
         """
 
     @abstractmethod
     def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
-        """Save the context of this model run to memory."""
+        """
+        Save the context of this model run to memory.
+
+        :param inputs: The input values.
+        :type inputs: Dict[str, Any]
+        :param outputs: The output values.
+        :type outputs: Dict[str, str]
+        """
 
     @abstractmethod
     def clear(self) -> None: