adjust code graph pipeline to new run_tasks signature

topoteretes · Feb 10, 2025 · 1e40d03 · 1e40d03
1 parent 4a1ef00
commit 1e40d03
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 8 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -10,7 +10,7 @@ repos:
     -   id: check-added-large-files
 - repo: https://github.com/astral-sh/ruff-pre-commit
   # Ruff version.
-  rev: v0.9.0
+  rev: v0.9.5
   hooks:
     # Run the linter.
     - id: ruff

diff --git a/cognee/api/v1/cognify/code_graph_pipeline.py b/cognee/api/v1/cognify/code_graph_pipeline.py
@@ -69,9 +69,19 @@ async def run_code_graph_pipeline(repo_path, include_docs=True):
             ),
         ]
 
+    pipeline_run_status = None
     if include_docs:
-        async for result in run_tasks(non_code_tasks, repo_path):
-            yield result
+        non_code_pipeline_run = run_tasks(non_code_tasks, None, repo_path, "cognify_pipeline")
+        async for run_status in non_code_pipeline_run:
+            pipeline_run_status = run_status
 
-    async for result in run_tasks(tasks, repo_path, "cognify_code_pipeline"):
-        yield result
+    from cognee.modules.data.methods import get_datasets
+
+    existing_datasets = await get_datasets(user.id)
+    code_pipeline_run = run_tasks(
+        tasks, existing_datasets[0].id, repo_path, "cognify_code_pipeline"
+    )
+    async for run_status in code_pipeline_run:
+        pipeline_run_status = run_status
+
+    return pipeline_run_status
diff --git a/cognee/tasks/storage/index_data_points.py b/cognee/tasks/storage/index_data_points.py
@@ -36,7 +36,7 @@ async def index_data_points(data_points: list[DataPoint]):
     for index_name_and_field, indexable_points in index_points.items():
         first_occurence = index_name_and_field.index("_")
         index_name = index_name_and_field[:first_occurence]
-        field_name = index_name_and_field[first_occurence + 1:]
+        field_name = index_name_and_field[first_occurence + 1 :]
         try:
             await vector_engine.index_data_points(index_name, field_name, indexable_points)
         except EmbeddingException as e:

diff --git a/examples/python/code_graph_example.py b/examples/python/code_graph_example.py
@@ -7,8 +7,7 @@
 
 
 async def main(repo_path, include_docs):
-    async for result in run_code_graph_pipeline(repo_path, include_docs):
-        print(result)
+    return await run_code_graph_pipeline(repo_path, include_docs)
 
 
 def parse_args():