From ca6e07567a4468ff6937580ede4f1be83adb2fbb Mon Sep 17 00:00:00 2001 From: Ishaan Date: Wed, 5 Aug 2026 09:04:01 +0000 Subject: [PATCH] feat(instructions_utils): add Jinja2-based templating with use_jinja2 flag The existing regex-based substitution in `inject_session_state` cannot express Signed-off-by: Ishaan --- src/google/adk/utils/instructions_utils.py | 79 ++++++++++++++++++- .../utils/test_instructions_utils.py | 66 ++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) diff --git a/src/google/adk/utils/instructions_utils.py b/src/google/adk/utils/instructions_utils.py index c42d674d752..18783ea94a4 100644 --- a/src/google/adk/utils/instructions_utils.py +++ b/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. diff --git a/tests/unittests/utils/test_instructions_utils.py b/tests/unittests/utils/test_instructions_utils.py index 78e84d82688..a563af25b58 100644 --- a/tests/unittests/utils/test_instructions_utils.py +++ b/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