diff --git a/src/google/adk/utils/instructions_utils.py b/src/google/adk/utils/instructions_utils.py index 92583dd10f..17bf017d09 100644 --- a/src/google/adk/utils/instructions_utils.py +++ b/src/google/adk/utils/instructions_utils.py @@ -16,6 +16,7 @@ import logging import re +from typing import Any from ..agents.readonly_context import ReadonlyContext from ..sessions.state import State @@ -46,7 +47,11 @@ async def build_instruction( ) -> str: return await inject_session_state( 'You can inject a state variable like {var_name} or an artifact ' - '{artifact.file_name} into the instruction template.', + '{artifact.file_name} into the instruction template.' + 'You can also inject a nested variable like {var_name.nested_var}.' + 'If a variable or nested attribute may be missing, append `?` to the ' + 'path or attribute name for optional handling, e.g. ' + '{var_name.optional_nested_var?}.', readonly_context, ) @@ -78,14 +83,52 @@ async def _async_sub(pattern, repl_async_fn, string) -> str: result.append(string[last_end:]) return ''.join(result) + def _get_nested_value(obj: Any, path: str) -> Any: + """Retrieve nested value from an object based on dot-separated path.""" + parts = path.split('.') + current = obj + + for part in parts: + if current is None: + return None + + optional = part.endswith('?') + key = part[:-1] if optional else part + + # Try dictionary access first + if hasattr(current, '__getitem__'): + try: + current = current[key] + continue + except (KeyError, TypeError): + # If dict access fails, fall through to try getattr + # UNLESS it's a pure dict which definitely doesn't have attributes + if isinstance(current, dict): + if optional: + return None + raise KeyError(f"Key '{key}' not found in path '{path}'") + pass + + # Try attribute access + try: + current = getattr(current, key) + except AttributeError: + # Both dict access and attribute access failed. + if optional: + return None + raise KeyError(f"Key '{key}' not found in path '{path}'") + + return current + async def _replace_match(match) -> str: - var_name = match.group().lstrip('{').rstrip('}').strip() - optional = False - if var_name.endswith('?'): - optional = True - var_name = var_name.removesuffix('?') - if var_name.startswith('artifact.'): - var_name = var_name.removeprefix('artifact.') + full_path = match.group().lstrip('{').rstrip('}').strip() + + if full_path.startswith('artifact.'): + var_name = full_path.removeprefix('artifact.') + optional = var_name.endswith('?') + if optional: + var_name = var_name[:-1] + if invocation_context.artifact_service is None: raise ValueError('Artifact service is not initialized.') artifact = await invocation_context.artifact_service.load_artifact( @@ -104,22 +147,17 @@ async def _replace_match(match) -> str: raise KeyError(f'Artifact {var_name} not found.') return str(artifact) else: - if not _is_valid_state_name(var_name): + if not _is_valid_state_name(full_path.split('.')[0].removesuffix('?')): return match.group() - if var_name in invocation_context.session.state: - value = invocation_context.session.state[var_name] + + try: + value = _get_nested_value(invocation_context.session.state, full_path) + if value is None: return '' return str(value) - else: - if optional: - logger.debug( - 'Context variable %s not found, replacing with empty string', - var_name, - ) - return '' - else: - raise KeyError(f'Context variable not found: `{var_name}`.') + except KeyError as e: + raise KeyError(f'Context variable not found: `{full_path}`.') from e return await _async_sub(r'{+[^{}]*}+', _replace_match, template) diff --git a/tests/unittests/utils/test_instructions_utils.py b/tests/unittests/utils/test_instructions_utils.py index 0a615aa5a5..2572df51b0 100644 --- a/tests/unittests/utils/test_instructions_utils.py +++ b/tests/unittests/utils/test_instructions_utils.py @@ -267,3 +267,218 @@ async def test_inject_session_state_with_optional_missing_state_returns_empty(): instruction_template, invocation_context ) assert populated_instruction == "Optional value: " + + +# Tests for nested state access feature +@pytest.mark.asyncio +async def test_inject_session_state_with_nested_dict_access(): + instruction_template = ( + "User name is {user.name} and role is {user.profile.role}" + ) + invocation_context = await _create_test_readonly_context( + state={ + "user": { + "name": "Alice", + "profile": {"role": "Engineer", "level": "Senior"}, + } + } + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "User name is Alice and role is Engineer" + + +@pytest.mark.asyncio +async def test_inject_session_state_with_deep_nested_access(): + instruction_template = "Deep value: {level1.level2.level3.value}" + invocation_context = await _create_test_readonly_context( + state={ + "level1": { + "level2": {"level3": {"value": "deep_data", "other": "ignored"}} + } + } + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Deep value: deep_data" + + +@pytest.mark.asyncio +async def test_inject_session_state_with_optional_nested_access_existing(): + instruction_template = "Name: {user?.name} Role: {user?.profile?.role}" + invocation_context = await _create_test_readonly_context( + state={ + "user": { + "name": "Bob", + "profile": {"role": "Developer"}, + } + } + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Name: Bob Role: Developer" + + +@pytest.mark.asyncio +async def test_inject_session_state_with_optional_nested_access_missing(): + instruction_template = "Name: {user?.name} Missing: {user?.missing?.field?}" + invocation_context = await _create_test_readonly_context( + state={"user": {"name": "Charlie"}} + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Name: Charlie Missing: " + + +@pytest.mark.asyncio +async def test_inject_session_state_with_optional_nested_missing_root(): + instruction_template = "Optional nested: {missing_root?.nested?.value?}" + invocation_context = await _create_test_readonly_context(state={}) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Optional nested: " + + +@pytest.mark.asyncio +async def test_inject_session_state_with_nested_none_value(): + instruction_template = "Value: {user.profile.role}" + invocation_context = await _create_test_readonly_context( + state={"user": {"profile": None}} + ) + + # When a value in the path is None, it returns empty string + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Value: " + + +@pytest.mark.asyncio +async def test_inject_session_state_with_optional_nested_none_value(): + instruction_template = "Value: {user.profile?.role?}" + invocation_context = await _create_test_readonly_context( + state={"user": {"profile": None}} + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Value: " + + +@pytest.mark.asyncio +async def test_inject_session_state_with_missing_nested_key_raises_error(): + instruction_template = "Value: {user.profile.missing_key}" + invocation_context = await _create_test_readonly_context( + state={"user": {"profile": {"role": "Engineer"}}} + ) + + with pytest.raises( + KeyError, match="Context variable not found: `user.profile.missing_key`" + ): + await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + + +@pytest.mark.asyncio +async def test_inject_session_state_with_required_parent_missing_raises_error(): + """Test that {user.profile?} raises error when 'user' (required) is missing. + + This verifies that optional chaining is per-segment, not for the whole path. + Even though 'profile?' is optional, 'user' is required and should raise error. + """ + instruction_template = "Value: {user.profile?}" + invocation_context = await _create_test_readonly_context(state={}) + + with pytest.raises( + KeyError, match="Context variable not found: `user.profile\\?`" + ): + await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + + +@pytest.mark.asyncio +async def test_inject_session_state_with_nested_and_prefixed_state(): + instruction_template = "User: {app:user.name} Temp: {temp:session.id}" + invocation_context = await _create_test_readonly_context( + state={ + "app:user": {"name": "Dana"}, + "temp:session": {"id": "session_123"}, + } + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "User: Dana Temp: session_123" + + +@pytest.mark.asyncio +async def test_inject_session_state_with_mixed_nested_and_flat_state(): + instruction_template = ( + "Flat: {simple_key}, Nested: {user.name}, Deep: {config.app.version}" + ) + invocation_context = await _create_test_readonly_context( + state={ + "simple_key": "simple_value", + "user": {"name": "Eve"}, + "config": {"app": {"version": "1.0.0"}}, + } + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Flat: simple_value, Nested: Eve, Deep: 1.0.0" + + +@pytest.mark.asyncio +async def test_inject_session_state_with_numeric_nested_values(): + instruction_template = "Age: {user.age}, Score: {user.metrics.score}" + invocation_context = await _create_test_readonly_context( + state={"user": {"age": 25, "metrics": {"score": 95.5}}} + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Age: 25, Score: 95.5" + + +@pytest.mark.asyncio +async def test_inject_session_state_with_nested_object_attribute_access(): + """Test accessing attributes on objects (not just dicts)""" + + class UserProfile: + + def __init__(self): + self.role = "Engineer" + self.department = "Engineering" + + class User: + + def __init__(self): + self.name = "Frank" + self.profile = UserProfile() + + instruction_template = "Name: {user.name}, Role: {user.profile.role}" + invocation_context = await _create_test_readonly_context( + state={"user": User()} + ) + + populated_instruction = await instructions_utils.inject_session_state( + instruction_template, invocation_context + ) + assert populated_instruction == "Name: Frank, Role: Engineer"