From bf34ba398e1d3dd39373a0e3b86f0c90e54ef8f7 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 27 Aug 2025 15:14:46 +0200
Subject: [PATCH] feat: adds temporal models for llm extraction

---
 cognee/modules/chunking/models/DocumentChunk.py |  7 ++++---
 cognee/modules/engine/models/Event.py           | 16 ++++++++++++++++
 cognee/modules/engine/models/Interval.py        |  7 +++++++
 cognee/modules/engine/models/Timestamp.py       | 13 +++++++++++++
 cognee/modules/engine/models/__init__.py        |  3 +++
 5 files changed, 43 insertions(+), 3 deletions(-)
 create mode 100644 cognee/modules/engine/models/Event.py
 create mode 100644 cognee/modules/engine/models/Interval.py
 create mode 100644 cognee/modules/engine/models/Timestamp.py

diff --git a/cognee/modules/chunking/models/DocumentChunk.py b/cognee/modules/chunking/models/DocumentChunk.py
index 475703265..9f8c57486 100644
--- a/cognee/modules/chunking/models/DocumentChunk.py
+++ b/cognee/modules/chunking/models/DocumentChunk.py
@@ -1,8 +1,9 @@
-from typing import List
+from typing import List, Union
 
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.data.processing.document_types import Document
 from cognee.modules.engine.models import Entity
+from cognee.tasks.temporal_graph.models import Event
 
 
 class DocumentChunk(DataPoint):
@@ -20,7 +21,7 @@ class DocumentChunk(DataPoint):
     - chunk_index: The index of the chunk in the original document.
     - cut_type: The type of cut that defined this chunk.
     - is_part_of: The document to which this chunk belongs.
-    - contains: A list of entities contained within the chunk (default is None).
+    - contains: A list of entities or events contained within the chunk (default is None).
     - metadata: A dictionary to hold meta information related to the chunk, including index
     fields.
     """
@@ -30,6 +31,6 @@ class DocumentChunk(DataPoint):
     chunk_index: int
     cut_type: str
     is_part_of: Document
-    contains: List[Entity] = None
+    contains: List[Union[Entity, Event]] = None
 
     metadata: dict = {"index_fields": ["text"]}
diff --git a/cognee/modules/engine/models/Event.py b/cognee/modules/engine/models/Event.py
new file mode 100644
index 000000000..88141e602
--- /dev/null
+++ b/cognee/modules/engine/models/Event.py
@@ -0,0 +1,16 @@
+from typing import Optional, Any
+from pydantic import SkipValidation
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.engine.models.Timestamp import Timestamp
+from cognee.modules.engine.models.Interval import Interval
+
+
+class Event(DataPoint):
+    name: str
+    description: Optional[str] = None
+    at: Optional[Timestamp] = None
+    during: Optional[Interval] = None
+    location: Optional[str] = None
+    attributes: SkipValidation[Any] = None
+
+    metadata: dict = {"index_fields": ["name"]}
\ No newline at end of file
diff --git a/cognee/modules/engine/models/Interval.py b/cognee/modules/engine/models/Interval.py
new file mode 100644
index 000000000..3666bf69d
--- /dev/null
+++ b/cognee/modules/engine/models/Interval.py
@@ -0,0 +1,7 @@
+from pydantic import Field
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.engine.models.Timestamp import Timestamp
+
+class Interval(DataPoint):
+    time_from: Timestamp = Field(...)
+    time_to: Timestamp = Field(...)
\ No newline at end of file
diff --git a/cognee/modules/engine/models/Timestamp.py b/cognee/modules/engine/models/Timestamp.py
new file mode 100644
index 000000000..38977c348
--- /dev/null
+++ b/cognee/modules/engine/models/Timestamp.py
@@ -0,0 +1,13 @@
+from pydantic import Field
+from cognee.infrastructure.engine import DataPoint
+
+
+class Timestamp(DataPoint):
+    time_at: int = Field(...)
+    year: int = Field(...)
+    month: int = Field(...)
+    day: int = Field(...)
+    hour: int = Field(...)
+    minute: int = Field(...)
+    second: int = Field(...)
+    timestamp_str: str = Field(...)
\ No newline at end of file
diff --git a/cognee/modules/engine/models/__init__.py b/cognee/modules/engine/models/__init__.py
index 2535f00f3..8d28ebf8a 100644
--- a/cognee/modules/engine/models/__init__.py
+++ b/cognee/modules/engine/models/__init__.py
@@ -4,3 +4,6 @@ from .TableRow import TableRow
 from .TableType import TableType
 from .node_set import NodeSet
 from .ColumnValue import ColumnValue
+from .Timestamp import Timestamp
+from .Interval import Interval
+from .Event import Event