Integrate graphiti's functionality as Tasks

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

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

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

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())