Wrap history manager based on settings

Signed-off-by: Ygal Blum <[email protected]>

cloud/helm/knowledgebase-slackbot/Chart.yaml

@@ -10,7 +10,7 @@ version: 0.1.0
 # incremented each time you make changes to the application. Versions are not expected to
 # follow Semantic Versioning. They should reflect the version the application is using.
 # It is recommended to use it with quotes.
-appVersion: "0.0.23"
+appVersion: "0.0.24"
 
 keywords:
 - slackbot

cloud/helm/knowledgebase-slackbot/templates/_knowledgebase-config.tpl

@@ -21,6 +21,9 @@ ollama:
   llm_model: {{ .Values.ollamaChatModel | quote }}
   embedding_model: {{ .Values.ollamaEmbeddingModel | quote }}
 
+history:
+  mode: redis
+
 redis:
   host: "{{ include "knowledgebase-slackbot.redis-host" . }}"
   password: $REDIS_PASSWORD|""

knowledge_base_gpt/apps/slackbot/slack_bot.py

@@ -6,7 +6,7 @@
 
 from knowledge_base_gpt.libs.settings.settings import Settings
 from knowledge_base_gpt.libs.gpt.private_chat import PrivateChat
-from knowledge_base_gpt.libs.history.redis import HistoryRedis
+from knowledge_base_gpt.libs.history.history import History
 
 
 class KnowledgeBaseSlackBotException(Exception):
@@ -18,7 +18,7 @@ class KnowledgeBaseSlackBot():  # pylint:disable=R0903
     """ Slackbot application backend """
 
     @inject
-    def __init__(self, settings: Settings, private_chat: PrivateChat, history: HistoryRedis) -> None:
+    def __init__(self, settings: Settings, private_chat: PrivateChat, history: History) -> None:
         self._private_chat = private_chat
         self._history = history
         self._handler = SocketModeHandler(App(token=settings.slackbot.bot_token), settings.slackbot.app_token)
@@ -52,11 +52,11 @@ def _got_message(self, message, say):
             text="On it. Be back with your answer soon"
         )
         answer = self._private_chat.answer_query(
-            self._history.get_messages(message['user']),
+            self._history.history.get_messages(message['user']),
             message['text'],
             chat_identifier=message['user']
         )
-        self._history.add_to_history(message['user'], answer)
+        self._history.history.add_to_history(message['user'], answer)
         say(answer['answer'])
 
     def _is_direct_message_channel(self, command):
@@ -74,7 +74,7 @@ def _reset_conversation(self, ack, say, command):  # pylint:disable=unused-argum
         if not self._is_direct_message_channel(command):
             return
 
-        self._history.reset(command['user_id'])
+        self._history.history.reset(command['user_id'])
 
         self._handler.app.client.chat_postEphemeral(
             channel=command['channel_id'],
@@ -96,7 +96,7 @@ def _forward_question(self, ack, say, command):  # pylint:disable=unused-argumen
         if not self._is_direct_message_channel(command):
             return
 
-        messages = self._history.get_messages(command['user_id'])
+        messages = self._history.history.get_messages(command['user_id'])
         if len(messages) == 0:
             msg = 'There is no active conversation'
         else:

knowledge_base_gpt/libs/history/history.py

@@ -0,0 +1,26 @@
+""" Manage history """
+from injector import inject, singleton
+
+from knowledge_base_gpt.libs.history.base import HistoryBase
+from knowledge_base_gpt.libs.history.redis import HistoryRedis
+from knowledge_base_gpt.libs.history.memory import HistoryMemory
+from knowledge_base_gpt.libs.settings.settings import Settings
+
+
+@singleton
+class History():  # pylint:disable=R0903
+    """ Wrap history manager based on settings """
+    @inject
+    def __init__(self, settings: Settings) -> None:
+        match settings.history.mode:
+            case 'memory':
+                self._history = HistoryMemory()
+            case 'redis':
+                self._history = HistoryRedis(settings.redis)
+            case _:
+                pass
+
+    @property
+    def history(self) -> HistoryBase:
+        """ Get history manager """
+        return self._history

knowledge_base_gpt/libs/history/memory.py

@@ -0,0 +1,42 @@
+""" Manage history in Memory """
+from typing import List, Dict, Any, Union
+
+from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
+
+from knowledge_base_gpt.libs.history.base import HistoryBase
+
+
+class HistoryMemory(HistoryBase):
+    """ Manage history in memory """
+    def __init__(self):
+        self._messages: Dict[str, List[BaseMessage]] = {}
+
+    def get_messages(self, session_id: str) -> List[BaseMessage]:
+        return self._messages.get(session_id, [])
+
+    def add_to_history(self, session_id: str, answer: Dict[str, Any]):
+        if self._messages.get(session_id) is None:
+            self._messages[session_id] = []
+        self._messages[session_id].extend(
+            [
+                self._get_human_message(answer['question']),
+                self._get_ai_message(answer['answer'])
+            ]
+
+        )
+
+    def reset(self, session_id: str):
+        if self._messages.get(session_id) is not None:
+            del self._messages[session_id]
+
+    @staticmethod
+    def _get_human_message(message: Union[HumanMessage, str]) -> HumanMessage:
+        if not isinstance(message, HumanMessage):
+            message = HumanMessage(content=message)
+        return message
+
+    @staticmethod
+    def _get_ai_message(message: Union[AIMessage, str]) -> AIMessage:
+        if not isinstance(message, AIMessage):
+            message = AIMessage(content=message)
+        return message

knowledge_base_gpt/libs/history/redis.py

@@ -1,20 +1,16 @@
 """ Manage history in Redis """
 from typing import List, Dict, Any
 
-from injector import inject, singleton
 from langchain_community.chat_message_histories import RedisChatMessageHistory
 from langchain_core.messages import BaseMessage
 
 from knowledge_base_gpt.libs.history.base import HistoryBase
-from knowledge_base_gpt.libs.settings.settings import RedisSettings, Settings
+from knowledge_base_gpt.libs.settings.settings import RedisSettings
 
 
-@singleton
 class HistoryRedis(HistoryBase):
     """ Manage history in Redis """
-    @inject
-    def __init__(self, settings: Settings):
-        redis_settings = settings.redis
+    def __init__(self, redis_settings: RedisSettings):
         self._url = self._build_url_string(redis_settings)
         self._ttl = redis_settings.ttl
 

knowledge_base_gpt/libs/settings/settings.py

@@ -129,6 +129,14 @@ class HuggingFaceSettings(BaseModel):
     )
 
 
+class HistorySettings(BaseModel):
+    """ History Settings """
+    mode: Literal['memory', 'redis'] = Field(
+        'memory',
+        description="Mode of History manager"
+    )
+
+
 class RedisSettings(BaseModel):
     """ Redis Settings """
     host: str = Field(
@@ -194,6 +202,7 @@ class Settings(BaseModel):
     slackbot: SlackBotSettings
     llm: LLMSettings
     ollama: OllamaSettings
+    history: HistorySettings
     redis: RedisSettings
     log: LogSettings
     content_loader: ContentLoaderSettings

settings.yaml

@@ -10,6 +10,8 @@ ollama:
   llm_model: $MODEL|mistral
   embedding_model: $EMBEDDINGS_MODEL_NAME|nomic-embed-text
 
+history: {}
+
 redis:
   host: $REDIS_HOST|localhost
   password: $REDIS_PASSWORD|""