diff --git a/examples/python/graphiti_example.py b/examples/python/graphiti_example.py index 156b372c1..7e44872cf 100644 --- a/examples/python/graphiti_example.py +++ b/examples/python/graphiti_example.py @@ -1,16 +1,20 @@ import asyncio import cognee -from cognee.api.v1.search import SearchType +import logging from cognee.modules.pipelines import Task, run_tasks -from cognee.tasks.temporal_awareness import ( - build_graph_with_temporal_awareness, - search_graph_with_temporal_awareness, -) +from cognee.shared.utils import setup_logging +from cognee.tasks.temporal_awareness import build_graph_with_temporal_awareness from cognee.infrastructure.databases.relational import ( create_db_and_tables as create_relational_db_and_tables, ) -from cognee.tasks.storage.index_graph_edges import index_graphiti_nodes_and_edges +from cognee.tasks.temporal_awareness.index_graphiti_objects import ( + index_and_transform_graphiti_nodes_and_edges, +) +from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search +from cognee.tasks.completion.graph_query_completion import retrieved_edges_to_string +from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt +from cognee.infrastructure.llm.get_llm_client import get_llm_client text_list = [ "Kamala Harris is the Attorney General of California. She was previously " @@ -36,8 +40,33 @@ async def main(): async for result in pipeline: print(result) - await index_graphiti_nodes_and_edges() + await index_and_transform_graphiti_nodes_and_edges() + + query = "When was Kamala Harris in office?" + triplets = await brute_force_triplet_search( + query=query, + top_k=3, + collections=["graphitinode_content", "graphitinode_name", "graphitinode_summary"], + ) + + args = { + "question": query, + "context": retrieved_edges_to_string(triplets), + } + + user_prompt = render_prompt("graph_context_for_question.txt", args) + system_prompt = read_query_prompt("answer_simple_question_restricted.txt") + + llm_client = get_llm_client() + computed_answer = await llm_client.acreate_structured_output( + text_input=user_prompt, + system_prompt=system_prompt, + response_model=str, + ) + + print(computed_answer) if __name__ == "__main__": + setup_logging(logging.INFO) asyncio.run(main())