✨ Add block assignment functionality in jinja package. (#1341)

This PR implements functionality of block assignments
(https://jinja.palletsprojects.com/en/stable/templates/#block-assignments)
in lexer, parser and engine=
I felt encouraged by the reply here:

#1087 (comment)

It was necessary to parse some chat_templates e.g.
https://huggingface.co/MadeAgents/Hammer2.1-1.5b/blob/main/tokenizer_config.json
(nicely formatted `chat_template`: https://pastebin.com/xBmhbwLi)

Let me know if you have any feedback.

---------

Co-authored-by: Joshua Lochner <[email protected]>
Co-authored-by: Joshua Lochner <[email protected]>

Author: pweglik

packages/jinja/src/ast.ts

@@ -51,7 +51,8 @@ export class SetStatement extends Statement {
 	override type = "Set";
 	constructor(
 		public assignee: Expression,
-		public value: Expression
+		public value: Expression | null,
+		public body: Statement[]
 	) {
 		super();
 	}

packages/jinja/src/lexer.ts

@@ -39,6 +39,7 @@ export const TOKEN_TYPES = Object.freeze({
 	Is: "Is",
 	NotIn: "NotIn",
 	Else: "Else",
+	EndSet: "EndSet",
 	EndIf: "EndIf",
 	ElseIf: "ElseIf",
 	EndFor: "EndFor",
@@ -61,6 +62,7 @@ const KEYWORDS = Object.freeze({
 	is: TOKEN_TYPES.Is,
 	if: TOKEN_TYPES.If,
 	else: TOKEN_TYPES.Else,
+	endset: TOKEN_TYPES.EndSet,
 	endif: TOKEN_TYPES.EndIf,
 	elif: TOKEN_TYPES.ElseIf,
 	endfor: TOKEN_TYPES.EndFor,

packages/jinja/src/parser.ts

@@ -131,11 +131,24 @@ export function parse(tokens: Token[]): Program {
 
 		if (is(TOKEN_TYPES.Equals)) {
 			++current;
-			const value = parseSetStatement();
+			const value = parseExpression();
+
+			return new SetStatement(left, value, []);
+		} else {
+			// parsing multiline set here
+			const body: Statement[] = [];
+			expect(TOKEN_TYPES.CloseStatement, "Expected %} token");
+			while (
+				!(tokens[current]?.type === TOKEN_TYPES.OpenStatement && tokens[current + 1]?.type === TOKEN_TYPES.EndSet)
+			) {
+				const another = parseAny();
+				body.push(another);
+			}
+			expect(TOKEN_TYPES.OpenStatement, "Expected {% token");
+			expect(TOKEN_TYPES.EndSet, "Expected endset token");
 
-			return new SetStatement(left, value);
+			return new SetStatement(left, null, body);
 		}
-		return left;
 	}
 
 	function parseIfStatement(): If {

packages/jinja/src/runtime.ts

@@ -587,6 +587,8 @@ export class Interpreter {
 						);
 					case "join":
 						return new StringValue(operand.value.map((x) => x.value).join(""));
+					case "string":
+						return new StringValue(toJSON(operand));
 					default:
 						throw new Error(`Unknown ArrayValue filter: ${filter.value}`);
 				}
@@ -916,7 +918,7 @@ export class Interpreter {
 	}
 
 	private evaluateSet(node: SetStatement, environment: Environment): NullValue {
-		const rhs = this.evaluate(node.value, environment);
+		const rhs = node.value ? this.evaluate(node.value, environment) : this.evaluateBlock(node.body, environment);
 		if (node.assignee.type === "Identifier") {
 			const variableName = (node.assignee as Identifier).value;
 			environment.setVariable(variableName, rhs);

packages/jinja/test/e2e.test.js

@@ -37,6 +37,7 @@ const TEST_STRINGS = {
 	// Set variables
 	VARIABLES: `{% set x = 'Hello' %}{% set y = 'World' %}{{ x + ' ' + y }}`,
 	VARIABLES_2: `{% set x = 'Hello'.split('el')[-1] %}{{ x }}`,
+	VARIABLES_BLOCK: `{% set x %}Hello!\nMultiline/block set!\n{% endset %}{{ x }}`,
 
 	// Numbers
 	NUMBERS: `|{{ 5 }}|{{ -5 }}|{{ add(3, -1) }}|{{ (3 - 1) + (a - 5) - (a + 5)}}|`,
@@ -90,6 +91,7 @@ const TEST_STRINGS = {
 	FILTER_OPERATOR_10: `|{{ " 1 \n 2 \n 3 \n\n " | indent }}|{{ " 1 \n 2 \n 3 \n\n " | indent(2) }}|{{ " 1 \n 2 \n 3 \n\n " | indent(first=True) }}|{{ " 1 \n 2 \n 3 \n\n " | indent(blank=True) }}|{{ " 1 \n 2 \n 3 \n\n " | indent(4, first=True) }}|`,
 	FILTER_OPERATOR_11: `{{ items | rejectattr('key') | length }}`,
 	FILTER_OPERATOR_12: `{{ messages | rejectattr('role', 'equalto', 'system') | length }}`,
+	FILTER_OPERATOR_13: `{{ tools | string }}`,
 
 	// Logical operators between non-Booleans
 	BOOLEAN_NUMERICAL: `|{{ 1 and 2 }}|{{ 1 and 0 }}|{{ 0 and 1 }}|{{ 0 and 0 }}|{{ 1 or 2 }}|{{ 1 or 0 }}|{{ 0 or 1 }}|{{ 0 or 0 }}|{{ not 1 }}|{{ not 0 }}|`,
@@ -705,6 +707,19 @@ const TEST_PARSED = {
 		{ value: "x", type: "Identifier" },
 		{ value: "}}", type: "CloseExpression" },
 	],
+	VARIABLES_BLOCK: [
+		{ value: "{%", type: "OpenStatement" },
+		{ value: "set", type: "Set" },
+		{ value: "x", type: "Identifier" },
+		{ value: "%}", type: "CloseStatement" },
+		{ value: `Hello!\nMultiline/block set!\n`, type: "Text" },
+		{ value: "{%", type: "OpenStatement" },
+		{ value: "endset", type: "EndSet" },
+		{ value: "%}", type: "CloseStatement" },
+		{ value: "{{", type: "OpenExpression" },
+		{ value: "x", type: "Identifier" },
+		{ value: "}}", type: "CloseExpression" },
+	],
 
 	// Numbers
 	NUMBERS: [
@@ -1682,6 +1697,13 @@ const TEST_PARSED = {
 		{ value: "length", type: "Identifier" },
 		{ value: "}}", type: "CloseExpression" },
 	],
+	FILTER_OPERATOR_13: [
+		{ value: "{{", type: "OpenExpression" },
+		{ value: "tools", type: "Identifier" },
+		{ value: "|", type: "Pipe" },
+		{ value: "string", type: "Identifier" },
+		{ value: "}}", type: "CloseExpression" },
+	],
 
 	// Logical operators between non-Booleans
 	BOOLEAN_NUMERICAL: [
@@ -3038,6 +3060,7 @@ const TEST_CONTEXT = {
 	// Set variables
 	VARIABLES: {},
 	VARIABLES_2: {},
+	VARIABLES_BLOCK: {},
 
 	// Numbers
 	NUMBERS: {
@@ -3173,6 +3196,9 @@ const TEST_CONTEXT = {
 	FILTER_OPERATOR_12: {
 		messages: [{ role: "system" }, { role: "user" }, { role: "assistant" }],
 	},
+	FILTER_OPERATOR_13: {
+		tools: [{ name: "some_tool", arguments: { some_name: "string" } }],
+	},
 
 	// Logical operators between non-Booleans
 	BOOLEAN_NUMERICAL: {},
@@ -3300,6 +3326,7 @@ const EXPECTED_OUTPUTS = {
 	// Set variables
 	VARIABLES: "Hello World",
 	VARIABLES_2: "lo",
+	VARIABLES_BLOCK: "Hello!\nMultiline/block set!\n",
 
 	// Numbers
 	NUMBERS: "|5|-5|2|-8|",
@@ -3353,6 +3380,7 @@ const EXPECTED_OUTPUTS = {
 	FILTER_OPERATOR_10: `| 1 \n     2 \n     3 \n\n     | 1 \n   2 \n   3 \n\n   |     1 \n     2 \n     3 \n\n     | 1 \n     2 \n     3 \n    \n     |     1 \n     2 \n     3 \n\n     |`,
 	FILTER_OPERATOR_11: `3`,
 	FILTER_OPERATOR_12: `2`,
+	FILTER_OPERATOR_13: `[{"name": "some_tool", "arguments": {"some_name": "string"}}]`,
 
 	// Logical operators between non-Booleans
 	BOOLEAN_NUMERICAL: `|2|0|0|0|1|1|1|0|false|true|`,