From 6e4addea17600a33ca187d0aea93d03425c4da64 Mon Sep 17 00:00:00 2001 From: Daniel Molnar Date: Tue, 13 May 2025 11:22:49 +0200 Subject: [PATCH] Fix graphiti example. --- .../index_graphiti_objects.py | 20 ++++++++++--------- examples/python/graphiti_example.py | 6 +++++- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/cognee/tasks/temporal_awareness/index_graphiti_objects.py b/cognee/tasks/temporal_awareness/index_graphiti_objects.py index 619ca958de..2124b830d8 100644 --- a/cognee/tasks/temporal_awareness/index_graphiti_objects.py +++ b/cognee/tasks/temporal_awareness/index_graphiti_objects.py @@ -20,20 +20,23 @@ async def index_and_transform_graphiti_nodes_and_edges(): logger.error("Failed to initialize engines: %s", e) raise RuntimeError("Initialization error") from e - await graph_engine.query("""MATCH (n) SET n.id = n.uuid RETURN n""") + await graph_engine.query("""MATCH (n) SET n.id = n.uuid RETURN n""", params={}) await graph_engine.query( """MATCH (source)-[r]->(target) SET r.source_node_id = source.id, r.target_node_id = target.id, - r.relationship_name = type(r) RETURN r""" + r.relationship_name = type(r) RETURN r""", + params={}, + ) + await graph_engine.query( + """MATCH (n) SET n.text = COALESCE(n.summary, n.content) RETURN n""", params={} ) - await graph_engine.query("""MATCH (n) SET n.text = COALESCE(n.summary, n.content) RETURN n""") - nodes_data, edges_data = await graph_engine.get_model_independent_graph_data() + nodes_data, edges_data = await graph_engine.get_graph_data() - for node_data in nodes_data[0]["nodes"]: + for node_id, node_data in nodes_data: graphiti_node = GraphitiNode( **{key: node_data[key] for key in ("content", "name", "summary") if key in node_data}, - id=node_data.get("uuid"), + id=node_id, ) data_point_type = type(graphiti_node) @@ -58,9 +61,8 @@ async def index_and_transform_graphiti_nodes_and_edges(): await vector_engine.index_data_points(index_name, field_name, indexable_points) edge_types = Counter( - edge[1][1] - for edge in edges_data[0]["elements"] - if isinstance(edge, list) and len(edge) == 3 + edge[2] # The edge key (relationship name) is at index 2 + for edge in edges_data ) for text, count in edge_types.items(): diff --git a/examples/python/graphiti_example.py b/examples/python/graphiti_example.py index 009c2f2f42..5ad964fcd7 100644 --- a/examples/python/graphiti_example.py +++ b/examples/python/graphiti_example.py @@ -14,6 +14,7 @@ from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt from cognee.infrastructure.llm.get_llm_client import get_llm_client +from cognee.modules.users.methods import get_default_user text_list = [ "Kamala Harris is the Attorney General of California. She was previously " @@ -27,6 +28,9 @@ async def main(): await cognee.prune.prune_system(metadata=True) await create_relational_db_and_tables() + # Initialize default user + user = await get_default_user() + for text in text_list: await cognee.add(text) @@ -34,7 +38,7 @@ async def main(): Task(build_graph_with_temporal_awareness, text_list=text_list), ] - pipeline = run_tasks(tasks) + pipeline = run_tasks(tasks, user=user) async for result in pipeline: print(result)