chore: formats cognee starter kit with ruff

Author: hajdul88

cognee-starter-kit/src/pipelines/custom-model.py

@@ -2,26 +2,24 @@
 import asyncio
 import pathlib
 from cognee import config, add, cognify, search, SearchType, prune, visualize_graph
+
 # from cognee.shared.utils import render_graph
 from cognee.low_level import DataPoint
 
+
 async def main():
     data_directory_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".data_storage")
-        ).resolve()
+        pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage")).resolve()
     )
     # Set up the data directory. Cognee will store files here.
     config.data_root_directory(data_directory_path)
 
     cognee_directory_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".cognee_system")
-        ).resolve()
+        pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system")).resolve()
     )
     # Set up the Cognee system directory. Cognee will store system files and databases here.
     config.system_root_directory(cognee_directory_path)
-    
+
     # Prune data and system metadata before running, only if we want "fresh" state.
     await prune.prune_data()
     await prune.prune_system(metadata=True)
@@ -31,7 +29,6 @@ async def main():
     # Add the text data to Cognee.
     await add(text)
 
-
     # Define a custom graph model for programming languages.
     class FieldType(DataPoint):
         name: str = "Field"
@@ -50,7 +47,6 @@ class ProgrammingLanguage(DataPoint):
         is_type: ProgrammingLanguageType
         metadata: dict = {"index_fields": ["name"]}
 
-
     # Cognify the text data.
     await cognify(graph_model=ProgrammingLanguage)
 
@@ -67,12 +63,16 @@ class ProgrammingLanguage(DataPoint):
     await visualize_graph(graph_file_path)
 
     # Completion query that uses graph data to form context.
-    graph_completion = await search(query_text="What is python?", query_type=SearchType.GRAPH_COMPLETION)
+    graph_completion = await search(
+        query_text="What is python?", query_type=SearchType.GRAPH_COMPLETION
+    )
     print("Graph completion result is:")
     print(graph_completion)
 
     # Completion query that uses document chunks to form context.
-    rag_completion = await search(query_text="What is Python?", query_type=SearchType.RAG_COMPLETION)
+    rag_completion = await search(
+        query_text="What is Python?", query_type=SearchType.RAG_COMPLETION
+    )
     print("Completion result is:")
     print(rag_completion)
 

cognee-starter-kit/src/pipelines/default.py

@@ -4,19 +4,16 @@
 from cognee import config, add, cognify, search, SearchType, prune, visualize_graph
 # from cognee.shared.utils import render_graph
 
+
 async def main():
     data_directory_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".data_storage")
-        ).resolve()
+        pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage")).resolve()
     )
     # Set up the data directory. Cognee will store files here.
     config.data_root_directory(data_directory_path)
 
     cognee_directory_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".cognee_system")
-        ).resolve()
+        pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".cognee_system")).resolve()
     )
     # Set up the Cognee system directory. Cognee will store system files and databases here.
     config.system_root_directory(cognee_directory_path)
@@ -46,12 +43,16 @@ async def main():
     await visualize_graph(graph_file_path)
 
     # Completion query that uses graph data to form context.
-    graph_completion = await search(query_text="What is python?", query_type=SearchType.GRAPH_COMPLETION)
+    graph_completion = await search(
+        query_text="What is python?", query_type=SearchType.GRAPH_COMPLETION
+    )
     print("Graph completion result is:")
     print(graph_completion)
 
     # Completion query that uses document chunks to form context.
-    rag_completion = await search(query_text="What is Python?", query_type=SearchType.RAG_COMPLETION)
+    rag_completion = await search(
+        query_text="What is Python?", query_type=SearchType.RAG_COMPLETION
+    )
     print("Completion result is:")
     print(rag_completion)
 

cognee-starter-kit/src/pipelines/low_level.py

@@ -10,6 +10,7 @@
 from cognee.tasks.storage.index_graph_edges import index_graph_edges
 from cognee.modules.users.methods import get_default_user
 
+
 class Person(DataPoint):
     name: str
     metadata: dict = {"index_fields": ["name"]}
@@ -85,7 +86,6 @@ async def main():
 
     await setup()
 
-
     # Generate a random dataset_id
     dataset_id = uuid.uuid4()
     user = await get_default_user()
@@ -98,7 +98,7 @@ async def main():
         dataset_id,
         None,
         user,
-        "demo_pipeline"
+        "demo_pipeline",
     )
 
     async for status in pipeline:

cognee/tests/test_starter_pipelines.py

@@ -9,16 +9,11 @@ class TestPipelines(unittest.TestCase):
 
     def setUp(self):
         # Ensure we're in the correct directory
-        self.project_root = os.path.abspath(os.path.join(
-            os.path.dirname(__file__), '../..'
-        ))
-        self.pipelines_dir = os.path.join(self.project_root, 'src', 'pipelines')
+        self.project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
+        self.pipelines_dir = os.path.join(self.project_root, "src", "pipelines")
 
         # Required environment variables
-        self.required_env_vars = [
-            'LLM_API_KEY',
-            'EMBEDDING_API_KEY'
-        ]
+        self.required_env_vars = ["LLM_API_KEY", "EMBEDDING_API_KEY"]
 
         # Check if required environment variables are set
         missing_vars = [var for var in self.required_env_vars if not os.environ.get(var)]
@@ -30,7 +25,7 @@ def _run_pipeline(self, script_name):
         script_path = os.path.join(self.pipelines_dir, script_name)
 
         # Use the Python executable from the virtual environment
-        python_exe = os.path.join(self.project_root, '.venv', 'bin', 'python')
+        python_exe = os.path.join(self.project_root, ".venv", "bin", "python")
         if not os.path.exists(python_exe):
             python_exe = sys.executable
 
@@ -40,30 +35,32 @@ def _run_pipeline(self, script_name):
                 check=True,
                 capture_output=True,
                 text=True,
-                timeout=300  # 5 minute timeout
+                timeout=300,  # 5 minute timeout
             )
             return result
         except subprocess.CalledProcessError as e:
-            self.fail(f"Pipeline {script_name} failed with code {e.returncode}. "
-                      f"Stdout: {e.stdout}, Stderr: {e.stderr}")
+            self.fail(
+                f"Pipeline {script_name} failed with code {e.returncode}. "
+                f"Stdout: {e.stdout}, Stderr: {e.stderr}"
+            )
         except subprocess.TimeoutExpired:
             self.fail(f"Pipeline {script_name} timed out after 300 seconds")
 
     def test_default_pipeline(self):
         """Test that the default pipeline runs successfully."""
-        result = self._run_pipeline('default.py')
+        result = self._run_pipeline("default.py")
         self.assertEqual(result.returncode, 0)
 
     def test_low_level_pipeline(self):
         """Test that the low-level pipeline runs successfully."""
-        result = self._run_pipeline('low_level.py')
+        result = self._run_pipeline("low_level.py")
         self.assertEqual(result.returncode, 0)
 
     def test_custom_model_pipeline(self):
         """Test that the custom model pipeline runs successfully."""
-        result = self._run_pipeline('custom-model.py')
+        result = self._run_pipeline("custom-model.py")
         self.assertEqual(result.returncode, 0)
 
 
-if __name__ == '__main__':
-    unittest.main()
+if __name__ == "__main__":
+    unittest.main()