From 82ac9fc26a01acba224d65d59559522155061e9b Mon Sep 17 00:00:00 2001
From: Vasilije <8619304+Vasilije1990@users.noreply.github.com>
Date: Tue, 6 Aug 2024 18:28:32 +0200
Subject: [PATCH] Fix to the task

---
 .../v1/tasks/graph_ontology/graph_ontology.py | 11 ----------
 cognee/tasks/graph_ontology/graph_ontology.py | 22 +++++++++++++++++++
 2 files changed, 22 insertions(+), 11 deletions(-)
 delete mode 100644 cognee/api/v1/tasks/graph_ontology/graph_ontology.py
 create mode 100644 cognee/tasks/graph_ontology/graph_ontology.py

diff --git a/cognee/api/v1/tasks/graph_ontology/graph_ontology.py b/cognee/api/v1/tasks/graph_ontology/graph_ontology.py
deleted file mode 100644
index 7a09b4f29..000000000
--- a/cognee/api/v1/tasks/graph_ontology/graph_ontology.py
+++ /dev/null
@@ -1,11 +0,0 @@
-
-
-from cognee.modules.pipelines.tasks.Task import Task
-
-from cognee.modules.data.extraction.knowledge_graph.establish_graph_topology import establish_graph_topology
-from cognee.shared.data_models import KnowledgeGraph
-
-
-
-async def ontology_task():
-    return Task(establish_graph_topology, topology_model = KnowledgeGraph, task_config = { "batch_size": 10 })
\ No newline at end of file
diff --git a/cognee/tasks/graph_ontology/graph_ontology.py b/cognee/tasks/graph_ontology/graph_ontology.py
new file mode 100644
index 000000000..1e502cd61
--- /dev/null
+++ b/cognee/tasks/graph_ontology/graph_ontology.py
@@ -0,0 +1,22 @@
+from typing import Type
+from pydantic import BaseModel
+
+from cognee.modules.data.processing.chunk_types.DocumentChunk import DocumentChunk
+from cognee.shared.data_models import KnowledgeGraph
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.modules.data.extraction.knowledge_graph.add_model_class_to_graph import add_model_class_to_graph
+
+
+async def establish_graph_topology(data_chunks: list[DocumentChunk], topology_model: Type[BaseModel]):
+    if topology_model == KnowledgeGraph:
+        return data_chunks
+
+    graph_engine = await get_graph_engine()
+
+    await add_model_class_to_graph(topology_model, graph_engine)
+
+    return data_chunks
+
+
+def generate_node_id(node_id: str) -> str:
+    return node_id.upper().replace(" ", "_").replace("'", "")