10112264105

Committed 26 Jul 2024 01:40PM UTC coverage: 90.045% (-0.001%) from 90.046%

Build # 10112264105

Build Type

Pull #8095

github

Committed by

web-flow

Commit Message

Merge e16cefc3a into 47f4db869

Pull Request Pull Request #8095: fix: Fix issue that could lead to RCE if using unsecure Jinja templates

Run Details

6793 of 7544 relevant lines covered (90.05%)

0.9 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

98.44

haystack/components/builders/chat_prompt_builder.py

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

from typing import Any, Dict, List, Optional, Set

from jinja2 import meta
from jinja2.sandbox import SandboxedEnvironment

from haystack import component, default_from_dict, default_to_dict, logging
from haystack.dataclasses.chat_message import ChatMessage, ChatRole

logger = logging.getLogger(__name__)


@component
class ChatPromptBuilder:
    """
    ChatPromptBuilder is a component that renders a chat prompt from a template string using Jinja2 templates.

    It is designed to construct prompts for the pipeline using static or dynamic templates: Users can change
    the prompt template at runtime by providing a new template for each pipeline run invocation if needed.

    The template variables found in the init template string are used as input types for the component and are all
    optional, unless explicitly specified. If an optional template variable is not provided as an input, it will be
    replaced with an empty string in the rendered prompt. Use `variable` and `required_variables` to specify the input
    types and required variables.

    Usage example with static prompt template:
    ```python
    template = [ChatMessage.from_user("Translate to {{ target_language }}. Context: {{ snippet }}; Translation:")]
    builder = ChatPromptBuilder(template=template)
    builder.run(target_language="spanish", snippet="I can't speak spanish.")
    ```

    Usage example of overriding the static template at runtime:
    ```python
    template = [ChatMessage.from_user("Translate to {{ target_language }}. Context: {{ snippet }}; Translation:")]
    builder = ChatPromptBuilder(template=template)
    builder.run(target_language="spanish", snippet="I can't speak spanish.")

    msg = "Translate to {{ target_language }} and summarize. Context: {{ snippet }}; Summary:"
    summary_template = [ChatMessage.from_user(msg)]
    builder.run(target_language="spanish", snippet="I can't speak spanish.", template=summary_template)
    ```

    Usage example with dynamic prompt template:
    ```python
    from haystack.components.builders import ChatPromptBuilder
    from haystack.components.generators.chat import OpenAIChatGenerator
    from haystack.dataclasses import ChatMessage
    from haystack import Pipeline
    from haystack.utils import Secret

    # no parameter init, we don't use any runtime template variables
    prompt_builder = ChatPromptBuilder()
    llm = OpenAIChatGenerator(api_key=Secret.from_token("<your-api-key>"), model="gpt-3.5-turbo")

    pipe = Pipeline()
    pipe.add_component("prompt_builder", prompt_builder)
    pipe.add_component("llm", llm)
    pipe.connect("prompt_builder.prompt", "llm.messages")

    location = "Berlin"
    language = "English"
    system_message = ChatMessage.from_system("You are an assistant giving 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": language},
                                        "template": messages}})
    print(res)

    >> {'llm': {'replies': [ChatMessage(content="Berlin is the capital city of Germany and one of the most vibrant
    and diverse cities in Europe. Here are some key things to know...Enjoy your time exploring the vibrant and dynamic
    capital of Germany!", role=<ChatRole.ASSISTANT: 'assistant'>, name=None, meta={'model': 'gpt-3.5-turbo-0613',
    'index': 0, 'finish_reason': 'stop', 'usage': {'prompt_tokens': 27, 'completion_tokens': 681, 'total_tokens':
    708}})]}}


    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"},
                                        "template": messages}})

    print(res)
    >> {'llm': {'replies': [ChatMessage(content="Here is the weather forecast for Berlin in the next 5
    days:\n\nDay 1: Mostly cloudy with a high of 22°C (72°F) and...so it's always a good idea to check for updates
    closer to your visit.", role=<ChatRole.ASSISTANT: 'assistant'>, name=None, meta={'model': 'gpt-3.5-turbo-0613',
    'index': 0, 'finish_reason': 'stop', 'usage': {'prompt_tokens': 37, 'completion_tokens': 201,
    'total_tokens': 238}})]}}
    ```

    Note how in the example above, we can dynamically change the prompt template by providing a new template to the
    run method of the pipeline.

    """

    def __init__(
        self,
        template: Optional[List[ChatMessage]] = None,
        required_variables: Optional[List[str]] = None,
        variables: Optional[List[str]] = None,
    ):
        """
        Constructs a ChatPromptBuilder component.

        :param template:
            A list of `ChatMessage` instances. All user and system messages are treated as potentially having jinja2
            templates and are rendered with the provided template variables. If not provided, the template
            must be provided at runtime using the `template` parameter of the `run` method.
        :param required_variables: An optional list of input variables that must be provided at all times.
            If not provided, an exception will be raised.
        :param variables:
            A list of template variable names you can use in prompt construction. For example,
            if `variables` contains the string `documents`, the component will create an input called
            `documents` of type `Any`. These variable names are used to resolve variables and their values during
            pipeline execution. The values associated with variables from the pipeline runtime are then injected into
            template placeholders of a prompt text template that is provided to the `run` method.
            If not provided, variables are inferred from `template`.
        """
        self._variables = variables
        self._required_variables = required_variables
        self.required_variables = required_variables or []
        self.template = template
        variables = variables or []
        self._env = SandboxedEnvironment()
        if template and not variables:
            for message in template:
                if message.is_from(ChatRole.USER) or message.is_from(ChatRole.SYSTEM):
                    # infere variables from template
                    ast = self._env.parse(message.content)
                    template_variables = meta.find_undeclared_variables(ast)
                    variables += list(template_variables)

        # setup inputs
        static_input_slots = {"template": Optional[str], "template_variables": Optional[Dict[str, Any]]}
        component.set_input_types(self, **static_input_slots)
        for var in variables:
            if var in self.required_variables:
                component.set_input_type(self, var, Any)
            else:
                component.set_input_type(self, var, Any, "")

    @component.output_types(prompt=List[ChatMessage])
    def run(
        self,
        template: Optional[List[ChatMessage]] = None,
        template_variables: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        """
        Executes the prompt building process.

        It applies the template variables to render the final prompt. You can provide variables either via pipeline
        (set through `variables` or inferred from `template` at initialization) or via additional template variables
        set directly to this method. On collision, the variables provided directly to this method take precedence.

        :param template:
            An optional list of ChatMessages to overwrite ChatPromptBuilder's default template. If None, the default
            template provided at initialization is used.
        :param template_variables:
            An optional dictionary of template variables. These are additional variables users can provide directly
            to this method in contrast to pipeline variables.
        :param kwargs:
            Pipeline variables (typically resolved from a pipeline) which are merged with the provided template
            variables.

        :returns: A dictionary with the following keys:
            - `prompt`: The updated list of `ChatMessage` instances after rendering the found templates.
        :raises ValueError:
            If `chat_messages` is empty or contains elements that are not instances of `ChatMessage`.
        """
        kwargs = kwargs or {}
        template_variables = template_variables or {}
        template_variables_combined = {**kwargs, **template_variables}

        if template is None:
            template = self.template

        if not template:
            raise ValueError(
                f"The {self.__class__.__name__} requires a non-empty list of ChatMessage instances. "
                f"Please provide a valid list of ChatMessage instances to render the prompt."
            )

        if not all(isinstance(message, ChatMessage) for message in template):
            raise ValueError(
                f"The {self.__class__.__name__} expects a list containing only ChatMessage instances. "
                f"The provided list contains other types. Please ensure that all elements in the list "
                f"are ChatMessage instances."
            )

        processed_messages = []
        for message in template:
            if message.is_from(ChatRole.USER) or message.is_from(ChatRole.SYSTEM):
                self._validate_variables(set(template_variables_combined.keys()))

                compiled_template = self._env.from_string(message.content)
                rendered_content = compiled_template.render(template_variables_combined)
                rendered_message = (
                    ChatMessage.from_user(rendered_content)
                    if message.is_from(ChatRole.USER)
                    else ChatMessage.from_system(rendered_content)
                )
                processed_messages.append(rendered_message)
            else:
                processed_messages.append(message)

        return {"prompt": processed_messages}

    def _validate_variables(self, provided_variables: Set[str]):
        """
        Checks if all the required template variables are provided.

        :param provided_variables:
            A set of provided template variables.
        :raises ValueError:
            If no template is provided or if all the required template variables are not provided.
        """
        missing_variables = [var for var in self.required_variables if var not in provided_variables]
        if missing_variables:
            missing_vars_str = ", ".join(missing_variables)
            raise ValueError(
                f"Missing required input variables in ChatPromptBuilder: {missing_vars_str}. "
                f"Required variables: {self.required_variables}. Provided variables: {provided_variables}."
            )

    def to_dict(self) -> Dict[str, Any]:
        """
        Returns a dictionary representation of the component.

        :returns:
            Serialized dictionary representation of the component.
        """
        if self.template is not None:
            template = [m.to_dict() for m in self.template]
        else:
            template = None

        return default_to_dict(
            self, template=template, variables=self._variables, required_variables=self._required_variables
        )

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ChatPromptBuilder":
        """
        Deserialize this component from a dictionary.

        :param data:
            The dictionary to deserialize and create the component.

        :returns:
            The deserialized component.
        """
        init_parameters = data["init_parameters"]
        template = init_parameters.get("template", [])
        init_parameters["template"] = [ChatMessage.from_dict(d) for d in template]

        return default_from_dict(cls, data)

1	# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	from typing import Any, Dict, List, Optional, Set	1✔
6
7	from jinja2 import meta	1✔
8	from jinja2.sandbox import SandboxedEnvironment	1✔
9
10	from haystack import component, default_from_dict, default_to_dict, logging	1✔
11	from haystack.dataclasses.chat_message import ChatMessage, ChatRole	1✔
12
13	logger = logging.getLogger(__name__)	1✔
14
15
16	@component	1✔
17	class ChatPromptBuilder:	1✔
18	"""
19	ChatPromptBuilder is a component that renders a chat prompt from a template string using Jinja2 templates.
20
21	It is designed to construct prompts for the pipeline using static or dynamic templates: Users can change
22	the prompt template at runtime by providing a new template for each pipeline run invocation if needed.
23
24	The template variables found in the init template string are used as input types for the component and are all
25	optional, unless explicitly specified. If an optional template variable is not provided as an input, it will be
26	replaced with an empty string in the rendered prompt. Use `variable` and `required_variables` to specify the input
27	types and required variables.
28
29	Usage example with static prompt template:
30	```python
31	template = [ChatMessage.from_user("Translate to {{ target_language }}. Context: {{ snippet }}; Translation:")]
32	builder = ChatPromptBuilder(template=template)
33	builder.run(target_language="spanish", snippet="I can't speak spanish.")
34	```
35
36	Usage example of overriding the static template at runtime:
37	```python
38	template = [ChatMessage.from_user("Translate to {{ target_language }}. Context: {{ snippet }}; Translation:")]
39	builder = ChatPromptBuilder(template=template)
40	builder.run(target_language="spanish", snippet="I can't speak spanish.")
41
42	msg = "Translate to {{ target_language }} and summarize. Context: {{ snippet }}; Summary:"
43	summary_template = [ChatMessage.from_user(msg)]
44	builder.run(target_language="spanish", snippet="I can't speak spanish.", template=summary_template)
45	```
46
47	Usage example with dynamic prompt template:
48	```python
49	from haystack.components.builders import ChatPromptBuilder
50	from haystack.components.generators.chat import OpenAIChatGenerator
51	from haystack.dataclasses import ChatMessage
52	from haystack import Pipeline
53	from haystack.utils import Secret
54
55	# no parameter init, we don't use any runtime template variables
56	prompt_builder = ChatPromptBuilder()
57	llm = OpenAIChatGenerator(api_key=Secret.from_token("<your-api-key>"), model="gpt-3.5-turbo")
58
59	pipe = Pipeline()
60	pipe.add_component("prompt_builder", prompt_builder)
61	pipe.add_component("llm", llm)
62	pipe.connect("prompt_builder.prompt", "llm.messages")
63
64	location = "Berlin"
65	language = "English"
66	system_message = ChatMessage.from_system("You are an assistant giving information to tourists in {{language}}")
67	messages = [system_message, ChatMessage.from_user("Tell me about {{location}}")]
68
69	res = pipe.run(data={"prompt_builder": {"template_variables": {"location": location, "language": language},
70	"template": messages}})
71	print(res)
72
73	>> {'llm': {'replies': [ChatMessage(content="Berlin is the capital city of Germany and one of the most vibrant
74	and diverse cities in Europe. Here are some key things to know...Enjoy your time exploring the vibrant and dynamic
75	capital of Germany!", role=<ChatRole.ASSISTANT: 'assistant'>, name=None, meta={'model': 'gpt-3.5-turbo-0613',
76	'index': 0, 'finish_reason': 'stop', 'usage': {'prompt_tokens': 27, 'completion_tokens': 681, 'total_tokens':
77	708}})]}}
78
79
80	messages = [system_message, ChatMessage.from_user("What's the weather forecast for {{location}} in the next
81	{{day_count}} days?")]
82
83	res = pipe.run(data={"prompt_builder": {"template_variables": {"location": location, "day_count": "5"},
84	"template": messages}})
85
86	print(res)
87	>> {'llm': {'replies': [ChatMessage(content="Here is the weather forecast for Berlin in the next 5
88	days:\n\nDay 1: Mostly cloudy with a high of 22°C (72°F) and...so it's always a good idea to check for updates
89	closer to your visit.", role=<ChatRole.ASSISTANT: 'assistant'>, name=None, meta={'model': 'gpt-3.5-turbo-0613',
90	'index': 0, 'finish_reason': 'stop', 'usage': {'prompt_tokens': 37, 'completion_tokens': 201,
91	'total_tokens': 238}})]}}
92	```
93
94	Note how in the example above, we can dynamically change the prompt template by providing a new template to the
95	run method of the pipeline.
96
97	"""
98
99	def __init__(	1✔
100	self,
101	template: Optional[List[ChatMessage]] = None,
102	required_variables: Optional[List[str]] = None,
103	variables: Optional[List[str]] = None,
104	):
105	"""
106	Constructs a ChatPromptBuilder component.
107
108	:param template:
109	A list of `ChatMessage` instances. All user and system messages are treated as potentially having jinja2
110	templates and are rendered with the provided template variables. If not provided, the template
111	must be provided at runtime using the `template` parameter of the `run` method.
112	:param required_variables: An optional list of input variables that must be provided at all times.
113	If not provided, an exception will be raised.
114	:param variables:
115	A list of template variable names you can use in prompt construction. For example,
116	if `variables` contains the string `documents`, the component will create an input called
117	`documents` of type `Any`. These variable names are used to resolve variables and their values during
118	pipeline execution. The values associated with variables from the pipeline runtime are then injected into
119	template placeholders of a prompt text template that is provided to the `run` method.
120	If not provided, variables are inferred from `template`.
121	"""
122	self._variables = variables	1✔
123	self._required_variables = required_variables	1✔
124	self.required_variables = required_variables or []	1✔
125	self.template = template	1✔
126	variables = variables or []	1✔
127	self._env = SandboxedEnvironment()	1✔
128	if template and not variables:	1✔
129	for message in template:	1✔
130	if message.is_from(ChatRole.USER) or message.is_from(ChatRole.SYSTEM):	1✔
131	# infere variables from template
132	ast = self._env.parse(message.content)	1✔
133	template_variables = meta.find_undeclared_variables(ast)	1✔
134	variables += list(template_variables)	1✔
135
136	# setup inputs
137	static_input_slots = {"template": Optional[str], "template_variables": Optional[Dict[str, Any]]}	1✔
138	component.set_input_types(self, **static_input_slots)	1✔
139	for var in variables:	1✔
140	if var in self.required_variables:	1✔
141	component.set_input_type(self, var, Any)	1✔
142	else:
143	component.set_input_type(self, var, Any, "")	1✔
144
145	@component.output_types(prompt=List[ChatMessage])	1✔
146	def run(	1✔
147	self,
148	template: Optional[List[ChatMessage]] = None,
149	template_variables: Optional[Dict[str, Any]] = None,
150	**kwargs,
151	):
152	"""
153	Executes the prompt building process.
154
155	It applies the template variables to render the final prompt. You can provide variables either via pipeline
156	(set through `variables` or inferred from `template` at initialization) or via additional template variables
157	set directly to this method. On collision, the variables provided directly to this method take precedence.
158
159	:param template:
160	An optional list of ChatMessages to overwrite ChatPromptBuilder's default template. If None, the default
161	template provided at initialization is used.
162	:param template_variables:
163	An optional dictionary of template variables. These are additional variables users can provide directly
164	to this method in contrast to pipeline variables.
165	:param kwargs:
166	Pipeline variables (typically resolved from a pipeline) which are merged with the provided template
167	variables.
168
169	:returns: A dictionary with the following keys:
170	- `prompt`: The updated list of `ChatMessage` instances after rendering the found templates.
171	:raises ValueError:
172	If `chat_messages` is empty or contains elements that are not instances of `ChatMessage`.
173	"""
174	kwargs = kwargs or {}	1✔
175	template_variables = template_variables or {}	1✔
176	template_variables_combined = {kwargs, template_variables}	1✔
177
178	if template is None:	1✔
179	template = self.template	1✔
180
181	if not template:	1✔
182	raise ValueError(	1✔
183	f"The {self.__class__.__name__} requires a non-empty list of ChatMessage instances. "
184	f"Please provide a valid list of ChatMessage instances to render the prompt."
185	)
186
187	if not all(isinstance(message, ChatMessage) for message in template):	1✔
188	raise ValueError(	1✔
189	f"The {self.__class__.__name__} expects a list containing only ChatMessage instances. "
190	f"The provided list contains other types. Please ensure that all elements in the list "
191	f"are ChatMessage instances."
192	)
193
194	processed_messages = []	1✔
195	for message in template:	1✔
196	if message.is_from(ChatRole.USER) or message.is_from(ChatRole.SYSTEM):	1✔
197	self._validate_variables(set(template_variables_combined.keys()))	1✔
198
199	compiled_template = self._env.from_string(message.content)	1✔
200	rendered_content = compiled_template.render(template_variables_combined)	1✔
201	rendered_message = (	1✔
202	ChatMessage.from_user(rendered_content)
203	if message.is_from(ChatRole.USER)
204	else ChatMessage.from_system(rendered_content)
205	)
206	processed_messages.append(rendered_message)	1✔
207	else:
208	processed_messages.append(message)	1✔
209
210	return {"prompt": processed_messages}	1✔
211
212	def _validate_variables(self, provided_variables: Set[str]):	1✔
213	"""
214	Checks if all the required template variables are provided.
215
216	:param provided_variables:
217	A set of provided template variables.
218	:raises ValueError:
219	If no template is provided or if all the required template variables are not provided.
220	"""
221	missing_variables = [var for var in self.required_variables if var not in provided_variables]	1✔
222	if missing_variables:	1✔
223	missing_vars_str = ", ".join(missing_variables)	1✔
224	raise ValueError(	1✔
225	f"Missing required input variables in ChatPromptBuilder: {missing_vars_str}. "
226	f"Required variables: {self.required_variables}. Provided variables: {provided_variables}."
227	)
228
229	def to_dict(self) -> Dict[str, Any]:	1✔
230	"""
231	Returns a dictionary representation of the component.
232
233	:returns:
234	Serialized dictionary representation of the component.
235	"""
236	if self.template is not None:	1✔
237	template = [m.to_dict() for m in self.template]	1✔
238	else:
239	template = None	×
240
241	return default_to_dict(	1✔
242	self, template=template, variables=self._variables, required_variables=self._required_variables
243	)
244
245	@classmethod	1✔
246	def from_dict(cls, data: Dict[str, Any]) -> "ChatPromptBuilder":	1✔
247	"""
248	Deserialize this component from a dictionary.
249
250	:param data:
251	The dictionary to deserialize and create the component.
252
253	:returns:
254	The deserialized component.
255	"""
256	init_parameters = data["init_parameters"]	1✔
257	template = init_parameters.get("template", [])	1✔
258	init_parameters["template"] = [ChatMessage.from_dict(d) for d in template]	1✔
259
260	return default_from_dict(cls, data)	1✔

deepset-ai / haystack / 10112264105

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous