Merge pull request #142 from hoomano/kroussel/tag_manager_remove_tags

tag manager refacto

backend/app/models/assistant/chat_assistant.py

@@ -1,3 +1,4 @@
+from mojodex_core.tag_manager import TagManager
 from models.assistant.execution_manager import ExecutionManager
 from mojodex_core.llm_engine.mpt import MPT
 from mojodex_core.llm_engine.providers.openai_vision_llm import VisionMessagesData
@@ -6,37 +7,21 @@
 
 
 class ChatAssistant(ABC):
-    language_start_tag, language_end_tag = "<user_language>", "</user_language>"
-
-    @staticmethod
-    def remove_tags_from_text(text, start_tag, end_tag):
-        """
-        Remove tags from text
-        :param text: text
-        :param start_tag: start tag
-        :param end_tag: end tag
-
-        :return: text without tags
-        """
-        try:
-            return text.split(start_tag)[1].split(end_tag)[0].strip() if start_tag in text else ""
-        except Exception as e:
-            raise Exception(
-                f"ChatAssistant: remove_tags_from_text :: text: {text} - start_tag: {start_tag} - end_tag: {end_tag} - {e}")
 
+
     def __init__(self, mojo_message_token_stream_callback, draft_token_stream_callback,
                  tag_proper_nouns, user_messages_are_audio, db_session, temperature=0, max_tokens=4000):
         try:
 
             self.db_session = db_session
-            self.execution_manager = ExecutionManager()
             self.tag_proper_nouns = tag_proper_nouns
             self.user_messages_are_audio = user_messages_are_audio
             self.mojo_message_token_stream_callback = mojo_message_token_stream_callback
             self.draft_token_stream_callback = draft_token_stream_callback
             self.default_temperature = temperature
             self.default_max_tokens = max_tokens
             self.language = None
+            self.tag_manager = TagManager("user_language")
 
         except Exception as e:
             raise Exception(f"{self.__class__.__name__} __init__ :: {e}")
@@ -119,20 +104,18 @@ def _manage_response_language_tags(self, response):
         :param response: response
         """
         try:
-            if self.language_start_tag in response:
+            if self.tag_manager.start_tag in response:
                 try:
-                    self.language = self.remove_tags_from_text(response,
-                                                               self.language_start_tag,
-                                                               self.language_end_tag).lower()
+                    self.language = self.tag_manager.remove_tags_from_text(response).lower()
                 except Exception as e:
                     pass
         except Exception as e:
             raise Exception(f"__manage_response_language_tags:: {e}")
 
     def _manage_execution_tags(self, response):
         try:
-            if ExecutionManager.execution_start_tag in response:
-                return self.execution_manager.manage_execution_text(response)
+            if ExecutionManager.tag_manager.start_tag in response:
+                return ExecutionManager.manage_execution_text(response)
         except Exception as e:
             raise Exception(f"_manage_execution_tags :: {e}")
 

backend/app/models/assistant/execution_manager.py

@@ -1,18 +1,13 @@
-from mojodex_backend_logger import MojodexBackendLogger
-
-
+from mojodex_core.tag_manager import TagManager
 
 class ExecutionManager:
-    execution_start_tag, execution_end_tag = "<execution>", "</execution>"
 
-    def __init__(self):
-        self.logger = MojodexBackendLogger(f"{self.__class__.__name__}")
+    tag_manager: TagManager = TagManager("execution")
 
-    def manage_execution_text(self, execution_text):
+    @staticmethod
+    def manage_execution_text(execution_text):
         try:
-            from models.assistant.chat_assistant import ChatAssistant
-            mojo_text = ChatAssistant.remove_tags_from_text(execution_text, self.execution_start_tag,
-                                                                             self.execution_end_tag)
+            mojo_text = ExecutionManager.tag_manager.remove_tags_from_text(execution_text)
             return mojo_text
         except Exception as e:
-            raise Exception(f"{self.__class__.__name__} :: manage_execution_text :: {e}")
+            raise Exception(f"{ExecutionManager.__class__.__name__} :: manage_execution_text :: {e}")

backend/app/models/assistant/home_chat_assistant.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 
+from mojodex_core.tag_manager import TagManager
 from models.knowledge.knowledge_manager import KnowledgeManager
 from models.assistant.chat_assistant import ChatAssistant
 from app import placeholder_generator, server_socket
@@ -32,7 +33,7 @@ def __init__(self, mojo_message_token_stream_callback, draft_token_stream_callba
                 Session).get(session_id)
             self.use_message_placeholder = use_message_placeholder
             self.task_manager = InstructTaskManager(self.session.session_id, self.user.user_id)
-
+            self.user_message_tag_manager = TagManager("user_message")
 
         except Exception as e:
             raise Exception(f"{self.__class__.__name__} __init__ :: {e}")
@@ -75,8 +76,8 @@ def _handle_placeholder(self):
             raise Exception(f"_handle_placeholder :: {e}")
 
     def _get_message_placeholder(self):
-        return f"{self.user_message_start_tag}{placeholder_generator.mojo_message}{self.user_message_end_tag}"
-
+        return self.user_message_tag_manager.add_tags_to_text(placeholder_generator.mojo_message)
+  
     @property
     def _mpt(self):
         try:
@@ -157,8 +158,7 @@ def _token_callback(self, partial_text):
                         # do not stop the stream, this will take effect at next run
 
             if self.user_message_start_tag in partial_text:
-                text = ChatAssistant.remove_tags_from_text(partial_text, self.user_message_start_tag,
-                                                           self.user_message_end_tag)
+                text = self.user_message_tag_manager.remove_tags_from_text(partial_text)
                 self.mojo_message_token_stream_callback(text)
 
             # else, task specific tags
@@ -175,8 +175,7 @@ def _manage_response_tags(self, response):
                                                                                  task_name=self.instruct_task_execution.task_name_in_user_language,
                                                                                  user_task_execution_pk=self.instruct_task_execution.user_task_execution_pk)
             if self.user_message_start_tag in response:
-                text = ChatAssistant.remove_tags_from_text(response, self.user_message_start_tag,
-                                                           self.user_message_end_tag)
+                text = self.user_message_tag_manager.remove_tags_from_text(response)
                 return {"text": text, 'text_with_tags': response}
             return self.task_manager.manage_response_task_tags(response)
         except Exception as e:

backend/app/models/assistant/session_controller.py

@@ -1,10 +1,9 @@
 import os
 from app import server_socket, time_manager, socketio_message_sender, main_logger
-from models.produced_text_managers.task_produced_text_manager import TaskProducedTextManager
+from mojodex_core.tag_manager import TagManager
+from mojodex_core.produced_text_managers.task_produced_text_manager import TaskProducedTextManager
 from models.assistant.instruct_task_assistant import InstructTaskAssistant
 from models.assistant.home_chat_assistant import HomeChatAssistant
-from models.assistant.chat_assistant import ChatAssistant
-
 from models.assistant.workflow_assistant import WorkflowAssistant
 from mojodex_core.entities.instruct_user_task_execution import InstructTaskExecution
 from mojodex_core.entities.message import Message
@@ -84,16 +83,14 @@ def _mojo_message_stream_callback(self, partial_text):
     @token_stream_callback('draft_token')
     def _produced_text_stream_callback(self, partial_text):
         try:
-            title = ChatAssistant.remove_tags_from_text(partial_text.strip(),
-                                                        TaskProducedTextManager.title_start_tag,
-                                                        TaskProducedTextManager.title_end_tag)
-            production = ChatAssistant.remove_tags_from_text(partial_text.strip(),
-                                                             TaskProducedTextManager.draft_start_tag,
-                                                             TaskProducedTextManager.draft_end_tag)
+            title_tag_manager = TagManager("title")
+            draft_tag_manager = TagManager("draft")
+            title = title_tag_manager.remove_tags_from_text(partial_text.strip())
+            production = draft_tag_manager.remove_tags_from_text(partial_text.strip())
             return {"produced_text_title": title,
                     "produced_text": production,
                     "session_id": self.session.session_id,
-                    "text": TaskProducedTextManager.remove_tags(partial_text)}
+                    "text": TaskProducedTextManager.get_produced_text_without_tags(partial_text)}
         except Exception as e:
             raise Exception(f"_produced_text_stream_callback :: {e}")
 

backend/app/models/assistant/workflow_assistant.py

@@ -1,5 +1,5 @@
 import json
-from models.produced_text_managers.task_produced_text_manager import TaskProducedTextManager
+from mojodex_core.produced_text_managers.task_produced_text_manager import TaskProducedTextManager
 from models.knowledge.knowledge_manager import KnowledgeManager
 from models.assistant.chat_assistant import ChatAssistant
 from app import placeholder_generator
@@ -78,10 +78,10 @@ def _mpt(self):
                        user_company_knowledge=self.workflow_execution.user.company_description,
                        infos_to_extract=self.workflow_execution.task.infos_to_extract,
                        workflow=self.workflow_execution.task,
-                       title_start_tag=TaskProducedTextManager.title_start_tag,
-                       title_end_tag=TaskProducedTextManager.title_end_tag,
-                       draft_start_tag=TaskProducedTextManager.draft_start_tag,
-                       draft_end_tag=TaskProducedTextManager.draft_end_tag,
+                       title_start_tag=TaskProducedTextManager.title_tag_manager.start_tag,
+                       title_end_tag=TaskProducedTextManager.title_tag_manager.end_tag,
+                       draft_start_tag=TaskProducedTextManager.draft_tag_manager.start_tag,
+                       draft_end_tag=TaskProducedTextManager.draft_tag_manager.end_tag,
                        user_workflow_inputs=self.workflow_execution.json_input_values,
                        produced_text_done=self.workflow_execution.produced_text_done,
                        audio_message=self.user_messages_are_audio,

backend/app/models/tasks/instruct_tasks/instruct_task_manager.py

@@ -1,34 +1,33 @@
 from models.tasks.task_executor import TaskExecutor
 
-from models.produced_text_managers.task_produced_text_manager import TaskProducedTextManager
+from mojodex_core.produced_text_managers.task_produced_text_manager import TaskProducedTextManager
 from models.assistant.execution_manager import ExecutionManager
 from app import placeholder_generator
-from models.assistant.chat_assistant import ChatAssistant
 
-from models.tasks.tag_manager import TagManager
+from mojodex_core.tag_manager import TagManager
 
 
 class InstructTaskManager:
 
     def __init__(self, session_id, user_id):
-        self.task_input_manager = TagManager("ask_user_primary_info", "Hello world!")
+        self.task_input_manager = TagManager("ask_user_primary_info")
         self.task_executor = TaskExecutor(session_id, user_id)
 
     @property
     def task_execution_placeholder(self):
-        return f"{ExecutionManager.execution_start_tag}" \
-               f"{TaskProducedTextManager.title_start_tag}{placeholder_generator.mojo_draft_title}{TaskProducedTextManager.title_end_tag}" \
-               f"{TaskProducedTextManager.draft_start_tag}{placeholder_generator.mojo_draft_body}{TaskProducedTextManager.draft_end_tag}" \
-               f"{ExecutionManager.execution_end_tag}"
+        return ExecutionManager.tag_manager.add_tags_to_text(
+               f"{TaskProducedTextManager.title_tag_manager.add_tags_to_text(placeholder_generator.mojo_draft_title)}"
+               f"{TaskProducedTextManager.draft_tag_manager.add_tags_to_text(placeholder_generator.mojo_draft_body)}")
 
     @property
     def task_message_placeholder(self):
-        return self.task_input_manager.placeholder
+        return self.task_input_manager.add_tags_to_text("Can you please provide more information?")
 
     def manage_response_task_tags(self, response):
         try:
             if self.task_input_manager.start_tag in response:
-                return self.task_input_manager.manage_text(response)
+                text_without_tags = self.task_input_manager.remove_tags_from_text(response)
+                return {"text": text_without_tags, "text_with_tags": response}
             return {"text": response}
         except Exception as e:
             raise Exception(f"{self.__class__.__name__} :: manage_response_task_tags :: {e}")
@@ -37,17 +36,13 @@ def manage_task_stream(self, partial_text, mojo_message_token_stream_callback, d
         try:
             text = None
             if self.task_input_manager.start_tag in partial_text:
-                text = ChatAssistant.remove_tags_from_text(partial_text,
-                                                           self.task_input_manager.start_tag,
-                                                           self.task_input_manager.end_tag)
+                text = self.task_input_manager.remove_tags_from_text(partial_text)
             if text and mojo_message_token_stream_callback:
                 mojo_message_token_stream_callback(text)
 
-            elif ExecutionManager.execution_start_tag in partial_text:
+            elif ExecutionManager.tag_manager.start_tag in partial_text:
                 # take the text between <execution> and </execution>
-                text = ChatAssistant.remove_tags_from_text(partial_text,
-                                                           ExecutionManager.execution_start_tag,
-                                                           ExecutionManager.execution_end_tag)
+                text = ExecutionManager.tag_manager.remove_tags_from_text(partial_text)
                 draft_token_stream_callback(text)
         except Exception as e:
             raise Exception(f"{self.__class__.__name__} :: manage_task_stream :: {e}")