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'] }}"