fix: context injected, not session object

Author: vmesel

dialog_lib/agents/abstract.py

@@ -25,7 +25,7 @@ def __init__(
         parent_session_id=None,
         dataset=None,
         llm_api_key=None,
-        dbsession=get_session(),
+        dbsession=get_session,
     ):
         """
         :param config: Configuration dictionary
@@ -132,12 +132,13 @@ def document_prompt(self):
 
     @property
     def retriever(self):
-        return DialogRetriever(
-            session=self.dbsession,
-            embedding_llm=self.embedding_llm,
-            threshold=self.cosine_similarity_threshold,
-            top_k=self.top_k
-        )
+        with self.dbsession() as session:
+            return DialogRetriever(
+                session=session,
+                embedding_llm=self.embedding_llm,
+                threshold=self.cosine_similarity_threshold,
+                top_k=self.top_k
+            )
 
     @property
     def model(self):
@@ -203,20 +204,22 @@ def answer_runnable(self):
 
     @property
     def memory(self):
-        return get_memory_instance(
-            session_id=self.session_id,
-            sqlalchemy_session=self.dbsession,
-            database_url=self.config.get("database_url")
-        )
+        with self.dbsession() as session:
+            return get_memory_instance(
+                session_id=self.session_id,
+                sqlalchemy_session=session,
+                database_url=self.config.get("database_url")
+            )
 
     def get_session_history(self, something):
-        return CustomPostgresChatMessageHistory(
-            connection_string=self.config.get("database_url"),
-            session_id=self.session_id,
-            parent_session_id=self.parent_session_id,
-            table_name="chat_messages",
-            dbsession=self.dbsession,
-        )
+        with self.dbsession() as session:
+            return CustomPostgresChatMessageHistory(
+                connection_string=self.config.get("database_url"),
+                session_id=self.session_id,
+                parent_session_id=self.parent_session_id,
+                table_name="chat_messages",
+                dbsession=session,
+            )
 
     def chain_router(self, input):
         return self.answer_runnable if len(input["relevant_contents"]) > 0 else self.fallback_chain

dialog_lib/db/memory.py

@@ -15,7 +15,7 @@ def __init__(
         self,
         *args,
         parent_session_id=None,
-        dbsession=get_session(),
+        dbsession=get_session,
         chats_model=Chat,
         chat_messages_model=ChatMessages,
         ssl_mode=None,
@@ -48,10 +48,10 @@ def create_tables(self) -> None:
 
     def add_tags(self, tags: str) -> None:
         """Add tags for a given session_id/uuid on chats table"""
-        self.dbsession.query(self.chats_model).where(
-            self.chats_model.session_id == self._session_id
-        ).update({getattr(self.chats_model, "tags"): tags})
-        self.dbsession.commit()
+        with self.dbsession() as session:
+            session.query(self.chats_model).where(
+                self.chats_model.session_id == self._session_id
+            ).update({getattr(self.chats_model, "tags"): tags})
 
     def add_message(self, message: BaseMessage) -> None:
         """Append the message to the record in PostgreSQL"""
@@ -61,13 +61,12 @@ def add_message(self, message: BaseMessage) -> None:
         if self.parent_session_id:
             message.parent = self.parent_session_id
         self.dbsession.add(message)
-        self.dbsession.commit()
 
 
 def generate_memory_instance(
     session_id,
     parent_session_id=None,
-    dbsession=get_session(),
+    dbsession=get_session,
     database_url=None,
     chats_model=Chat,
     chat_messages_model=ChatMessages,
@@ -88,29 +87,31 @@ def generate_memory_instance(
 
 
 def add_user_message_to_message_history(
-    session_id, message, memory=None, dbsession=get_session(), database_url=None
+    session_id, message, memory=None, dbsession=get_session, database_url=None
 ):
     """
     Add a user message to the message history and returns the updated
     memory instance
     """
-    if not memory:
-        memory = generate_memory_instance(
-            session_id, dbsession=dbsession, database_url=database_url
-        )
+    with dbsession() as session:
+        if not memory:
+            memory = generate_memory_instance(
+                session_id, dbsession=session, database_url=database_url
+            )
 
-    memory.add_user_message(message)
-    return memory
+        memory.add_user_message(message)
+        return memory
 
 
-def get_messages(session_id, dbsession=get_session(), database_url=None):
+def get_messages(session_id, dbsession=get_session, database_url=None):
     """
     Get all messages for a given session_id
     """
-    memory = generate_memory_instance(
-        session_id, dbsession=dbsession, database_url=database_url
-    )
-    return memory.messages
+    with dbsession() as session:
+        memory = generate_memory_instance(
+            session_id, dbsession=session, database_url=database_url
+        )
+        return memory.messages
 
 def get_memory_instance(session_id, sqlalchemy_session, database_url):
     return generate_memory_instance(

dialog_lib/db/session.py

@@ -1,10 +1,22 @@
 import os
 
 import sqlalchemy as sa
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
+from contextlib import contextmanager
+
+engine = sa.create_engine(os.environ.get("DATABASE_URL"))
+Session = sessionmaker(bind=engine)
+
+@contextmanager
 def get_session():
-    engine = sa.create_engine(os.environ.get("DATABASE_URL"))
-    session = Session(engine)
-    yield session
-    session.close()
+    session = Session()
+    try:
+        yield session
+        session.flush()
+        session.commit()
+    except Exception as e:
+        session.rollback()
+        raise e
+    finally:
+        session.close()

dialog_lib/db/utils.py

@@ -3,14 +3,14 @@
 from .models import Chat
 
 
-def create_chat_session(identifier=None, dbsession=get_session(), model=Chat):
+def create_chat_session(identifier=None, dbsession=get_session, model=Chat):
     if identifier is None:
         identifier = uuid.uuid4().hex
 
-    chat = dbsession.query(model).filter_by(session_id=identifier).first()
-    if not chat:
-        chat = model(session_id=identifier)
-        dbsession.add(chat)
-        dbsession.commit()
+    with dbsession() as session:
+        chat = session.query(model).filter_by(session_id=identifier).first()
+        if not chat:
+            chat = model(session_id=identifier)
+            session.add(chat)
 
     return {"chat_id": chat.session_id}

dialog_lib/loaders/csv.py

@@ -7,7 +7,7 @@
 
 
 def load_csv(
-        file_path, dbsession=get_session(), embeddings_model_instance=None,
+        file_path, dbsession=get_session, embeddings_model_instance=None,
         embedding_llm_model=None, embedding_llm_api_key=None, company_id=None
     ):
 
@@ -20,22 +20,22 @@ def load_csv(
         else:
             raise ValueError("Invalid embeddings model")
 
-    for csv_content in contents:
-        content = {}
+    with dbsession() as session:
+        for csv_content in contents:
+            content = {}
 
-        for line in csv_content.page_content.split("\n"):
-            values = line.split(": ")
-            content[values[0]] = values[1]
+            for line in csv_content.page_content.split("\n"):
+                values = line.split(": ")
+                content[values[0]] = values[1]
 
-        company_content = CompanyContent(
-            category="csv",
-            subcategory="csv-content",
-            question=content["question"],
-            content=content["content"],
-            dataset=company_id,
-            embedding=generate_embedding(csv_content.page_content, embeddings_model_instance)
-        )
-        dbsession.add(company_content)
+            company_content = CompanyContent(
+                category="csv",
+                subcategory="csv-content",
+                question=content["question"],
+                content=content["content"],
+                dataset=company_id,
+                embedding=generate_embedding(csv_content.page_content, embeddings_model_instance)
+            )
+            session.add(company_content)
 
-    dbsession.commit()
     return company_content

dialog_lib/loaders/web.py

@@ -4,21 +4,21 @@
 from langchain_community.document_loaders import WebBaseLoader
 
 
-def load_webpage(url, embeddings_model_instance, dbsession=get_session(), company_id=None):
+def load_webpage(url, embeddings_model_instance, dbsession=get_session, company_id=None):
     loader = WebBaseLoader(url)
     contents = loader.load()
 
-    for url_content in contents:
-        company_content = CompanyContent(
-            link=url,
-            category="web",
-            subcategory="website-content",
-            question=url_content.metadata["title"],
-            content=url_content.page_content,
-            dataset=company_id,
-            embedding=generate_embedding(url_content.page_content, embeddings_model_instance)
-        )
-        dbsession.add(company_content)
-        dbsession.flush()
+    with dbsession() as session:
+        for url_content in contents:
+            company_content = CompanyContent(
+                link=url,
+                category="web",
+                subcategory="website-content",
+                question=url_content.metadata["title"],
+                content=url_content.page_content,
+                dataset=company_id,
+                embedding=generate_embedding(url_content.page_content, embeddings_model_instance)
+            )
+            session.add(company_content)
 
     return company_content

dialog_lib/tests/conftest.py

@@ -5,19 +5,18 @@
 from aioresponses import aioresponses
 from sqlalchemy.orm import Session
 from dialog_lib.db.models import Base
+from dialog_lib.db import get_session
 
 
 @pytest.fixture
 def db_engine():
     return sqlalchemy.create_engine(os.environ.get('DATABASE_URL'))
 
+
 @pytest.fixture
-def db_session(db_engine):
-    Base.metadata.create_all(db_engine)
-    session = Session(db_engine)
-    yield session
-    session.rollback()
-    session.close()
+def db_session():
+    return get_session
+
 
 @pytest.fixture
 def mock_aioresponse():

dialog_lib/tests/loaders/test_web_loader.py

@@ -9,9 +9,10 @@ def test_load_web_content(mock_aioresponse, db_session, mocker):
     mocker.patch('dialog_lib.loaders.web.generate_embedding', return_value=[0] * 1536)
     mock_aioresponse.get('http://example.com', body='Hello, world!')
 
-    content = load_webpage('http://example.com', None, db_session, 1)
-    assert content.question == "Example Domain"
-    assert content.embedding == [0] * 1536
+    load_webpage('http://example.com', None, db_session, 1)
+
+    with db_session() as session:
+        content = session.query(CompanyContent).first()
+        assert content.question == "Example Domain"
+        assert content.embedding.tolist() == [0]*1536
 
-    content = db_session.query(CompanyContent).all()
-    assert len(content) == 1