Skip to content

Commit 3514b60

Browse files
authored
Fix datetime warning (#267)
1 parent 0b93788 commit 3514b60

File tree

3 files changed

+15
-6
lines changed

3 files changed

+15
-6
lines changed

src/neo4j_graphrag/experimental/pipeline/notification.py

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
# limitations under the License.
1515
from __future__ import annotations
1616

17-
import datetime
1817
from typing import Any, Optional
1918

2019
from neo4j_graphrag.experimental.pipeline.types import (
@@ -41,7 +40,6 @@ async def notify_pipeline_started(
4140
event = PipelineEvent(
4241
event_type=EventType.PIPELINE_STARTED,
4342
run_id=run_id,
44-
timestamp=datetime.datetime.utcnow(),
4543
message=None,
4644
payload=input_data,
4745
)
@@ -53,7 +51,6 @@ async def notify_pipeline_finished(
5351
event = PipelineEvent(
5452
event_type=EventType.PIPELINE_FINISHED,
5553
run_id=run_id,
56-
timestamp=datetime.datetime.utcnow(),
5754
message=None,
5855
payload=output_data,
5956
)
@@ -69,7 +66,6 @@ async def notify_task_started(
6966
event_type=EventType.TASK_STARTED,
7067
run_id=run_id,
7168
task_name=task_name,
72-
timestamp=datetime.datetime.utcnow(),
7369
message=None,
7470
payload=input_data,
7571
)
@@ -85,7 +81,6 @@ async def notify_task_finished(
8581
event_type=EventType.TASK_FINISHED,
8682
run_id=run_id,
8783
task_name=task_name,
88-
timestamp=datetime.datetime.utcnow(),
8984
message=None,
9085
payload=output_data.result.model_dump()
9186
if output_data and output_data.result

src/neo4j_graphrag/experimental/pipeline/types.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,9 @@ class Event(BaseModel):
8282
event_type: EventType
8383
run_id: str
8484
"""Pipeline unique run_id, same as the one returned in PipelineResult after pipeline.run"""
85-
timestamp: datetime.datetime
85+
timestamp: datetime.datetime = Field(
86+
default_factory=lambda: datetime.datetime.now(datetime.timezone.utc)
87+
)
8688
message: Optional[str] = None
8789
"""Optional information about the status"""
8890
payload: Optional[dict[str, Any]] = None

tests/unit/experimental/pipeline/test_pipeline.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
PipelineEvent,
3131
RunResult,
3232
TaskEvent,
33+
Event,
3334
)
3435

3536
from .components import (
@@ -479,3 +480,14 @@ async def test_pipeline_event_notification() -> None:
479480
if previous_ts:
480481
assert actual_event.timestamp > previous_ts
481482
previous_ts = actual_event.timestamp
483+
484+
485+
def test_event_model_no_warning(recwarn: Sized) -> None:
486+
event = Event(
487+
event_type=EventType.PIPELINE_STARTED,
488+
run_id="run_id",
489+
message=None,
490+
payload=None,
491+
)
492+
assert event.timestamp is not None
493+
assert len(recwarn) == 0

0 commit comments

Comments
 (0)