From b5c6628cfe5a28b86a9349c954b46bbeba77e9ce Mon Sep 17 00:00:00 2001
From: vujke <stefanvujovic93@gmail.com>
Date: Sat, 1 Jun 2024 18:43:04 +0200
Subject: [PATCH] add_coginitive_layer_graphs fix

---
 .../infrastructure/databases/graph/config.py   | 18 +++++++++++-------
 .../graph/add_cognitive_layer_graphs.py        |  2 ++
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/cognee/infrastructure/databases/graph/config.py b/cognee/infrastructure/databases/graph/config.py
index ea852947c..bdb122823 100644
--- a/cognee/infrastructure/databases/graph/config.py
+++ b/cognee/infrastructure/databases/graph/config.py
@@ -1,9 +1,11 @@
 """ This module contains the configuration for the graph database. """
+
 import os
 from functools import lru_cache
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from cognee.infrastructure.databases.relational.config import get_relationaldb_config
-from cognee.shared.data_models import DefaultGraphModel, GraphDBType
+from cognee.shared.data_models import DefaultGraphModel, GraphDBType, KnowledgeGraph
+
 
 class GraphConfig(BaseSettings):
     graph_filename: str = "cognee_graph.pkl"
@@ -12,11 +14,13 @@ class GraphConfig(BaseSettings):
     graph_database_username: str = ""
     graph_database_password: str = ""
     graph_database_port: int = 123
-    graph_file_path: str = os.path.join(get_relationaldb_config().db_path, graph_filename)
-    graph_engine: object =  GraphDBType.NETWORKX
-    graph_model: object = DefaultGraphModel
+    graph_file_path: str = os.path.join(
+        get_relationaldb_config().db_path, graph_filename
+    )
+    graph_engine: object = GraphDBType.NETWORKX
+    graph_model: object = KnowledgeGraph
 
-    model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
+    model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
     def to_dict(self) -> dict:
         return {
@@ -28,10 +32,10 @@ class GraphConfig(BaseSettings):
             "graph_database_username": self.graph_database_username,
             "graph_database_password": self.graph_database_password,
             "graph_database_port": self.graph_database_port,
-            "graph_engine": self.graph_engine
-
+            "graph_engine": self.graph_engine,
         }
 
+
 @lru_cache
 def get_graph_config():
     return GraphConfig()
diff --git a/cognee/modules/cognify/graph/add_cognitive_layer_graphs.py b/cognee/modules/cognify/graph/add_cognitive_layer_graphs.py
index 0b05d5fe5..8d2e3fb7c 100644
--- a/cognee/modules/cognify/graph/add_cognitive_layer_graphs.py
+++ b/cognee/modules/cognify/graph/add_cognitive_layer_graphs.py
@@ -3,10 +3,12 @@ from uuid import uuid4
 from typing import List, Tuple, TypedDict
 from pydantic import BaseModel
 from cognee.infrastructure.databases.vector import DataPoint
+
 # from cognee.utils import extract_pos_tags, extract_named_entities, extract_sentiment_vader
 from cognee.infrastructure.databases.graph.config import get_graph_config
 from cognee.infrastructure.databases.vector.config import get_vectordb_config
 
+
 class GraphLike(TypedDict):
     nodes: List
     edges: List