fix: graphiti example (#816)
<!-- .github/pull_request_template.md --> ## Description <!-- Provide a clear description of the changes in this PR --> ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin.
This commit is contained in:
parent
f93463e616
commit
8ea00971ad
2 changed files with 16 additions and 10 deletions
|
|
@ -20,20 +20,23 @@ async def index_and_transform_graphiti_nodes_and_edges():
|
||||||
logger.error("Failed to initialize engines: %s", e)
|
logger.error("Failed to initialize engines: %s", e)
|
||||||
raise RuntimeError("Initialization error") from 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(
|
await graph_engine.query(
|
||||||
"""MATCH (source)-[r]->(target) SET r.source_node_id = source.id,
|
"""MATCH (source)-[r]->(target) SET r.source_node_id = source.id,
|
||||||
r.target_node_id = target.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(
|
graphiti_node = GraphitiNode(
|
||||||
**{key: node_data[key] for key in ("content", "name", "summary") if key in node_data},
|
**{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)
|
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)
|
await vector_engine.index_data_points(index_name, field_name, indexable_points)
|
||||||
|
|
||||||
edge_types = Counter(
|
edge_types = Counter(
|
||||||
edge[1][1]
|
edge[2] # The edge key (relationship name) is at index 2
|
||||||
for edge in edges_data[0]["elements"]
|
for edge in edges_data
|
||||||
if isinstance(edge, list) and len(edge) == 3
|
|
||||||
)
|
)
|
||||||
|
|
||||||
for text, count in edge_types.items():
|
for text, count in edge_types.items():
|
||||||
|
|
|
||||||
|
|
@ -14,6 +14,7 @@ from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_forc
|
||||||
from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
|
from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
|
||||||
from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt
|
from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt
|
||||||
from cognee.infrastructure.llm.get_llm_client import get_llm_client
|
from cognee.infrastructure.llm.get_llm_client import get_llm_client
|
||||||
|
from cognee.modules.users.methods import get_default_user
|
||||||
|
|
||||||
text_list = [
|
text_list = [
|
||||||
"Kamala Harris is the Attorney General of California. She was previously "
|
"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 cognee.prune.prune_system(metadata=True)
|
||||||
await create_relational_db_and_tables()
|
await create_relational_db_and_tables()
|
||||||
|
|
||||||
|
# Initialize default user
|
||||||
|
user = await get_default_user()
|
||||||
|
|
||||||
for text in text_list:
|
for text in text_list:
|
||||||
await cognee.add(text)
|
await cognee.add(text)
|
||||||
|
|
||||||
|
|
@ -34,7 +38,7 @@ async def main():
|
||||||
Task(build_graph_with_temporal_awareness, text_list=text_list),
|
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:
|
async for result in pipeline:
|
||||||
print(result)
|
print(result)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue