feat(instructions_utils): add Jinja2-based templating with use_jinja2 flag

src/google/adk/utils/instructions_utils.py

@@ -20,6 +20,7 @@
 from typing import Callable
 from typing import Union
 
+import jinja2
 from typing_extensions import TypeAlias
 
 from ..agents.readonly_context import ReadonlyContext
@@ -42,6 +43,7 @@
 async def inject_session_state(
     template: str,
     readonly_context: ReadonlyContext,
+    use_jinja2: bool = False,
 ) -> str:
   """Populates values in the instruction template, e.g. state, artifact, etc.
 
@@ -69,13 +71,43 @@ async def build_instruction(
   )
   ```
 
+  For more expressive templates with conditionals and loops, set
+  ``use_jinja2=True``.  Session state variables are available directly by
+  name (``{{ var_name }}``) and artifacts can be loaded with the async
+  ``artifact`` helper (``{{ artifact("file_name") | await }}``).
+
+  e.g.
+  ```
+  async def build_instruction(
+      readonly_context: ReadonlyContext,
+  ) -> str:
+    return await inject_session_state(
+        '{% if user_name %}Hello {{ user_name }}!{% endif %}',
+        readonly_context,
+        use_jinja2=True,
+    )
+  ```
+
   Args:
     template: The instruction template.
-    readonly_context: The read-only context
+    readonly_context: The read-only context.
+    use_jinja2: If True, render the template with Jinja2 instead of the
+      default regex-based engine.  Defaults to False for backward
+      compatibility.
 
   Returns:
     The instruction template with values populated.
   """
+  if use_jinja2:
+    return await _render_with_jinja2(template, readonly_context)
+  return await _render_with_regex(template, readonly_context)
+
+
+async def _render_with_regex(
+    template: str,
+    readonly_context: ReadonlyContext,
+) -> str:
+  """Renders *template* using the legacy regex-based substitution engine."""
 
   # The substitution pattern requires a '{', so a template without one can
   # never match. Return it as-is to avoid the regex scan on every LLM call,
@@ -142,6 +174,51 @@ async def _replace_match(match) -> str:
   return await _async_sub(r'{+[^{}]*}+', _replace_match, template)
 
 
+async def _render_with_jinja2(
+    template: str,
+    readonly_context: ReadonlyContext,
+) -> str:
+  """Renders *template* using a Jinja2 environment.
+
+  Session state variables are exposed as top-level template variables.
+  Artifacts can be loaded with the ``artifact(filename)`` async callable
+  available inside the template.
+
+  Args:
+    template: A Jinja2 template string.
+    readonly_context: The read-only context.
+
+  Returns:
+    The rendered string.
+  """
+  invocation_context = readonly_context._invocation_context
+
+  async def _load_artifact(filename: str) -> str:
+    if invocation_context.artifact_service is None:
+      raise ValueError('Artifact service is not initialized.')
+    artifact = await invocation_context.artifact_service.load_artifact(
+        app_name=invocation_context.session.app_name,
+        user_id=invocation_context.session.user_id,
+        session_id=invocation_context.session.id,
+        filename=filename,
+    )
+    if artifact is None:
+      raise KeyError(f'Artifact {filename} not found.')
+    return str(artifact)
+
+  env = jinja2.Environment(
+      enable_async=True,
+      undefined=jinja2.StrictUndefined,
+      autoescape=False,
+  )
+  jinja_template = env.from_string(template)
+
+  context_vars = dict(invocation_context.session.state)
+  context_vars['artifact'] = _load_artifact
+
+  return await jinja_template.render_async(**context_vars)
+
+
 def _is_valid_state_name(var_name):
   """Checks if the variable name is a valid state name.
 

tests/unittests/utils/test_instructions_utils.py

@@ -284,6 +284,72 @@ async def test_inject_session_state_with_optional_missing_state_returns_empty():
   assert populated_instruction == "Optional value: "
 
 
+@pytest.mark.asyncio
+async def test_inject_session_state_jinja2_basic_variable():
+  instruction_template = "Hello {{ user_name }}, you are in {{ app_state }} state."
+  invocation_context = await _create_test_readonly_context(
+      state={"user_name": "Foo", "app_state": "active"}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context, use_jinja2=True
+  )
+  assert populated_instruction == "Hello Foo, you are in active state."
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_jinja2_conditional():
+  instruction_template = "{% if show_hint %}Hint: read the docs.{% endif %}"
+  invocation_context = await _create_test_readonly_context(
+      state={"show_hint": True}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context, use_jinja2=True
+  )
+  assert populated_instruction == "Hint: read the docs."
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_jinja2_for_loop():
+  instruction_template = "{% for item in items %}{{ item }} {% endfor %}"
+  invocation_context = await _create_test_readonly_context(
+      state={"items": ["a", "b", "c"]}
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context, use_jinja2=True
+  )
+  assert populated_instruction == "a b c "
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_jinja2_artifact():
+  instruction_template = "Content: {{ artifact('my_file') | await }}"
+  mock_artifact_service = MockArtifactService(
+      {"my_file": "artifact data"}
+  )
+  invocation_context = await _create_test_readonly_context(
+      artifact_service=mock_artifact_service
+  )
+
+  populated_instruction = await instructions_utils.inject_session_state(
+      instruction_template, invocation_context, use_jinja2=True
+  )
+  assert populated_instruction == "Content: artifact data"
+
+
+@pytest.mark.asyncio
+async def test_inject_session_state_jinja2_undefined_variable_raises():
+  instruction_template = "Hello {{ missing_var }}!"
+  invocation_context = await _create_test_readonly_context()
+
+  with pytest.raises(Exception):
+    await instructions_utils.inject_session_state(
+        instruction_template, invocation_context, use_jinja2=True
+    )
+
+
 def test_module_exposes_instruction_provider_alias():
   assert instructions_utils.InstructionProvider is InstructionProvider