diff --git a/.github/workflows/test_memgraph.yml b/.github/workflows/test_memgraph.yml
index c9134bc1f..a6f549a22 100644
--- a/.github/workflows/test_memgraph.yml
+++ b/.github/workflows/test_memgraph.yml
@@ -56,4 +56,4 @@ jobs:
           GRAPH_DATABASE_URL: "bolt://localhost:7687"
           GRAPH_DATABASE_PASSWORD: "memgraph"
           GRAPH_DATABASE_USERNAME: "memgraph"
-        run: poetry run python ./cognee/tests/test_memgraph.py
+        run: poetry run python ./cognee/tests/unit/infrastructure/databases/vector/test_memgraph.py
diff --git a/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py
index 73762bbf2..988569565 100644
--- a/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py
+++ b/cognee/infrastructure/llm/structured_output_framework/baml_src/extraction/extract_summary.py
@@ -36,13 +36,17 @@ async def extract_summary(content: str, response_model: Type[BaseModel]):
     config = get_llm_config()
 
     # Use BAML's SummarizeContent function
-    summary_result = await b.SummarizeContent(content, baml_options={"client_registry": config.baml_registry})
+    summary_result = await b.SummarizeContent(
+        content, baml_options={"client_registry": config.baml_registry}
+    )
 
     # Convert BAML result to the expected response model
     if response_model is SummarizedCode:
         # If it's asking for SummarizedCode but we got SummarizedContent,
         # we need to use SummarizeCode instead
-        code_result = await b.SummarizeCode(content, baml_options={"client_registry": config.baml_registry})
+        code_result = await b.SummarizeCode(
+            content, baml_options={"client_registry": config.baml_registry}
+        )
         return code_result
     else:
         # For other models, return the summary result
@@ -70,7 +74,9 @@ async def extract_code_summary(content: str):
     else:
         try:
             config = get_llm_config()
-            result = await b.SummarizeCode(content, baml_options={"client_registry": config.baml_registry})
+            result = await b.SummarizeCode(
+                content, baml_options={"client_registry": config.baml_registry}
+            )
         except Exception as e:
             logger.error(
                 "Failed to extract code summary with BAML, falling back to mock summary", exc_info=e
diff --git a/cognee/tests/unit/infrastructure/databases/vector/__init__.py b/cognee/tests/unit/infrastructure/databases/vector/__init__.py
new file mode 100644
index 000000000..9399921ac
--- /dev/null
+++ b/cognee/tests/unit/infrastructure/databases/vector/__init__.py
@@ -0,0 +1 @@
+# Vector database tests module
diff --git a/cognee/tests/test_memgraph.py b/cognee/tests/unit/infrastructure/databases/vector/test_memgraph.py
similarity index 98%
rename from cognee/tests/test_memgraph.py
rename to cognee/tests/unit/infrastructure/databases/vector/test_memgraph.py
index 49f3e08db..00a1c2875 100644
--- a/cognee/tests/test_memgraph.py
+++ b/cognee/tests/unit/infrastructure/databases/vector/test_memgraph.py
@@ -32,7 +32,7 @@ async def main():
     dataset_name = "cs_explanations"
 
     explanation_file_path = os.path.join(
-        pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
+        pathlib.Path(__file__).parent, "../../../../../test_data/Natural_language_processing.txt"
     )
     await cognee.add([explanation_file_path], dataset_name)