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