diff --git a/cognee/tasks/temporal_awareness/__init__.py b/cognee/tasks/temporal_awareness/__init__.py new file mode 100644 index 000000000..9b6717f5f --- /dev/null +++ b/cognee/tasks/temporal_awareness/__init__.py @@ -0,0 +1,4 @@ +from .build_graph_with_temporal_awareness import \ + build_graph_with_temporal_awareness +from .search_graph_with_temporal_awareness import \ + search_graph_with_temporal_awareness diff --git a/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py b/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py new file mode 100644 index 000000000..4a7df6db0 --- /dev/null +++ b/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py @@ -0,0 +1,22 @@ +from datetime import datetime + +from graphiti_core import Graphiti +from graphiti_core.nodes import EpisodeType + + +async def build_graph_with_temporal_awareness(text_list): + + graphiti = Graphiti("bolt://localhost:7687", "neo4j", "pleaseletmein") + await graphiti.build_indices_and_constraints() + print("Graph database initialized.") + + for i, text in enumerate(text_list): + await graphiti.add_episode( + name=f"episode_{i}", + episode_body=text, + source=EpisodeType.text, + source_description="input", + reference_time=datetime.now() + ) + print(f"Added text: {text[:35]}...") + return graphiti \ No newline at end of file diff --git a/cognee/tasks/temporal_awareness/search_graph_with_temporal_awareness.py b/cognee/tasks/temporal_awareness/search_graph_with_temporal_awareness.py new file mode 100644 index 000000000..8e2cd0fc8 --- /dev/null +++ b/cognee/tasks/temporal_awareness/search_graph_with_temporal_awareness.py @@ -0,0 +1,6 @@ + + +async def search_graph_with_temporal_awareness(graphiti, query): + search_result = await graphiti.search(query) + await graphiti.close() + return search_result \ No newline at end of file diff --git a/examples/python/graphiti_example.py b/examples/python/graphiti_example.py new file mode 100644 index 000000000..a48fa9b4c --- /dev/null +++ b/examples/python/graphiti_example.py @@ -0,0 +1,29 @@ +import asyncio + +import cognee +from cognee.api.v1.search import SearchType +from cognee.modules.pipelines import Task, run_tasks +from cognee.tasks.temporal_awareness import ( + build_graph_with_temporal_awareness, search_graph_with_temporal_awareness) + +text_list = [ + "Kamala Harris is the Attorney General of California. She was previously " + "the district attorney for San Francisco.", + "As AG, Harris was in office from January 3, 2011 – January 3, 2017", +] + +async def main(): + + tasks = [ + Task(build_graph_with_temporal_awareness, text_list=text_list), + Task(search_graph_with_temporal_awareness, query='Who was the California Attorney General?') + ] + + pipeline = run_tasks(tasks) + + async for result in pipeline: + print(result) + + +if __name__ == '__main__': + asyncio.run(main())