gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix to the task

Browse source
This commit is contained in:
Vasilije 2024-08-06 18:28:32 +02:00
parent 64f54a32c6
commit 82ac9fc26a
2 changed files with 22 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
cognee/api/v1/tasks/graph_ontology/graph_ontology.py
View file

@ -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 })

22
cognee/tasks/graph_ontology/graph_ontology.py Normal file
View file

@ -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("'", "")