Fix datetime node

CHANGELOG.md

@@ -1,6 +1,7 @@
 # CHANGELOG
 
 ## NEXT
+* Fix datetime node to always return the current time
 
 ## v0.9.0
 * Update langgraph and langchain versions to improve security

assistant_core/builder/datetime.py

@@ -4,7 +4,7 @@
 from zoneinfo import ZoneInfo
 
 from assistant_core.builder import BaseBuilder
-from assistant_core.nodes import PromptNode
+from assistant_core.nodes import DataNode
 
 UTC = ZoneInfo("UTC")
 
@@ -27,17 +27,24 @@ def get_current_date_prompt(TZ: ZoneInfo) -> str:
     )
 
 
-class DateTimeNode(PromptNode):
+class DateTimeNode(DataNode):
     """Includes a prompt with the current date and time."""
 
-    def __init__(self, TZ: ZoneInfo = UTC, *args, **kwargs):
+    def __init__(
+        self, name: str = "date_time_node", TZ: ZoneInfo = UTC, *args, **kwargs
+    ):
         """Initialize the date and time node."""
         super().__init__(
-            name="date_time_node",
-            prompt=get_current_date_prompt(TZ),
+            name=name,
             *args,
             **kwargs,
         )
+        self.TZ = TZ
+
+    async def __call__(self, state, config):
+        """Return the prompt with the current date and time."""
+        prompt = get_current_date_prompt(self.TZ)
+        return {"messages": [self.system_message(prompt)]}
 
 
 class DateTimeBuilder(BaseBuilder):

tests/unit/test_datetime_builder.py

@@ -40,12 +40,15 @@ def test_helpers_return_expected_strings():
         assert get_current_date_prompt(UTC) == expected_prompt
 
 
-def test_datetime_node_has_prompt():
+async def test_datetime_node_has_prompt():
     with mock.patch("assistant_core.builder.datetime.datetime.datetime", FixedDateTime):
         node = DateTimeNode(TZ=UTC)
         assert isinstance(node, DateTimeNode)
-        assert "Today is Friday" in node.prompt
-        assert _make_fixed_dt().isoformat() in node.prompt
+        output = await node(state=None, config=None)
+        expected_prompt = get_current_date_prompt(UTC)
+        assert "messages" in output
+        system_message = output["messages"][0]
+        assert expected_prompt in system_message.content
 
 
 def test_datetime_builder_registers_node_and_entrypoint(builder_context):
@@ -64,3 +67,18 @@ def test_datetime_builder_registers_node_and_entrypoint(builder_context):
     # entrypoint should be set to the new node and an edge should have been created
     assert builder_context.entrypoint == "date_time_node"
     assert ("date_time_node", "test_agent") in builder_context.graph_builder.edges
+
+
+async def test_datetime_prompt_refresh():
+    current_time = datetime.datetime.now(UTC)
+    node = DateTimeNode(TZ=UTC)
+    for i in range(3):
+        test_time = current_time + datetime.timedelta(seconds=i)
+        with mock.patch(
+            "assistant_core.builder.datetime.get_current_time",
+            return_value=test_time.isoformat(),
+        ):
+            messages = await node(state=None, config=None)
+
+        system_message = messages["messages"][0]
+        assert test_time.strftime("%Y-%m-%dT%H:%M:%S") in system_message.content