From 68cc4002b8a6f7a704225e79b6321fd8c610af06 Mon Sep 17 00:00:00 2001
From: Pawel <pawel.rzepecki@intel.com>
Date: Fri, 8 Aug 2025 14:44:37 +0200
Subject: [PATCH 1/2] added jinja2 extension class to enable Phi4-reasoning
 loading

---
 src/llm/servable_initializer.cpp | 50 +++++++++++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/llm/servable_initializer.cpp b/src/llm/servable_initializer.cpp
index 431e5a54f8..d2d9b0814b 100644
--- a/src/llm/servable_initializer.cpp
+++ b/src/llm/servable_initializer.cpp
@@ -61,13 +61,61 @@ void GenAiServableInitializer::loadPyTemplateProcessor(std::shared_ptr<GenAiServ
             import json
             from pathlib import Path
 
+            global contextmanager
+            from contextlib import contextmanager
+
             global jinja2
             import jinja2
+            global ImmutableSandboxedEnvironment
             from jinja2.sandbox import ImmutableSandboxedEnvironment
+            from jinja2.ext import Extension
 
             def raise_exception(message):
                 raise jinja2.exceptions.TemplateError(message)
 
+            class AssistantTracker(Extension):
+                # This extension is used to track the indices of assistant-generated tokens in the rendered chat
+                tags = {"generation"}
+
+                def __init__(self, environment: ImmutableSandboxedEnvironment):
+                    # The class is only initiated by jinja.
+                    super().__init__(environment)
+                    environment.extend(activate_tracker=self.activate_tracker)
+                    self._rendered_blocks = None
+                    self._generation_indices = None
+
+                def parse(self, parser: jinja2.parser.Parser) -> jinja2.nodes.CallBlock:
+                    lineno = next(parser.stream).lineno
+                    body = parser.parse_statements(["name:endgeneration"], drop_needle=True)
+                    return jinja2.nodes.CallBlock(self.call_method("_generation_support"), [], [], body).set_lineno(lineno)
+
+                @jinja2.pass_eval_context
+                def _generation_support(self, context: jinja2.nodes.EvalContext, caller: jinja2.runtime.Macro) -> str:
+                    rv = caller()
+                    if self.is_active():
+                        # Only track generation indices if the tracker is active
+                        start_index = len("".join(self._rendered_blocks))
+                        end_index = start_index + len(rv)
+                        self._generation_indices.append((start_index, end_index))
+                    return rv
+
+                def is_active(self) -> bool:
+                    return self._rendered_blocks or self._generation_indices
+
+                @contextmanager
+                def activate_tracker(self, rendered_blocks: list[int], generation_indices: list[int]):
+                    try:
+                        if self.is_active():
+                            raise ValueError("AssistantTracker should not be reused before closed")
+                        self._rendered_blocks = rendered_blocks
+                        self._generation_indices = generation_indices
+
+                        yield
+                    finally:
+                        self._rendered_blocks = None
+                        self._generation_indices = None
+
+
             # Default chat template accepts only single message and outputs only it's 'content'
             # effectively turning it into a regular prompt. 
             default_chat_template = "{% if messages|length != 1 %} {{ raise_exception('This servable accepts only single message requests') }}{% endif %}{{ messages[0]['content'] }}"
@@ -83,7 +131,7 @@ void GenAiServableInitializer::loadPyTemplateProcessor(std::shared_ptr<GenAiServ
             # Try to read template from template.jinja file
             jinja_file = Path(templates_directory + "/template.jinja")
             template_loader = jinja2.FileSystemLoader(searchpath=templates_directory)
-            jinja_env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True, loader=template_loader)
+            jinja_env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True, extensions=[AssistantTracker, jinja2.ext.loopcontrols], loader=template_loader)
             jinja_env.policies["json.dumps_kwargs"]["ensure_ascii"] = False
             jinja_env.globals["raise_exception"] = raise_exception     
             if jinja_file.is_file():

From 7f26615175a36def8e0797f1eb23b5ba6d2171e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Rzepecki?= <pawel.rzepecki@intel.com>
Date: Mon, 11 Aug 2025 06:47:34 +0100
Subject: [PATCH 2/2] Added comment about the source of the new class

---
 src/llm/servable_initializer.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/llm/servable_initializer.cpp b/src/llm/servable_initializer.cpp
index d2d9b0814b..edcbcb3a67 100644
--- a/src/llm/servable_initializer.cpp
+++ b/src/llm/servable_initializer.cpp
@@ -72,7 +72,9 @@ void GenAiServableInitializer::loadPyTemplateProcessor(std::shared_ptr<GenAiServ
 
             def raise_exception(message):
                 raise jinja2.exceptions.TemplateError(message)
-
+            
+            # Following the logic from:
+            # https://github.com/huggingface/transformers/blob/7188e2e28c6d663284634732564143b820a03f8b/src/transformers/utils/chat_template_utils.py#L398
             class AssistantTracker(Extension):
                 # This extension is used to track the indices of assistant-generated tokens in the rendered chat
                 tags = {"generation"}