Fix testing structure for load_csv and removes unnecessary commits from session

Author: vmesel

src/dialog/db/__init__.py

@@ -15,12 +15,9 @@ def session_scope():
     with Session(bind=engine) as session:
         try:
             yield session
-            session.commit()
         except Exception as exc:
             session.rollback()
             raise exc
-        finally:
-            session.close()
 
 def get_session():
     with session_scope() as session:

src/dialog/routers/openai.py

@@ -55,7 +55,7 @@ async def ask_question_to_llm(message: OpenAIChat, session: Session = Depends(ge
             session_id = Settings().OPENWEB_UI_SESSION,
         )
         session.add(new_chat)
-        session.flush()
+        session.commit()
     else:
         logging.info("Using old chat entity")
         new_chat = chat_entity
@@ -72,7 +72,7 @@ async def ask_question_to_llm(message: OpenAIChat, session: Session = Depends(ge
             message=_message.content,
         )
         session.add(new_message)
-    session.flush()
+    session.commit()
 
     process_user_message_args = {
         "message": non_empty_messages[-1].content,

src/load_csv.py

@@ -150,6 +150,7 @@ def load_csv_and_generate_embeddings(
     ]
     session.add_all(company_contents)
     session.commit()
+    return session.query(CompanyContent).all()
 
 
 if __name__ == "__main__":

src/tests/conftest.py

@@ -30,7 +30,6 @@ def dbsession(mocker):
 
     with Session() as session:
         yield session
-        session.rollback()
 
     Base.metadata.drop_all(bind=engine)
 

src/tests/test_load_csv.py

@@ -47,9 +47,8 @@ def test_load_csv(mocker, dbsession, csv_file: str):
         [0.2] * 1536,
     ]  # 1536 is the expected dimension of the embeddings
 
-    load_csv.load_csv_and_generate_embeddings(csv_file, cleardb=True)
+    result = load_csv.load_csv_and_generate_embeddings(csv_file, cleardb=True)
 
-    result = dbsession.query(load_csv.CompanyContent).all()
     assert len(result) == 2
 
 
@@ -66,7 +65,7 @@ def test_multiple_columns_embedding(mocker, dbsession, csv_file: str):
 
     mock_generate_embeddings.assert_called_with(
         [
-            "category: cat1\nsubcategory: subcat1\ncontent: content1", 
+            "category: cat1\nsubcategory: subcat1\ncontent: content1",
             "category: cat2\nsubcategory: subcat2\ncontent: content2"
         ],
         embedding_llm_instance=load_csv.EMBEDDINGS_LLM,
@@ -80,20 +79,17 @@ def test_clear_db(mocker, dbsession, csv_file: str):
         [0.2] * 1536,
     ]  # 1536 is the expected dimension of the embeddings
 
-    load_csv.load_csv_and_generate_embeddings(csv_file, cleardb=True)
-    initial_run = dbsession.query(load_csv.CompanyContent).all()
+    initial_run = load_csv.load_csv_and_generate_embeddings(csv_file, cleardb=True)
 
-    load_csv.load_csv_and_generate_embeddings(csv_file, cleardb=True)
-    clear_db_run = dbsession.query(load_csv.CompanyContent).all()
+    clear_db_run = load_csv.load_csv_and_generate_embeddings(csv_file, cleardb=True)
 
     other_csv_file = _create_csv(
         data=[
             ["cat3", "subcat3", "q3", "content3", "dataset3"],
             ["cat4", "subcat4", "q4", "content4", "dataset4"],
         ]
     )
-    load_csv.load_csv_and_generate_embeddings(other_csv_file, cleardb=False)
-    dont_clear_db_run = dbsession.query(load_csv.CompanyContent).all()
+    dont_clear_db_run = load_csv.load_csv_and_generate_embeddings(other_csv_file, cleardb=False)
 
     assert len(initial_run) == 2
     assert len(clear_db_run) == 2
@@ -122,7 +118,7 @@ def test_documents_to_company_content():
             "link": "http://test_link"
         }
     )
-    
+
     # Define a mock embedding
     embedding = [0.1] * 1536  # Example embedding
 
@@ -155,7 +151,7 @@ def test_get_document_pk():
             "link": "http://test_link"
         }
     )
-    
+
     # Define the fields to be used for primary key generation
     pk_metadata_fields = ["category", "subcategory", "question"]
 
@@ -172,7 +168,7 @@ def test_get_document_pk():
 def test_load_csv_with_metadata(csv_file: str):
     metadata_columns = ["category", "subcategory", "question", "dataset"]
     embed_columns = ["content"]
-    
+
     # Call the function to test
     docs = load_csv.load_csv_with_metadata(csv_file, embed_columns, metadata_columns)
 
@@ -195,7 +191,7 @@ def test_load_csv_with_metadata(csv_file: str):
         "content": "content2",
     }
 
-def test_retrieve_docs_from_vectordb(mocker):
+def test_retrieve_docs_from_vectordb(mocker, dbsession):
     # Create mock CompanyContent objects
     mock_company_contents = [
         CompanyContent(

src/tests/test_views.py

@@ -31,7 +31,7 @@ def test_ask_question_no_session_id(client, mocker, llm_mock, dbsession):
 def test_get_chat_content(client, chat_session, dbsession):
     chat = ChatMessages(session_id=chat_session["chat_id"], message="Hello")
     dbsession.add(chat)
-    dbsession.flush()
+    dbsession.commit()
     response = client.get(f"/chat/{chat_session['chat_id']}")
     assert response.status_code == 200
     assert "message" in response.json()