Coverage for haystack/utils/jinja2_chat_extension.py: 99%
146 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import json
6import secrets
7from collections.abc import Callable
8from typing import Any
10from jinja2 import TemplateSyntaxError, nodes, pass_environment
11from jinja2.ext import Extension
12from markupsafe import Markup
14from haystack import logging
15from haystack.dataclasses.chat_message import (
16 ChatMessage,
17 ChatMessageContentT,
18 ChatRole,
19 FileContent,
20 ImageContent,
21 ReasoningContent,
22 TextContent,
23 ToolCall,
24 ToolCallResult,
25 _deserialize_content_part,
26 _serialize_content_part,
27)
29logger = logging.getLogger(__name__)
31_TAG_NAME = "haystack_content_part"
32_NONCE_ATTR = "_haystack_content_part_nonce"
35def _sentinel_tags(nonce: str) -> tuple[str, str]:
36 """Build the (opening, closing) sentinel tags for the given nonce"""
37 return f"<{_TAG_NAME}:{nonce}>", f"</{_TAG_NAME}:{nonce}>"
40class _TemplatizedPart(Markup):
41 """Marker type for content produced by `templatize_part`."""
43 pass
46def _finalize(value: object) -> str:
47 """
48 Jinja2 `finalize` callback that prevents sentinel tag injection.
50 Called automatically on every `{{ }}` expression result during template rendering.
51 Legitimate structured content from the `templatize_part` filter is wrapped in `_TemplatizedPart` and passes.
52 Any other value containing sentinel tags has those tags replaced with harmless HTML entities so that
53 `_parse_content_parts` will not treat them as structured content.
54 """
55 if isinstance(value, _TemplatizedPart):
56 return value
58 # escape the leading "<" of any sentinel tag so that `_parse_content_parts` cannot recognize it
59 return str(value).replace(f"<{_TAG_NAME}", f"<{_TAG_NAME}").replace(f"</{_TAG_NAME}", f"</{_TAG_NAME}")
62def _redact_nonce(text: str, start_tag: str, end_tag: str) -> str:
63 return text.replace(start_tag, f"<{_TAG_NAME}:redacted>").replace(end_tag, f"</{_TAG_NAME}:redacted>")
66class ChatMessageExtension(Extension):
67 """
68 A Jinja2 extension for creating structured chat messages with mixed content types.
70 This extension provides a custom `{% message %}` tag that allows creating chat messages
71 with different attributes (role, name, meta) and mixed content types (text, images, etc.).
73 Inspired by [Banks](https://github.com/masci/banks).
75 Example:
76 ```
77 {% message role="system" %}
78 You are a helpful assistant. You like to talk with {{user_name}}.
79 {% endmessage %}
81 {% message role="user" %}
82 Hello! I am {{user_name}}. Please describe the images.
83 {% for image in images %}
84 {{ image | templatize_part }}
85 {% endfor %}
86 {% endmessage %}
87 ```
89 This extension also provides an `{% insert %}` placeholder tag that evaluates an expression to a `ChatMessage`
90 or a list of `ChatMessage` objects and expands it into the prompt, so a runtime conversation can be interleaved
91 with literal `{% message %}` blocks:
93 ```
94 {% message role="system" %}You are a helpful assistant.{% endmessage %}
95 {% insert messages %}
96 {% message role="user" %}{{ query }}{% endmessage %}
97 ```
99 The expression can be a plain variable (`{% insert messages %}`), a slice or index
100 (`{% insert messages[-1:] %}`, `{% insert messages[-1] %}`), or a combination of variables
101 (`{% insert previous + current %}`).
103 ### How it works
104 1. The `{% message %}` tag is used to define a chat message.
105 2. The message can contain text and other structured content parts.
106 3. To include a structured content part in the message, the `| templatize_part` filter is used.
107 The filter serializes the content part into a JSON string and wraps it in a `<haystack_content_part>` tag.
108 4. The `_build_chat_message_json` method of the extension parses the message content parts,
109 converts them into a ChatMessage object and serializes it to a JSON string.
110 5. The obtained JSON string is usable in the ChatPromptBuilder component, where templates are rendered to actual
111 ChatMessage objects.
112 """
114 SUPPORTED_ROLES = [role.value for role in ChatRole]
116 tags = {"message", "insert"}
118 def __init__(self, environment: Any) -> None:
119 super().__init__(environment)
120 # a fresh random nonce per environment to produce sentinel tags only usable by `templatize_part`
121 setattr(environment, _NONCE_ATTR, secrets.token_hex(16))
122 # values not produced by `templatize_part` get their sentinel-like tags escaped
123 environment.finalize = _finalize
124 environment.filters["templatize_part"] = templatize_part
126 def parse(self, parser: Any) -> nodes.Node | list[nodes.Node]:
127 """
128 Dispatch parsing based on the tag that triggered the extension.
130 Handles both the single `{% message %}` block tag and the `{% insert %}` placeholder tag.
132 :param parser: The Jinja2 parser instance
133 :return: A CallBlock node containing the parsed configuration
134 """
135 tag = next(parser.stream)
136 if tag.value == "insert":
137 return self._parse_insert_tag(parser, tag.lineno)
138 return self._parse_message_tag(parser, tag.lineno)
140 def _parse_insert_tag(self, parser: Any, lineno: int) -> nodes.Node:
141 """
142 Parse the `{% insert %}` placeholder tag.
144 This bodyless tag evaluates an expression to a `ChatMessage` or a list of `ChatMessage` objects and expands
145 it into the same JSON-line format produced by `{% message %}` blocks, so messages provided at runtime can be
146 interleaved with literal message blocks (for example a system message above and a user message below).
148 The expression can be a plain variable (`{% insert messages %}`), a slice or index
149 (`{% insert messages[-1:] %}`, `{% insert messages[-1] %}`), or a combination of variables
150 (`{% insert previous + current %}`).
152 :param parser: The Jinja2 parser instance
153 :param lineno: The line number of the tag, used for error reporting.
154 :return: A CallBlock node that expands the evaluated expression.
155 :raises TemplateSyntaxError: If the tag is not given an expression.
156 """
157 if parser.stream.current.test("block_end"):
158 raise TemplateSyntaxError(
159 "The 'insert' tag requires an expression that evaluates to a ChatMessage or a list of ChatMessage "
160 "objects, for example '{% insert messages %}' or '{% insert messages[-1:] %}'.",
161 lineno,
162 )
163 expr = parser.parse_expression()
164 # Bodyless tag: empty body, no matching end tag required.
165 return nodes.CallBlock(
166 self.call_method(name="_build_inserted_messages_json", args=[expr]), [], [], []
167 ).set_lineno(lineno)
169 def _parse_message_tag(self, parser: Any, lineno: int) -> nodes.Node | list[nodes.Node]:
170 """
171 Parse the message tag and its attributes in the Jinja2 template.
173 This method handles the parsing of role (mandatory), name (optional), meta (optional) and message body content.
175 :param parser: The Jinja2 parser instance
176 :param lineno: The line number of the tag, used for error reporting.
177 :return: A CallBlock node containing the parsed message configuration
178 :raises TemplateSyntaxError: If an invalid role is provided
179 """
181 # Parse role attribute (mandatory)
182 parser.stream.expect("name:role")
183 parser.stream.expect("assign")
184 role_expr = parser.parse_expression()
186 if isinstance(role_expr, nodes.Const):
187 role = role_expr.value
188 if role not in self.SUPPORTED_ROLES:
189 raise TemplateSyntaxError(f"Role must be one of: {', '.join(self.SUPPORTED_ROLES)}", lineno)
191 # Parse optional name attribute
192 name_expr = None
193 if parser.stream.current.test("name:name"):
194 parser.stream.skip()
195 parser.stream.expect("assign")
196 name_expr = parser.parse_expression()
197 if not isinstance(name_expr.value, str):
198 raise TemplateSyntaxError("name must be a string", lineno)
200 # Parse optional meta attribute
201 meta_expr = None
202 if parser.stream.current.test("name:meta"):
203 parser.stream.skip()
204 parser.stream.expect("assign")
205 meta_expr = parser.parse_expression()
206 if not isinstance(meta_expr, nodes.Dict):
207 raise TemplateSyntaxError("meta must be a dictionary", lineno)
209 # Parse message body
210 body = parser.parse_statements(("name:endmessage",), drop_needle=True)
212 # Build message node with all parameters
213 return nodes.CallBlock(
214 self.call_method(
215 name="_build_chat_message_json",
216 args=[role_expr, name_expr or nodes.Const(None), meta_expr or nodes.Dict([])],
217 ),
218 [],
219 [],
220 body,
221 ).set_lineno(lineno)
223 def _build_chat_message_json(self, role: str, name: str | None, meta: dict, caller: Callable[[], str]) -> str:
224 """
225 Build a ChatMessage object from template content and serialize it to a JSON string.
227 This method is called by Jinja2 when processing a `{% message %}` tag.
228 It takes the rendered content from the template, converts XML blocks into ChatMessageContentT objects,
229 creates a ChatMessage object and serializes it to a JSON string.
231 :param role: The role of the message
232 :param name: Optional name for the message sender
233 :param meta: Optional metadata dictionary
234 :param caller: Callable that returns the rendered content
235 :return: A JSON string representation of the ChatMessage object
236 """
238 content = caller()
239 start_tag, end_tag = _sentinel_tags(getattr(self.environment, _NONCE_ATTR))
240 parts = self._parse_content_parts(content, start_tag, end_tag)
241 if not parts:
242 raise ValueError(
243 f"Message template produced content that couldn't be parsed into any message parts. "
244 f"Content: {_redact_nonce(content, start_tag, end_tag)!r}"
245 )
247 chat_message = self._validate_build_chat_message(parts=parts, role=role, meta=meta, name=name)
249 return json.dumps(chat_message.to_dict()) + "\n"
251 def _build_inserted_messages_json(
252 self,
253 messages: list[ChatMessage] | ChatMessage,
254 caller: Callable[[], str], # noqa: ARG002
255 ) -> str:
256 """
257 Expand a list of ChatMessage objects into newline-separated JSON, one message per line.
259 This method is called by Jinja2 when processing an `{% insert %}` tag. It produces the same JSON-line format
260 as `_build_chat_message_json`, so the messages are parsed back into ChatMessage objects by the
261 ChatPromptBuilder alongside any literal `{% message %}` blocks. The full `ChatMessage.to_dict()` payload is
262 serialized so that all content types (tool calls, tool call results, images, reasoning, name and meta) round
263 trip without loss.
265 :param messages: The value the `{% insert %}` expression evaluated to. A missing or empty value expands to
266 nothing. A single ChatMessage is also accepted, since indexing with an integer (for example
267 `{% insert messages[-1] %}`) yields one message rather than a list. The value is validated at render time
268 because it comes from untrusted template input.
269 :param caller: Callable that returns the (empty) rendered body. Unused.
270 :return: Newline-terminated JSON lines, one per message, or an empty string if there are no messages.
271 :raises ValueError: If the value is not a ChatMessage or a list of ChatMessage objects.
272 """
273 if isinstance(messages, ChatMessage):
274 messages = [messages]
275 if not messages:
276 return ""
277 if not isinstance(messages, (list, tuple)) or not all(isinstance(m, ChatMessage) for m in messages):
278 raise ValueError(
279 "The '{% insert %}' expression must evaluate to a ChatMessage or a list of ChatMessage objects. "
280 f"Got: {type(messages).__name__}."
281 )
282 return "".join(json.dumps(message.to_dict()) + "\n" for message in messages)
284 @staticmethod
285 def _parse_content_parts(content: str, start_tag: str, end_tag: str) -> list[ChatMessageContentT]:
286 """
287 Parse a string into a sequence of ChatMessageContentT objects.
289 This method handles:
290 - Plain text content, converted to TextContent objects
291 - Structured content parts wrapped in sentinel tags, converted to ChatMessageContentT objects
293 :param content: Input string containing mixed text and content parts
294 :param start_tag: The opening sentinel tag (including the nonce)
295 :param end_tag: The closing sentinel tag (including the nonce)
296 :return: A list of ChatMessageContentT objects
297 :raises ValueError: If the content is empty or contains only whitespace characters or if a
298 `<haystack_content_part>` tag is found without a matching closing tag.
299 """
300 if not content.strip():
301 raise ValueError(
302 f"Message content in template is empty or contains only whitespace characters. "
303 f"Content: {_redact_nonce(content, start_tag, end_tag)!r}"
304 )
306 parts: list[ChatMessageContentT] = []
307 cursor = 0
308 total_length = len(content)
310 while cursor < total_length:
311 tag_start = content.find(start_tag, cursor)
313 if tag_start == -1:
314 # No more tags, add remaining text if any
315 remaining_text = content[cursor:].strip()
316 if remaining_text:
317 parts.append(TextContent(text=remaining_text))
318 break
320 # Add text before tag if any
321 if tag_start > cursor:
322 plain_text = content[cursor:tag_start].strip()
323 if plain_text:
324 parts.append(TextContent(text=plain_text))
326 content_start = tag_start + len(start_tag)
327 tag_end = content.find(end_tag, content_start)
329 if tag_end == -1:
330 snippet = _redact_nonce(content, start_tag, end_tag)[tag_start : tag_start + 50]
331 raise ValueError(
332 f"Found unclosed <haystack_content_part> tag at position {tag_start}. Content: '{snippet}...'"
333 )
335 json_content = content[content_start:tag_end]
336 data = json.loads(json_content)
337 parts.append(_deserialize_content_part(data))
339 cursor = tag_end + len(end_tag)
341 return parts
343 @staticmethod
344 def _validate_build_chat_message(
345 parts: list[ChatMessageContentT], role: str, meta: dict, name: str | None = None
346 ) -> ChatMessage:
347 """
348 Validate the parts of a chat message and build a ChatMessage object.
350 :param parts: Content parts of the message
351 :param role: The role of the message
352 :param meta: The metadata of the message
353 :param name: The optional name of the message
354 :return: A ChatMessage object
356 :raises ValueError: If content parts don't allow to build a valid ChatMessage object or the role is not
357 supported
358 """
360 if role == "user":
361 valid_parts = [part for part in parts if isinstance(part, (TextContent, str, ImageContent, FileContent))]
362 if len(parts) != len(valid_parts):
363 raise ValueError(
364 "User message must contain only TextContent, string, ImageContent or FileContent parts."
365 )
366 return ChatMessage.from_user(meta=meta, name=name, content_parts=valid_parts)
368 if role == "system":
369 if not isinstance(parts[0], TextContent):
370 raise ValueError("System message must contain a text part.")
371 text = parts[0].text
372 if len(parts) > 1:
373 raise ValueError("System message must contain only one text part.")
374 return ChatMessage.from_system(meta=meta, name=name, text=text)
376 if role == "assistant":
377 texts = [part.text for part in parts if isinstance(part, TextContent)]
378 tool_calls = [part for part in parts if isinstance(part, ToolCall)]
379 reasoning = [part for part in parts if isinstance(part, ReasoningContent)]
380 if len(texts) > 1:
381 raise ValueError("Assistant message must contain one text part at most.")
382 if len(texts) == 0 and len(tool_calls) == 0:
383 raise ValueError("Assistant message must contain at least one text or tool call part.")
384 if len(parts) > len(texts) + len(tool_calls) + len(reasoning):
385 raise ValueError("Assistant message must contain only text, tool call or reasoning parts.")
386 return ChatMessage.from_assistant(
387 meta=meta,
388 name=name,
389 text=texts[0] if texts else None,
390 tool_calls=tool_calls or None,
391 reasoning=reasoning[0] if reasoning else None,
392 )
394 if role == "tool":
395 tool_call_results = [part for part in parts if isinstance(part, ToolCallResult)]
396 if len(tool_call_results) == 0 or len(tool_call_results) > 1 or len(parts) > len(tool_call_results):
397 raise ValueError("Tool message must contain only one tool call result.")
399 tool_result = tool_call_results[0].result
400 origin = tool_call_results[0].origin
401 error = tool_call_results[0].error
403 return ChatMessage.from_tool(meta=meta, tool_result=tool_result, origin=origin, error=error)
405 raise ValueError(f"Unsupported role: {role}")
408@pass_environment
409def templatize_part(environment: Any, value: ChatMessageContentT) -> "_TemplatizedPart":
410 """
411 Jinja filter to convert a ChatMessageContentT object into a JSON string wrapped in sentinel content tags.
413 :param environment: The Jinja2 environment
414 :param value: The ChatMessageContentT object to convert
415 :return: A `_TemplatizedPart` holding a JSON string wrapped in special XML content tags
416 :raises ValueError: If the value is not an instance of ChatMessageContentT
417 """
418 start_tag, end_tag = _sentinel_tags(getattr(environment, _NONCE_ATTR))
419 return _TemplatizedPart(f"{start_tag}{json.dumps(_serialize_content_part(value))}{end_tag}")