diff --git a/haystack/components/builders/chat_prompt_builder.py b/haystack/components/builders/chat_prompt_builder.py index a097c9319..21a90f785 100644 --- a/haystack/components/builders/chat_prompt_builder.py +++ b/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): diff --git a/releasenotes/notes/feat-time-extension-chatprompt-builder-fb159207088943d8.yaml b/releasenotes/notes/feat-time-extension-chatprompt-builder-fb159207088943d8.yaml new file mode 100644 index 000000000..f1c707c0d --- /dev/null +++ b/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. diff --git a/test/components/builders/test_chat_prompt_builder.py b/test/components/builders/test_chat_prompt_builder.py index fcee86b17..838766918 100644 --- a/test/components/builders/test_chat_prompt_builder.py +++ b/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):