File tree Expand file tree Collapse file tree 3 files changed +15
-6
lines changed
src/neo4j_graphrag/experimental/pipeline
tests/unit/experimental/pipeline Expand file tree Collapse file tree 3 files changed +15
-6
lines changed Original file line number Diff line number Diff line change 14
14
# limitations under the License.
15
15
from __future__ import annotations
16
16
17
- import datetime
18
17
from typing import Any , Optional
19
18
20
19
from neo4j_graphrag .experimental .pipeline .types import (
@@ -41,7 +40,6 @@ async def notify_pipeline_started(
41
40
event = PipelineEvent (
42
41
event_type = EventType .PIPELINE_STARTED ,
43
42
run_id = run_id ,
44
- timestamp = datetime .datetime .utcnow (),
45
43
message = None ,
46
44
payload = input_data ,
47
45
)
@@ -53,7 +51,6 @@ async def notify_pipeline_finished(
53
51
event = PipelineEvent (
54
52
event_type = EventType .PIPELINE_FINISHED ,
55
53
run_id = run_id ,
56
- timestamp = datetime .datetime .utcnow (),
57
54
message = None ,
58
55
payload = output_data ,
59
56
)
@@ -69,7 +66,6 @@ async def notify_task_started(
69
66
event_type = EventType .TASK_STARTED ,
70
67
run_id = run_id ,
71
68
task_name = task_name ,
72
- timestamp = datetime .datetime .utcnow (),
73
69
message = None ,
74
70
payload = input_data ,
75
71
)
@@ -85,7 +81,6 @@ async def notify_task_finished(
85
81
event_type = EventType .TASK_FINISHED ,
86
82
run_id = run_id ,
87
83
task_name = task_name ,
88
- timestamp = datetime .datetime .utcnow (),
89
84
message = None ,
90
85
payload = output_data .result .model_dump ()
91
86
if output_data and output_data .result
Original file line number Diff line number Diff line change @@ -82,7 +82,9 @@ class Event(BaseModel):
82
82
event_type : EventType
83
83
run_id : str
84
84
"""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
+ )
86
88
message : Optional [str ] = None
87
89
"""Optional information about the status"""
88
90
payload : Optional [dict [str , Any ]] = None
Original file line number Diff line number Diff line change 30
30
PipelineEvent ,
31
31
RunResult ,
32
32
TaskEvent ,
33
+ Event ,
33
34
)
34
35
35
36
from .components import (
@@ -479,3 +480,14 @@ async def test_pipeline_event_notification() -> None:
479
480
if previous_ts :
480
481
assert actual_event .timestamp > previous_ts
481
482
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
You can’t perform that action at this time.
0 commit comments