Changes session

Author: vmesel

src/dialog/llm/abstract_llm.py

@@ -1,6 +1,7 @@
 from langchain.chains.llm import LLMChain
 from langchain.memory.chat_memory import BaseChatMemory
 from langchain.prompts import ChatPromptTemplate
+from dialog.models.db import get_session
 
 
 class AbstractLLM:
@@ -25,7 +26,7 @@ def __init__(self, config, session_id=None, parent_session_id=None, dataset=None
         self.dataset = dataset
         self.llm_key = llm_key
         self.parent_session_id = parent_session_id
-        self.dbsession = dbsession
+        self.dbsession = dbsession or next(get_session())
 
     def get_prompt(self, input) -> ChatPromptTemplate:
         """

src/dialog/llm/embeddings.py

@@ -24,6 +24,9 @@ def generate_embedding(document: str):
 
 
 def get_most_relevant_contents_from_message(message, top=5, dataset=None, session=None):
+    if session is None:
+        session = next(get_session())
+
     message_embedding = generate_embedding(message)
     filters = [
         CompanyContent.embedding.cosine_distance(message_embedding) < Settings().COSINE_SIMILARITY_THRESHOLD,

src/dialog/llm/memory.py

@@ -13,7 +13,7 @@ class CustomPostgresChatMessageHistory(PostgresChatMessageHistory):
 
     def __init__(self, *args, parent_session_id=None, dbsession=None, **kwargs):
         self.parent_session_id = parent_session_id
-        self.dbsession = dbsession
+        self.dbsession = dbsession or next(get_session())
         super().__init__(*args, **kwargs)
 
     def _create_table_if_not_exists(self) -> None:
@@ -52,6 +52,9 @@ def generate_memory_instance(session_id, parent_session_id=None, dbsession=None)
     """
     Generate a memory instance for a given session_id
     """
+    if not dbsession:
+        dbsession = next(get_session())
+
     return CustomPostgresChatMessageHistory(
         connection_string=Settings().DATABASE_URL,
         session_id=session_id,
@@ -69,6 +72,9 @@ def add_user_message_to_message_history(session_id, message, memory=None, dbsess
     if not memory:
         memory = generate_memory_instance(session_id)
 
+    if not dbsession:
+        dbsession = next(get_session())
+
     memory.add_user_message(message)
     return memory
 
@@ -77,5 +83,8 @@ def get_messages(session_id, dbsession=None):
     """
     Get all messages for a given session_id
     """
+    if not dbsession:
+        dbsession = next(get_session())
+
     memory = generate_memory_instance(session_id, dbsession=dbsession)
     return memory.messages