From 5a43751e61ab218f340eccd533742443f2197ed0 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 27 Aug 2025 15:17:57 +0200
Subject: [PATCH] feat: adds entity and event extraction task

---
 cognee/tasks/temporal_graph/__init__.py       |  2 ++
 .../extract_events_and_entities.py            | 20 +++++++++++++++++++
 2 files changed, 22 insertions(+)
 create mode 100644 cognee/tasks/temporal_graph/__init__.py
 create mode 100644 cognee/tasks/temporal_graph/extract_events_and_entities.py

diff --git a/cognee/tasks/temporal_graph/__init__.py b/cognee/tasks/temporal_graph/__init__.py
new file mode 100644
index 000000000..163fb6840
--- /dev/null
+++ b/cognee/tasks/temporal_graph/__init__.py
@@ -0,0 +1,2 @@
+from .extract_events_and_entities import extract_events_and_entities
+
diff --git a/cognee/tasks/temporal_graph/extract_events_and_entities.py b/cognee/tasks/temporal_graph/extract_events_and_entities.py
new file mode 100644
index 000000000..37e113d56
--- /dev/null
+++ b/cognee/tasks/temporal_graph/extract_events_and_entities.py
@@ -0,0 +1,20 @@
+import asyncio
+from typing import Type, List
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.modules.chunking.models import DocumentChunk
+from cognee.tasks.temporal_graph.models import EventList
+from cognee.modules.engine.utils.generate_event_datapoint import generate_event_datapoint
+
+
+async def extract_events_and_entities(data_chunks: List[DocumentChunk]) -> List[DocumentChunk]:
+    """Extracts events and entities from a chunk of documents."""
+    events = await asyncio.gather(
+        *[LLMGateway.extract_event_graph(chunk.text, EventList) for chunk in data_chunks]
+    )
+
+    for data_chunk, event_list in zip(data_chunks, events):
+        for event in event_list.events:
+            event_datapoint = generate_event_datapoint(event)
+            data_chunk.contains.append(event_datapoint)
+
+    return data_chunks
\ No newline at end of file