From 2e7cfd8ab78c4f19be6e525f9240d7bd806d215e Mon Sep 17 00:00:00 2001
From: Vasilije <8619304+Vasilije1990@users.noreply.github.com>
Date: Sat, 25 May 2024 20:05:25 +0200
Subject: [PATCH] fix db engine config

---
 cognee/modules/tasks/create_task_status_table.py | 6 +++---
 cognee/tests/test_library.py                     | 1 -
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/cognee/modules/tasks/create_task_status_table.py b/cognee/modules/tasks/create_task_status_table.py
index ec9e045aa..26beceb1c 100644
--- a/cognee/modules/tasks/create_task_status_table.py
+++ b/cognee/modules/tasks/create_task_status_table.py
@@ -1,7 +1,7 @@
-from cognee.infrastructure.databases.relational.config import get_relationaldb_config
-config = get_relationaldb_config()
+from cognee.infrastructure.InfrastructureConfig import infrastructure_config
+
 def create_task_status_table():
-    db_engine = config.db_engine
+    db_engine = infrastructure_config.get_config("database_engine")
 
     db_engine.create_table("cognee_task_status", [
         dict(name = "data_id", type = "STRING"),
diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py
index cc02e533a..eaac864bc 100755
--- a/cognee/tests/test_library.py
+++ b/cognee/tests/test_library.py
@@ -25,7 +25,6 @@ async def  main():
     dataset_name = "cs_explanations"
 
     explanation_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt")
-    #
     await cognee.add([explanation_file_path], dataset_name)
     
     # dataset_name = "short_stories"