haystack/test/components/builders/test_dynamic_chat_prompt_builder.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
from typing import List

import pytest

from haystack import Pipeline, component
from haystack.components.builders import DynamicChatPromptBuilder
from haystack.dataclasses import ChatMessage


class TestDynamicChatPromptBuilder:
    def test_initialization(self):
        runtime_variables = ["var1", "var2", "var3"]
        builder = DynamicChatPromptBuilder(runtime_variables)
        assert builder.runtime_variables == runtime_variables

        # we have inputs that contain: prompt_source, template_variables + runtime_variables
        expected_keys = set(runtime_variables + ["prompt_source", "template_variables"])
        assert set(builder.__haystack_input__._sockets_dict.keys()) == expected_keys

        # response is always prompt regardless of chat mode
        assert set(builder.__haystack_output__._sockets_dict.keys()) == {"prompt"}

        # prompt_source is a list of ChatMessage
        assert builder.__haystack_input__._sockets_dict["prompt_source"].type == List[ChatMessage]

        # output is always prompt, but the type is different depending on the chat mode
        assert builder.__haystack_output__._sockets_dict["prompt"].type == List[ChatMessage]

    def test_non_empty_chat_messages(self):
        prompt_builder = DynamicChatPromptBuilder(runtime_variables=["documents"])
        prompt_source = [ChatMessage.from_system(content="Hello"), ChatMessage.from_user(content="Hello, {{ who }}!")]
        template_variables = {"who": "World"}

        result = prompt_builder.run(prompt_source, template_variables)

        assert result == {
            "prompt": [ChatMessage.from_system(content="Hello"), ChatMessage.from_user(content="Hello, World!")]
        }

    def test_single_chat_message(self):
        prompt_builder = DynamicChatPromptBuilder(runtime_variables=["documents"])
        prompt_source = [ChatMessage.from_user(content="Hello, {{ who }}!")]
        template_variables = {"who": "World"}

        result = prompt_builder.run(prompt_source, template_variables)

        assert result == {"prompt": [ChatMessage.from_user(content="Hello, World!")]}

    def test_empty_chat_message_list(self):
        prompt_builder = DynamicChatPromptBuilder(runtime_variables=["documents"])

        with pytest.raises(ValueError):
            prompt_builder.run(prompt_source=[], template_variables={})

    def test_chat_message_list_with_mixed_object_list(self):
        prompt_builder = DynamicChatPromptBuilder(runtime_variables=["documents"])

        with pytest.raises(ValueError):
            prompt_builder.run(prompt_source=[ChatMessage.from_user("Hello"), "there world"], template_variables={})

    def test_chat_message_list_with_missing_variables(self):
        prompt_builder = DynamicChatPromptBuilder(runtime_variables=["documents"])
        prompt_source = [ChatMessage.from_user(content="Hello, {{ who }}!")]

        # Call the _process_chat_messages method and expect a ValueError
        with pytest.raises(ValueError):
            prompt_builder.run(prompt_source, template_variables={})

    def test_missing_template_variables(self):
        prompt_builder = DynamicChatPromptBuilder(runtime_variables=["documents"])

        # missing template variable city
        with pytest.raises(ValueError):
            prompt_builder._validate_template("Hello, I'm {{ name }}, and I live in {{ city }}.", {"name"})

        # missing template variable name
        with pytest.raises(ValueError):
            prompt_builder._validate_template("Hello, I'm {{ name }}, and I live in {{ city }}.", {"city"})

        # completely unknown template variable
        with pytest.raises(ValueError):
            prompt_builder._validate_template("Hello, I'm {{ name }}, and I live in {{ city }}.", {"age"})

    def test_provided_template_variables(self):
        prompt_builder = DynamicChatPromptBuilder(runtime_variables=["documents"])

        # both variables are provided
        prompt_builder._validate_template("Hello, I'm {{ name }}, and I live in {{ city }}.", {"name", "city"})

        # provided variables are a superset of the required variables
        prompt_builder._validate_template("Hello, I'm {{ name }}, and I live in {{ city }}.", {"name", "city", "age"})

    def test_multiple_templated_chat_messages(self):
        prompt_builder = DynamicChatPromptBuilder()
        language = "French"
        location = "Berlin"
        messages = [
            ChatMessage.from_system("Write your response in this language:{{language}}"),
            ChatMessage.from_user("Tell me about {{location}}"),
        ]

        result = prompt_builder.run(
            template_variables={"language": language, "location": location}, prompt_source=messages
        )
        assert result["prompt"] == [
            ChatMessage.from_system("Write your response in this language:French"),
            ChatMessage.from_user("Tell me about Berlin"),
        ], "The templated messages should match the expected output."

    def test_multiple_templated_chat_messages_in_place(self):
        prompt_builder = DynamicChatPromptBuilder()
        language = "French"
        location = "Berlin"
        messages = [
            ChatMessage.from_system("Write your response ins this language:{{language}}"),
            ChatMessage.from_user("Tell me about {{location}}"),
        ]

        res = prompt_builder.run(
            template_variables={"language": language, "location": location}, prompt_source=messages
        )
        assert res == {
            "prompt": [
                ChatMessage.from_system("Write your response ins this language:French"),
                ChatMessage.from_user("Tell me about Berlin"),
            ]
        }, "The templated messages should match the expected output."

    def test_some_templated_chat_messages(self):
        prompt_builder = DynamicChatPromptBuilder()
        language = "English"
        location = "Paris"
        messages = [
            ChatMessage.from_system("Please, respond in the following language: {{language}}."),
            ChatMessage.from_user("I would like to learn more about {{location}}."),
            ChatMessage.from_assistant("Yes, I can help you with that {{subject}}"),
            ChatMessage.from_user("Ok so do so please, be elaborate."),
        ]

        result = prompt_builder.run(
            template_variables={"language": language, "location": location}, prompt_source=messages
        )

        expected_messages = [
            ChatMessage.from_system("Please, respond in the following language: English."),
            ChatMessage.from_user("I would like to learn more about Paris."),
            ChatMessage.from_assistant(
                "Yes, I can help you with that {{subject}}"
            ),  # assistant message should not be templated
            ChatMessage.from_user("Ok so do so please, be elaborate."),
        ]

        assert result["prompt"] == expected_messages, "The templated messages should match the expected output."

    def test_example_in_pipeline(self):
        prompt_builder = DynamicChatPromptBuilder()

        pipe = Pipeline()
        pipe.add_component("prompt_builder", prompt_builder)

        location = "Berlin"
        system_message = ChatMessage.from_system(
            "You are a helpful assistant giving out valuable information to tourists."
        )
        messages = [system_message, ChatMessage.from_user("Tell me about {{location}}")]

        res = pipe.run(
            data={"prompt_builder": {"template_variables": {"location": location}, "prompt_source": messages}}
        )
        assert res == {
            "prompt_builder": {
                "prompt": [
                    ChatMessage.from_system("You are a helpful assistant giving out valuable information to tourists."),
                    ChatMessage.from_user("Tell me about Berlin"),
                ]
            }
        }

        messages = [
            system_message,
            ChatMessage.from_user("What's the weather forecast for {{location}} in the next {{day_count}} days?"),
        ]

        res = pipe.run(
            data={
                "prompt_builder": {
                    "template_variables": {"location": location, "day_count": "5"},
                    "prompt_source": messages,
                }
            }
        )
        assert res == {
            "prompt_builder": {
                "prompt": [
                    ChatMessage.from_system("You are a helpful assistant giving out valuable information to tourists."),
                    ChatMessage.from_user("What's the weather forecast for Berlin in the next 5 days?"),
                ]
            }
        }

    def test_example_in_pipeline_with_multiple_templated_messages(self):
        # no parameter init, we don't use any runtime template variables
        prompt_builder = DynamicChatPromptBuilder()

        pipe = Pipeline()
        pipe.add_component("prompt_builder", prompt_builder)

        location = "Berlin"
        system_message = ChatMessage.from_system(
            "You are a helpful assistant giving out valuable information to tourists in {{language}}."
        )
        messages = [system_message, ChatMessage.from_user("Tell me about {{location}}")]

        res = pipe.run(
            data={
                "prompt_builder": {
                    "template_variables": {"location": location, "language": "German"},
                    "prompt_source": messages,
                }
            }
        )
        assert res == {
            "prompt_builder": {
                "prompt": [
                    ChatMessage.from_system(
                        "You are a helpful assistant giving out valuable information to tourists in German."
                    ),
                    ChatMessage.from_user("Tell me about Berlin"),
                ]
            }
        }

        messages = [
            system_message,
            ChatMessage.from_user("What's the weather forecast for {{location}} in the next {{day_count}} days?"),
        ]

        res = pipe.run(
            data={
                "prompt_builder": {
                    "template_variables": {"location": location, "day_count": "5", "language": "English"},
                    "prompt_source": messages,
                }
            }
        )
        assert res == {
            "prompt_builder": {
                "prompt": [
                    ChatMessage.from_system(
                        "You are a helpful assistant giving out valuable information to tourists in English."
                    ),
                    ChatMessage.from_user("What's the weather forecast for Berlin in the next 5 days?"),
                ]
            }
        }

    def test_pipeline_complex(self):
        @component
        class ValueProducer:
            def __init__(self, value_to_produce: str):
                self.value_to_produce = value_to_produce

            @component.output_types(value_output=str)
            def run(self):
                return {"value_output": self.value_to_produce}

        pipe = Pipeline()
        pipe.add_component("prompt_builder", DynamicChatPromptBuilder(runtime_variables=["value_output"]))
        pipe.add_component("value_producer", ValueProducer(value_to_produce="Berlin"))
        pipe.connect("value_producer.value_output", "prompt_builder")

        messages = [
            ChatMessage.from_system("You give valuable information to tourists."),
            ChatMessage.from_user("Tell me about {{value_output}}"),
        ]

        res = pipe.run(data={"prompt_source": messages})
        assert res == {
            "prompt_builder": {
                "prompt": [
                    ChatMessage.from_system("You give valuable information to tourists."),
                    ChatMessage.from_user("Tell me about Berlin"),
                ]
            }
        }