Added graph intefrace, added neo4j + networkx structure and updates to the notebook
This commit is contained in:
parent
30f758a6a3
commit
b530f19e49
2 changed files with 91 additions and 3 deletions
|
|
@ -14,10 +14,27 @@ class GraphDBInterface(Protocol):
|
|||
@abstractmethod
|
||||
async def update_graph(
|
||||
self,
|
||||
collection_name: str,
|
||||
collection_config: object
|
||||
graph_name: str,
|
||||
graph_config: object
|
||||
): raise NotImplementedError
|
||||
|
||||
|
||||
""" Save and Load Graphs """
|
||||
|
||||
@abstractmethod
|
||||
async def save_graph(
|
||||
self,
|
||||
graph_name: str
|
||||
): raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def load_graph(
|
||||
self,
|
||||
graph_name: str
|
||||
): raise NotImplementedError
|
||||
|
||||
""" Collections """
|
||||
|
||||
# @abstractmethod
|
||||
# async def delete_collection(
|
||||
# self,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,76 @@
|
|||
import pickle
|
||||
from datetime import datetime
|
||||
import aiofiles
|
||||
import networkx as nx
|
||||
from databases.graph.graph_db_interface import GraphDBInterface
|
||||
|
||||
|
||||
class NetworXDB(GraphDBInterface):
|
||||
pass
|
||||
def __init__(self, filename="cognee_graph.pkl"):
|
||||
self.filename = filename
|
||||
self.graph = nx.MultiDiGraph()
|
||||
|
||||
|
||||
async def save_graph(self):
|
||||
try:
|
||||
async with aiofiles.open(self.filename, "wb") as f:
|
||||
await f.write(pickle.dumps(self.graph))
|
||||
self.logger.info("Graph saved successfully.")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to save graph: {e}")
|
||||
|
||||
async def load_graph(self):
|
||||
try:
|
||||
async with aiofiles.open(self.filename, "rb") as f:
|
||||
data = await f.read()
|
||||
self.graph = pickle.loads(data)
|
||||
self.logger.info("Graph loaded successfully.")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to load graph: {e}")
|
||||
|
||||
async def create_graph(self, user_id, custom_user_properties=None, required_layers=None, default_fields=None):
|
||||
"""Asynchronously create or update a user content graph based on given parameters."""
|
||||
# Assume required_layers is a dictionary-like object; use more robust validation in production
|
||||
category_name = required_layers['name']
|
||||
subgroup_names = [subgroup['name'] for subgroup in required_layers['cognitive_subgroups']]
|
||||
|
||||
# Construct the additional_categories structure
|
||||
additional_categories = {category_name: subgroup_names}
|
||||
|
||||
# Define default fields for all nodes if not provided
|
||||
if default_fields is None:
|
||||
default_fields = {
|
||||
'created_at': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
'updated_at': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
|
||||
# Merge custom user properties with default fields; custom properties take precedence
|
||||
user_properties = {**default_fields, **(custom_user_properties or {})}
|
||||
|
||||
# Default content categories and update with any additional categories provided
|
||||
content_categories = {
|
||||
"Temporal": ["Historical events", "Schedules and timelines"],
|
||||
"Positional": ["Geographical locations", "Spatial data"],
|
||||
"Propositions": ["Hypotheses and theories", "Claims and arguments"],
|
||||
"Personalization": ["User preferences", "User information"]
|
||||
}
|
||||
content_categories.update(additional_categories)
|
||||
|
||||
# Ensure the user node exists with properties
|
||||
self.graph.add_node(user_id, **user_properties, exist=True)
|
||||
|
||||
# Add or update content category nodes and their edges
|
||||
for category, subclasses in content_categories.items():
|
||||
category_properties = {**default_fields, 'type': 'category'}
|
||||
self.graph.add_node(category, **category_properties, exist=True)
|
||||
self.graph.add_edge(user_id, category, relationship='created')
|
||||
|
||||
# Add or update subclass nodes and their edges
|
||||
for subclass in subclasses:
|
||||
subclass_node_id = f"{category}:{subclass}"
|
||||
subclass_properties = {**default_fields, 'type': 'subclass', 'content': subclass}
|
||||
self.graph.add_node(subclass_node_id, **subclass_properties, exist=True)
|
||||
self.graph.add_edge(category, subclass_node_id, relationship='includes')
|
||||
|
||||
# Save the graph asynchronously after modifications
|
||||
await self.save_graph()
|
||||
Loading…
Add table
Reference in a new issue