Add support of reply and do not get borred if prev message is borred. Manage memory.

fjerkroa_bot/ai_responder.py

@@ -150,13 +150,19 @@ class AIResponder(AIResponderBase):
     def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None:
         super().__init__(config, channel)
         self.history: List[Dict[str, Any]] = []
+        self.memory: str = 'I am an assistant.'
         self.rate_limit_backoff = exponential_backoff()
         self.history_file: Optional[Path] = None
+        self.memory_file: Optional[Path] = None
         if 'history-directory' in self.config:
             self.history_file = Path(self.config['history-directory']).expanduser() / f'{self.channel}.dat'
             if self.history_file.exists():
                 with open(self.history_file, 'rb') as fd:
                     self.history = pickle.load(fd)
+            self.memory_file = Path(self.config['history-directory']).expanduser() / f'{self.channel}.memory'
+            if self.memory_file.exists():
+                with open(self.memory_file, 'rb') as fd:
+                    self.memory = pickle.load(fd)
 
     def message(self, message: AIMessage, limit: Optional[int] = None) -> List[Dict[str, Any]]:
         messages = []
@@ -168,6 +174,7 @@ class AIResponder(AIResponderBase):
             with open(news_feed) as fd:
                 news_feed = fd.read().strip()
             system = system.replace('{news}', news_feed)
+        system = system.replace('{memory}', self.memory)
         messages.append({"role": "system", "content": system})
         if limit is not None:
             while len(self.history) > limit:
@@ -240,6 +247,9 @@ class AIResponder(AIResponderBase):
     async def fix(self, answer: str) -> str:
         raise NotImplementedError()
 
+    async def memory_rewrite(self, memory: str, user: str, question: str, answer: str) -> str:
+        raise NotImplementedError()
+
     async def translate(self, text: str, language: str = "english") -> str:
         raise NotImplementedError()
 
@@ -268,6 +278,19 @@ class AIResponder(AIResponderBase):
             with open(self.history_file, 'wb') as fd:
                 pickle.dump(self.history, fd)
 
+    def update_memory(self, memory) -> None:
+        if self.memory_file is not None:
+            with open(self.memory_file, 'wb') as fd:
+                pickle.dump(self.memory, fd)
+
+    async def handle_picture(self, response: Dict) -> bool:
+        if not isinstance(response.get("picture"), (type(None), str)):
+            logging.warning(f"picture key is wrong in response: {pp(response)}")
+            return False
+        if response.get("picture") is not None:
+            response["picture"] = await self.translate(response["picture"])
+        return True
+
     async def send(self, message: AIMessage) -> AIResponse:
         # Get the history limit from the configuration
         limit = self.config["history-limit"]
@@ -305,15 +328,10 @@ class AIResponder(AIResponderBase):
                     retries -= 1
                     continue
 
-            # Check if the response has the correct picture format
-            if not isinstance(response.get("picture"), (type(None), str)):
-                logging.warning(f"picture key is wrong in response: {pp(response)}")
+            if not await self.handle_picture(response):
                 retries -= 1
                 continue
 
-            if response.get("picture") is not None:
-                response["picture"] = await self.translate(response["picture"])
-
             # Post-process the message and update the answer's content
             answer_message = await self.post_process(message, response)
             answer['content'] = str(answer_message)
@@ -322,6 +340,11 @@ class AIResponder(AIResponderBase):
             self.update_history(messages[-1], answer, limit, message.historise_question)
             logging.info(f"got this answer:\n{str(answer_message)}")
 
+            # Update memory
+            if answer_message.answer is not None:
+                self.memory = await self.memory_rewrite(self.memory, message.user, message.message, answer_message.answer)
+                self.update_memory(self.memory)
+
             # Return the updated answer message
             return answer_message