feat: time extension for ChatPromptBuilder (#9001)

* feat: time extension for ChatPromptBuilder

* chore: release notes

* Fix comment

---------

Co-authored-by: Sebastian Husch Lee <sjrl423@gmail.com>

haystack/components/builders/chat_prompt_builder.py

@@ -10,6 +10,7 @@ from jinja2.sandbox import SandboxedEnvironment
 
 from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.dataclasses.chat_message import ChatMessage, ChatRole, TextContent
+from haystack.utils import Jinja2TimeExtension
 
 logger = logging.getLogger(__name__)
 
@@ -124,7 +125,13 @@ class ChatPromptBuilder:
         self.required_variables = required_variables or []
         self.template = template
         variables = variables or []
-        self._env = SandboxedEnvironment()
+        try:
+            # The Jinja2TimeExtension needs an optional dependency to be installed.
+            # If it's not available we can do without it and use the ChatPromptBuilder as is.
+            self._env = SandboxedEnvironment(extensions=[Jinja2TimeExtension])
+        except ImportError:
+            self._env = SandboxedEnvironment()
+
         if template and not variables:
             for message in template:
                 if message.is_from(ChatRole.USER) or message.is_from(ChatRole.SYSTEM):

releasenotes/notes/feat-time-extension-chatprompt-builder-fb159207088943d8.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Users can now work with date and time in the ChatPromptBuilder.
+    In the same way as the PromptBuilder, the ChatPromptBuilder now supports arrow to work with datetime.

test/components/builders/test_chat_prompt_builder.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict, List, Optional
 from jinja2 import TemplateSyntaxError
+import arrow
 import logging
 import pytest
 
@@ -346,6 +347,87 @@ class TestChatPromptBuilder:
             )
             assert "ChatPromptBuilder has 2 prompt variables, but `required_variables` is not set. " in caplog.text
 
+    def test_with_custom_dateformat(self) -> None:
+        template = [ChatMessage.from_user("Formatted date: {% now 'UTC', '%Y-%m-%d' %}")]
+        builder = ChatPromptBuilder(template=template)
+
+        result = builder.run()["prompt"]
+
+        now_formatted = f"Formatted date: {arrow.now('UTC').strftime('%Y-%m-%d')}"
+        assert len(result) == 1
+        assert result[0].role == "user"
+        assert result[0].text == now_formatted
+
+    def test_with_different_timezone(self) -> None:
+        template = [ChatMessage.from_user("Current time in New York is: {% now 'America/New_York' %}")]
+        builder = ChatPromptBuilder(template=template)
+
+        result = builder.run()["prompt"]
+
+        now_ny = f"Current time in New York is: {arrow.now('America/New_York').strftime('%Y-%m-%d %H:%M:%S')}"
+        assert len(result) == 1
+        assert result[0].role == "user"
+        assert result[0].text == now_ny
+
+    def test_date_with_addition_offset(self) -> None:
+        template = [ChatMessage.from_user("Time after 2 hours is: {% now 'UTC' + 'hours=2' %}")]
+        builder = ChatPromptBuilder(template=template)
+
+        result = builder.run()["prompt"]
+
+        now_plus_2 = f"Time after 2 hours is: {(arrow.now('UTC').shift(hours=+2)).strftime('%Y-%m-%d %H:%M:%S')}"
+        assert len(result) == 1
+        assert result[0].role == "user"
+        assert result[0].text == now_plus_2
+
+    def test_date_with_subtraction_offset(self) -> None:
+        template = [ChatMessage.from_user("Time after 12 days is: {% now 'UTC' - 'days=12' %}")]
+        builder = ChatPromptBuilder(template=template)
+
+        result = builder.run()["prompt"]
+
+        now_minus_12 = f"Time after 12 days is: {(arrow.now('UTC').shift(days=-12)).strftime('%Y-%m-%d %H:%M:%S')}"
+        assert len(result) == 1
+        assert result[0].role == "user"
+        assert result[0].text == now_minus_12
+
+    def test_invalid_timezone(self) -> None:
+        template = [ChatMessage.from_user("Current time is: {% now 'Invalid/Timezone' %}")]
+        builder = ChatPromptBuilder(template=template)
+
+        # Expect ValueError for invalid timezone
+        with pytest.raises(ValueError, match="Invalid timezone"):
+            builder.run()
+
+    def test_invalid_offset(self) -> None:
+        template = [ChatMessage.from_user("Time after invalid offset is: {% now 'UTC' + 'invalid_offset' %}")]
+        builder = ChatPromptBuilder(template=template)
+
+        # Expect ValueError for invalid offset
+        with pytest.raises(ValueError, match="Invalid offset or operator"):
+            builder.run()
+
+    def test_multiple_messages_with_date_template(self) -> None:
+        template = [
+            ChatMessage.from_user("Current date is: {% now 'UTC' %}"),
+            ChatMessage.from_assistant("Thank you for providing the date"),
+            ChatMessage.from_user("Yesterday was: {% now 'UTC' - 'days=1' %}"),
+        ]
+        builder = ChatPromptBuilder(template=template)
+
+        result = builder.run()["prompt"]
+
+        now = f"Current date is: {arrow.now('UTC').strftime('%Y-%m-%d %H:%M:%S')}"
+        yesterday = f"Yesterday was: {(arrow.now('UTC').shift(days=-1)).strftime('%Y-%m-%d %H:%M:%S')}"
+
+        assert len(result) == 3
+        assert result[0].role == "user"
+        assert result[0].text == now
+        assert result[1].role == "assistant"
+        assert result[1].text == "Thank you for providing the date"
+        assert result[2].role == "user"
+        assert result[2].text == yesterday
+
 
 class TestChatPromptBuilderDynamic:
     def test_multiple_templated_chat_messages(self):