From dd94781033b79d72dc7613fb5f3ac313da45776a Mon Sep 17 00:00:00 2001 From: Rita Aleksziev Date: Wed, 4 Dec 2024 16:33:26 +0100 Subject: [PATCH 1/2] Integrate graphiti's functionality as Tasks --- cognee/tasks/temporal_awareness/__init__.py | 4 +++ .../build_graph_with_temporal_awareness.py | 22 ++++++++++++++ .../search_graph_with_temporal_awareness.py | 6 ++++ examples/python/graphiti_example.py | 29 +++++++++++++++++++ 4 files changed, 61 insertions(+) create mode 100644 cognee/tasks/temporal_awareness/__init__.py create mode 100644 cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py create mode 100644 cognee/tasks/temporal_awareness/search_graph_with_temporal_awareness.py create mode 100644 examples/python/graphiti_example.py 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()) From 5d71059f6c7c070585d8a7c089fe6d49d433fb06 Mon Sep 17 00:00:00 2001 From: Rita Aleksziev Date: Wed, 4 Dec 2024 17:00:48 +0100 Subject: [PATCH 2/2] fix credentials --- .../build_graph_with_temporal_awareness.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py b/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py index 4a7df6db0..ac534623b 100644 --- a/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py +++ b/cognee/tasks/temporal_awareness/build_graph_with_temporal_awareness.py @@ -1,3 +1,4 @@ +import os from datetime import datetime from graphiti_core import Graphiti @@ -6,7 +7,10 @@ async def build_graph_with_temporal_awareness(text_list): - graphiti = Graphiti("bolt://localhost:7687", "neo4j", "pleaseletmein") + url = os.getenv("GRAPH_DATABASE_URL") + password = os.getenv("GRAPH_DATABASE_PASSWORD") + graphiti = Graphiti(url, "neo4j", password) + await graphiti.build_indices_and_constraints() print("Graph database initialized.")