diff --git a/src/google/adk/utils/instructions_utils.py b/src/google/adk/utils/instructions_utils.py index c42d674d75..41239c7109 100644 --- a/src/google/adk/utils/instructions_utils.py +++ b/src/google/adk/utils/instructions_utils.py @@ -42,6 +42,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 +70,43 @@ async def build_instruction( ) ``` + To use Jinja2 templating for advanced use cases such as conditionals and + loops, set ``use_jinja2=True``: + + ``` + 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, + ) + ``` + + When ``use_jinja2=True`` the session state keys are available directly as + template variables. An async ``artifact(filename)`` callable is also + injected so you can load artifacts inline: + ``{{ artifact('my_file.txt') }}``. + Args: template: The instruction template. - readonly_context: The read-only context + readonly_context: The read-only context. + use_jinja2: When True, render the template with Jinja2 instead of the + default regex-based substitution. Defaults to False to preserve + 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 the template using the default regex-based substitution.""" # 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 +173,54 @@ 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 Jinja2. + + Session state keys are exposed directly as template variables. An async + ``artifact(filename)`` helper is also available inside the template. + + Args: + template: A Jinja2 template string. + readonly_context: The read-only context. + + Returns: + The rendered string. + """ + try: + from jinja2 import Environment # pylint: disable=g-import-not-at-top + except ImportError as e: + raise ImportError( + 'Jinja2 is required when use_jinja2=True. ' + 'Install it with: pip install jinja2' + ) from e + + invocation_context = readonly_context._invocation_context + + async def _artifact(filename: str) -> str: + """Async helper exposed to Jinja2 templates for loading artifacts.""" + if invocation_context.artifact_service is None: + raise ValueError('Artifact service is not initialized.') + result = 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 result is None: + raise KeyError(f'Artifact {filename} not found.') + return str(result) + + env = Environment(enable_async=True) # pylint: disable=invalid-name + jinja_template = env.from_string(template) + + ctx = dict(invocation_context.session.state) + ctx['artifact'] = _artifact + return await jinja_template.render_async(**ctx) + + 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 78e84d8268..99226e3e85 100644 --- a/tests/unittests/utils/test_instructions_utils.py +++ b/tests/unittests/utils/test_instructions_utils.py @@ -292,3 +292,87 @@ def test_llm_agent_reexports_same_instruction_provider(): # Existing importers rely on `from ...llm_agent import InstructionProvider`; # it must remain the exact same object after moving the alias here. assert LlmAgentInstructionProvider is InstructionProvider + + +# --------------------------------------------------------------------------- +# Jinja2-based templating (use_jinja2=True) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_basic_variable(): + invocation_context = await _create_test_readonly_context( + state={"user_name": "Alice"} + ) + result = await instructions_utils.inject_session_state( + "Hello {{ user_name }}!", + invocation_context, + use_jinja2=True, + ) + assert result == "Hello Alice!" + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_conditional(): + invocation_context = await _create_test_readonly_context( + state={"is_admin": True} + ) + result = await instructions_utils.inject_session_state( + "{% if is_admin %}Admin mode{% else %}User mode{% endif %}", + invocation_context, + use_jinja2=True, + ) + assert result == "Admin mode" + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_for_loop(): + invocation_context = await _create_test_readonly_context( + state={"items": ["a", "b", "c"]} + ) + result = await instructions_utils.inject_session_state( + "{% for item in items %}{{ item }}{% endfor %}", + invocation_context, + use_jinja2=True, + ) + assert result == "abc" + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_artifact(): + mock_artifact_service = MockArtifactService({"doc.txt": "file content"}) + invocation_context = await _create_test_readonly_context( + artifact_service=mock_artifact_service + ) + result = await instructions_utils.inject_session_state( + "Content: {{ artifact('doc.txt') }}", + invocation_context, + use_jinja2=True, + ) + assert result == "Content: file content" + + +@pytest.mark.asyncio +async def test_inject_session_state_jinja2_missing_artifact_raises(): + mock_artifact_service = MockArtifactService({}) + invocation_context = await _create_test_readonly_context( + artifact_service=mock_artifact_service + ) + with pytest.raises(KeyError, match="Artifact missing.txt not found."): + await instructions_utils.inject_session_state( + "{{ artifact('missing.txt') }}", + invocation_context, + use_jinja2=True, + ) + + +@pytest.mark.asyncio +async def test_inject_session_state_default_still_uses_regex(): + """Passing no use_jinja2 flag must still use the regex path.""" + invocation_context = await _create_test_readonly_context( + state={"key": "value"} + ) + result = await instructions_utils.inject_session_state( + "The key is {key}.", invocation_context + ) + assert result == "The key is value."