cognee/cognitive_architecture/modules/cognify/graph/add_propositions.py
2024-03-09 18:39:00 +01:00

97 lines
3.4 KiB
Python

""" Here we update semantic graph with content that classifier produced"""
import uuid
from datetime import datetime
from enum import Enum, auto
from typing import Type, Optional, Any
from pydantic import BaseModel
from cognitive_architecture.infrastructure.databases.graph.get_graph_client import get_graph_client
from cognitive_architecture.shared.data_models import GraphDBType, DefaultGraphModel, Document, DocumentType, Category, Relationship, UserProperties, UserLocation
def add_propositions(G, category_name, subclass_content, layer_description, new_data, layer_uuid,
layer_decomposition_uuid):
""" Add nodes and edges to the graph for the given LLM knowledge graph and the layer"""
# Find the node ID for the subclass within the category
G.load_graph_from_file()
G = graph_client.graph
subclass_node_id = None
for node, data in G.nodes(data=True):
if subclass_content in node:
subclass_node_id = node
print(subclass_node_id)
if not subclass_node_id:
print(f"Subclass '{subclass_content}' under category '{category_name}' not found in the graph.")
return G
# Mapping from old node IDs to new node IDs
node_id_mapping = {}
# Add nodes from the Pydantic object
for node in new_data.nodes:
unique_node_id = uuid.uuid4()
new_node_id = f"{node.description} - {str(layer_uuid)} - {str(layer_decomposition_uuid)} - {str(unique_node_id)}"
G.add_node(new_node_id,
created_at=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
updated_at=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
description=node.description,
category=node.category,
memory_type=node.memory_type,
layer_uuid=str(layer_uuid),
layer_description=str(layer_description),
layer_decomposition_uuid=str(layer_decomposition_uuid),
id=str(unique_node_id),
type='detail')
G.add_edge(subclass_node_id, new_node_id, relationship='detail')
# Store the mapping from old node ID to new node ID
node_id_mapping[node.id] = new_node_id
# Add edges from the Pydantic object using the new node IDs
for edge in new_data.edges:
# Use the mapping to get the new node IDs
source_node_id = node_id_mapping.get(edge.source)
target_node_id = node_id_mapping.get(edge.target)
if source_node_id and target_node_id:
G.add_edge(source_node_id, target_node_id, description=edge.description, relationship='relation')
else:
print(f"Could not find mapping for edge from {edge.source} to {edge.target}")
return G
if __name__ == "__main__":
import asyncio
# Assuming all necessary imports and GraphDBType, get_graph_client, Document, DocumentType, etc. are defined
# Initialize the graph client
graph_client = get_graph_client(GraphDBType.NETWORKX)
G = asyncio.run(add_propositions(graph_client, 'category_name', 'subclass_content', 'layer_description', 'new_data', 'layer_uuid',
'layer_decomposition_uuid'))
if __name__ == "__main__":
import asyncio
# Assuming all necessary imports and GraphDBType, get_graph_client, Document, DocumentType, etc. are defined
# Initialize the graph client
graph_client = get_graph_client(GraphDBType.NETWORKX)