diff --git a/CHANGELOG.md b/CHANGELOG.md index 36dbd0f..e21b3f0 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/assistant_core/builder/datetime.py b/assistant_core/builder/datetime.py index 42329d5..bf2b039 100644 --- a/assistant_core/builder/datetime.py +++ b/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): diff --git a/tests/unit/test_datetime_builder.py b/tests/unit/test_datetime_builder.py index 1e7a3ca..9b4e9eb 100644 --- a/tests/unit/test_datetime_builder.py +++ b/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