From f99a01e08ab34afb70d82679fc4de14f583d223a Mon Sep 17 00:00:00 2001 From: vaibhav-patel Date: Mon, 22 Jun 2026 15:00:28 +0530 Subject: [PATCH 1/2] feat(utils): add opt-in Jinja2 instruction templating via use_jinja2 flag Add a `use_jinja2: bool = False` flag to `inject_session_state` that, when enabled, renders the instruction template with a sandboxed Jinja2 environment instead of the existing regex-based `{var}` substitution. This unlocks control flow (conditionals, loops) and filters in agent instructions while keeping the default behavior fully backward compatible. The regex rendering logic is extracted unchanged into `_render_with_regex`, and a new `_render_with_jinja2` coroutine handles the Jinja2 path. It exposes the session `state` mapping and an async `artifact(name)` accessor to templates; the environment runs with `enable_async=True` so the artifact coroutine is awaited automatically, and missing artifacts render as empty strings. jinja2 is imported lazily inside the Jinja2 path so the base install is not forced to depend on it, and a dedicated `jinja` optional extra documents the install path (`pip install google-adk[jinja]`). A `jinja2.sandbox.SandboxedEnvironment` is used because instruction templates may carry user/session data. Adds unit tests for variable substitution, conditionals, loops, filters, artifact loading, the uninitialized-artifact-service error, the sandbox blocking unsafe attribute access, and confirmation that the default path is unchanged. Fixes #2942. --- pyproject.toml | 3 + src/google/adk/utils/instructions_utils.py | 113 ++++++++++++- .../utils/test_instructions_utils.py | 157 ++++++++++++++++++ 3 files changed, 272 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 80e4af5fb34..85101c94778 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -179,6 +179,9 @@ optional-dependencies.gcp = [ "pyarrow>=14", "python-dateutil>=2.9.0.post0,<3", ] +optional-dependencies.jinja = [ + "jinja2>=3.1.4,<4", # For Jinja2-based instruction templating (inject_session_state(use_jinja2=True)). +] optional-dependencies.mcp = [ "anyio>=4.9,<5", "mcp>=1.24,<2", diff --git a/src/google/adk/utils/instructions_utils.py b/src/google/adk/utils/instructions_utils.py index 0dfe0a2b7fe..1c61cd89eed 100644 --- a/src/google/adk/utils/instructions_utils.py +++ b/src/google/adk/utils/instructions_utils.py @@ -30,6 +30,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. @@ -57,9 +58,55 @@ async def build_instruction( ) ``` + When ``use_jinja2`` is ``True``, the template is rendered with a sandboxed + Jinja2 environment instead of the default regex-based substitution. This + enables control flow such as conditionals and loops in addition to plain + variable injection. Inside a Jinja2 template, session state is available as + the ``state`` mapping (e.g. ``{{ state['var_name'] }}``), while artifacts are + loaded with the ``artifact`` helper (e.g. ``{{ artifact('file_name') }}``). + The artifact helper is asynchronous and is awaited automatically by the + async Jinja2 environment. + + e.g. + ``` + return await inject_session_state( + '{% if state["is_premium"] %}Premium user.{% else %}Free user.{% endif %}' + '{% for item in state["items"] %}- {{ item }}\\n{% endfor %}', + 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 a sandboxed Jinja2 environment. + If False (the default), use the regex-based ``{var}`` substitution. The + default preserves backward-compatible behavior. + + 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 the template using the regex-based ``{var}`` substitution. + + This is the default, backward-compatible rendering path. It replaces + ``{var_name}`` with the matching session state value and + ``{artifact.file_name}`` with the loaded artifact content. A trailing ``?`` + (e.g. ``{var_name?}``) marks the reference as optional, replacing it with an + empty string when the value is missing instead of raising. + + Args: + template: The instruction template. + readonly_context: The read-only context. Returns: The instruction template with values populated. @@ -124,6 +171,70 @@ 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 the template using a sandboxed Jinja2 environment. + + Unlike the regex-based path, this supports full Jinja2 control flow such as + conditionals (``{% if %}``), loops (``{% for %}``) and filters, in addition + to variable injection. + + The following names are exposed to the template: + - ``state``: the session state mapping, e.g. ``{{ state['var_name'] }}`` or + ``{% if state['flag'] %}...{% endif %}``. + - ``artifact``: an async accessor that loads an artifact by filename, e.g. + ``{{ artifact('file_name') }}``. The environment runs with + ``enable_async=True``, so the returned coroutine is awaited + automatically; a missing artifact renders as an empty string. + + A ``jinja2.sandbox.SandboxedEnvironment`` is used because instruction + templates may include user- or session-provided data, and the sandbox blocks + access to unsafe attributes and operations. ``jinja2`` is imported lazily so + that installations that never use this path are not required to have it. + + Args: + template: The instruction template. + readonly_context: The read-only context. + + Returns: + The instruction template with values populated. + + Raises: + ValueError: If the artifact service is required but not initialized. + """ + try: + from jinja2.sandbox import SandboxedEnvironment + except ImportError as e: + raise ImportError( + 'jinja2 is required to use Jinja2-based instruction templating' + ' (use_jinja2=True). Install it with `pip install google-adk[jinja]`' + ' (or `pip install jinja2`).' + ) from e + + invocation_context = readonly_context._invocation_context + session_state = invocation_context.session.state + + async def _artifact(name: 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=name, + ) + return str(artifact) if artifact is not None else '' + + env = SandboxedEnvironment(enable_async=True) + jinja_template = env.from_string(template) + return await jinja_template.render_async( + state=session_state, + artifact=_artifact, + ) + + 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 9e176241bc4..40ed676197f 100644 --- a/tests/unittests/utils/test_instructions_utils.py +++ b/tests/unittests/utils/test_instructions_utils.py @@ -267,3 +267,160 @@ async def test_inject_session_state_with_optional_missing_state_returns_empty(): instruction_template, invocation_context ) assert populated_instruction == "Optional value: " + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_basic_substitution(): + instruction_template = "Hello {{ state['user_name'] }}, welcome back." + invocation_context = await _create_test_readonly_context( + state={"user_name": "Foo"} + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context, use_jinja2=True + ) + assert populated_instruction == "Hello Foo, welcome back." + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_conditional(): + instruction_template = ( + "{% if state['is_premium'] %}Premium user.{% else %}Free user." + "{% endif %}" + ) + premium_context = await _create_test_readonly_context( + state={"is_premium": True} + ) + free_context = await _create_test_readonly_context( + state={"is_premium": False} + ) + + assert ( + await instructions_utils.inject_session_state( + instruction_template, premium_context, use_jinja2=True + ) + == "Premium user." + ) + assert ( + await instructions_utils.inject_session_state( + instruction_template, free_context, use_jinja2=True + ) + == "Free user." + ) + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_loop(): + instruction_template = ( + "Items:{% for item in state['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 == "Items: a b c" + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_filter(): + instruction_template = "Name: {{ state['name'] | upper }}" + invocation_context = await _create_test_readonly_context( + state={"name": "foo"} + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context, use_jinja2=True + ) + assert populated_instruction == "Name: FOO" + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_artifact(): + # The async Jinja2 environment auto-awaits the artifact coroutine, so the + # template uses a plain call without an `await` keyword. + instruction_template = "Artifact: {{ artifact('my_file') }}" + mock_artifact_service = MockArtifactService( + {"my_file": "This is my artifact content."} + ) + 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 == "Artifact: This is my artifact content." + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_missing_artifact_returns_empty(): + instruction_template = "Artifact: [{{ artifact('missing') }}]" + mock_artifact_service = MockArtifactService({}) + 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 == "Artifact: []" + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_artifact_service_not_initialized(): + instruction_template = "Artifact: {{ artifact('my_file') }}" + invocation_context = await _create_test_readonly_context() + + with pytest.raises(ValueError, match="Artifact service is not initialized."): + await instructions_utils.inject_session_state( + instruction_template, invocation_context, use_jinja2=True + ) + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_sandbox_blocks_unsafe_access(): + # Instruction templates may carry user/session data, so rendering runs in a + # SandboxedEnvironment. The classic sandbox-escape vector (reaching an + # object's class MRO) must be rejected. + from jinja2.exceptions import SecurityError + + instruction_template = "{{ ''.__class__.__mro__ }}" + invocation_context = await _create_test_readonly_context() + + with pytest.raises(SecurityError): + await instructions_utils.inject_session_state( + instruction_template, invocation_context, use_jinja2=True + ) + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_does_not_touch_brace_var_syntax(): + # The legacy {var} syntax is not interpolated by the Jinja2 path; only + # Jinja2 delimiters are. This documents the opt-in boundary between engines. + instruction_template = "Legacy {user_name} and jinja {{ state['user_name'] }}" + invocation_context = await _create_test_readonly_context( + state={"user_name": "Foo"} + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context, use_jinja2=True + ) + assert populated_instruction == "Legacy {user_name} and jinja Foo" + + +@pytest.mark.asyncio +async def test_inject_session_state_default_flag_uses_regex_path(): + # Explicitly assert the default (use_jinja2=False) keeps the regex behavior: + # Jinja2 delimiters are left untouched and {var} is substituted. + instruction_template = "Regex {user_name}, jinja {{ state['user_name'] }}" + invocation_context = await _create_test_readonly_context( + state={"user_name": "Foo"} + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Regex Foo, jinja {{ state['user_name'] }}" From 611b177b77a741cc93b03a1820bbd130a27eddb5 Mon Sep 17 00:00:00 2001 From: vaibhav-patel Date: Mon, 22 Jun 2026 15:34:01 +0530 Subject: [PATCH 2/2] chore: re-trigger CI checks