Skip to content

Commit 7137d8c

Browse files
feat: support dot-separated nested state values in instruction templates
Adds support for referencing nested dictionary values in instruction templates using dot notation, e.g. {user.profile.name} resolves to session.state['user']['profile']['name']. This is useful when agents receive structured LLM responses via output_schema and need to reference nested fields in instructions without manually flattening state. - Add _resolve_nested() helper for dot-separated path traversal - Update _replace_match() to use nested resolution when path contains dots - Update _is_valid_state_name() to accept dot-separated identifiers - Preserve backward compatibility: flat keys, artifact. prefix, and app:/user:/temp: namespaces all work unchanged Fixes #575
1 parent 9199189 commit 7137d8c

File tree

2 files changed

+121
-7
lines changed

2 files changed

+121
-7
lines changed

src/google/adk/utils/instructions_utils.py

Lines changed: 47 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,29 @@
2626

2727
logger = logging.getLogger('google_adk.' + __name__)
2828

29+
_MISSING = object()
30+
31+
32+
def _resolve_nested(state, path: str, optional: bool = False):
33+
"""Traverse a nested dict/State using a dot-separated path.
34+
35+
Args:
36+
state: The state mapping to traverse.
37+
path: A dot-separated key path, e.g. "user.profile.name".
38+
optional: If True, return _MISSING for missing keys instead of raising.
39+
40+
Returns:
41+
The resolved value, or _MISSING if not found and optional is True.
42+
"""
43+
keys = path.split('.')
44+
value = state
45+
for key in keys:
46+
if isinstance(value, dict) and key in value:
47+
value = value[key]
48+
else:
49+
return _MISSING
50+
return value
51+
2952

3053
async def inject_session_state(
3154
template: str,
@@ -36,6 +59,13 @@ async def inject_session_state(
3659
This method is intended to be used in InstructionProvider based instruction
3760
and global_instruction which are called with readonly_context.
3861
62+
Supports dot-separated paths for nested state values, e.g.
63+
``{user.profile.name}`` resolves to
64+
``session.state['user']['profile']['name']``.
65+
66+
Use ``?`` suffix for optional variables that may not exist, e.g.
67+
``{user.preferences?}`` returns empty string if not found.
68+
3969
e.g.
4070
```
4171
...
@@ -45,7 +75,8 @@ async def build_instruction(
4575
readonly_context: ReadonlyContext,
4676
) -> str:
4777
return await inject_session_state(
48-
'You can inject a state variable like {var_name} or an artifact '
78+
'You can inject a state variable like {var_name} or a nested '
79+
'value like {user.profile.name} or an artifact '
4980
'{artifact.file_name} into the instruction template.',
5081
readonly_context,
5182
)
@@ -106,12 +137,16 @@ async def _replace_match(match) -> str:
106137
else:
107138
if not _is_valid_state_name(var_name):
108139
return match.group()
109-
if var_name in invocation_context.session.state:
110-
value = invocation_context.session.state[var_name]
111-
if value is None:
112-
return ''
113-
return str(value)
140+
state = invocation_context.session.state
141+
if '.' in var_name and not var_name.startswith(
142+
(State.APP_PREFIX, State.USER_PREFIX, State.TEMP_PREFIX)
143+
):
144+
value = _resolve_nested(state, var_name, optional)
145+
elif var_name in state:
146+
value = state[var_name]
114147
else:
148+
value = _MISSING
149+
if value is _MISSING:
115150
if optional:
116151
logger.debug(
117152
'Context variable %s not found, replacing with empty string',
@@ -120,6 +155,9 @@ async def _replace_match(match) -> str:
120155
return ''
121156
else:
122157
raise KeyError(f'Context variable not found: `{var_name}`.')
158+
if value is None:
159+
return ''
160+
return str(value)
123161

124162
return await _async_sub(r'{+[^{}]*}+', _replace_match, template)
125163

@@ -129,6 +167,7 @@ def _is_valid_state_name(var_name):
129167
130168
Valid state is either:
131169
- Valid identifier
170+
- Dot-separated valid identifiers (e.g. "user.profile.name")
132171
- <Valid prefix>:<Valid identifier>
133172
All the others will just return as it is.
134173
@@ -140,7 +179,8 @@ def _is_valid_state_name(var_name):
140179
"""
141180
parts = var_name.split(':')
142181
if len(parts) == 1:
143-
return var_name.isidentifier()
182+
# Support dot-separated nested paths like "user.profile.name"
183+
return all(seg.isidentifier() for seg in var_name.split('.'))
144184

145185
if len(parts) == 2:
146186
prefixes = [State.APP_PREFIX, State.USER_PREFIX, State.TEMP_PREFIX]

tests/unittests/utils/test_instructions_utils.py

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,3 +267,77 @@ async def test_inject_session_state_with_optional_missing_state_returns_empty():
267267
instruction_template, invocation_context
268268
)
269269
assert populated_instruction == "Optional value: "
270+
271+
272+
@pytest.mark.asyncio
273+
async def test_inject_session_state_with_nested_value():
274+
instruction_template = "Hello {user.profile.name}, age {user.profile.age}."
275+
invocation_context = await _create_test_readonly_context(
276+
state={"user": {"profile": {"name": "Alice", "age": 30}}}
277+
)
278+
279+
populated_instruction = await instructions_utils.inject_session_state(
280+
instruction_template, invocation_context
281+
)
282+
assert populated_instruction == "Hello Alice, age 30."
283+
284+
285+
@pytest.mark.asyncio
286+
async def test_inject_session_state_with_nested_missing_raises_key_error():
287+
instruction_template = "Value: {user.missing.key}"
288+
invocation_context = await _create_test_readonly_context(
289+
state={"user": {"profile": {"name": "Alice"}}}
290+
)
291+
292+
with pytest.raises(
293+
KeyError, match="Context variable not found: `user.missing.key`."
294+
):
295+
await instructions_utils.inject_session_state(
296+
instruction_template, invocation_context
297+
)
298+
299+
300+
@pytest.mark.asyncio
301+
async def test_inject_session_state_with_nested_optional_missing():
302+
instruction_template = "Value: {user.missing.key?}"
303+
invocation_context = await _create_test_readonly_context(
304+
state={"user": {"profile": {"name": "Alice"}}}
305+
)
306+
307+
populated_instruction = await instructions_utils.inject_session_state(
308+
instruction_template, invocation_context
309+
)
310+
assert populated_instruction == "Value: "
311+
312+
313+
@pytest.mark.asyncio
314+
async def test_inject_session_state_nested_does_not_conflict_with_artifact():
315+
instruction_template = (
316+
"Name: {user.name}, Artifact: {artifact.my_file}"
317+
)
318+
mock_artifact_service = MockArtifactService(
319+
{"my_file": "artifact content"}
320+
)
321+
invocation_context = await _create_test_readonly_context(
322+
state={"user": {"name": "Bob"}},
323+
artifact_service=mock_artifact_service,
324+
)
325+
326+
populated_instruction = await instructions_utils.inject_session_state(
327+
instruction_template, invocation_context
328+
)
329+
assert populated_instruction == "Name: Bob, Artifact: artifact content"
330+
331+
332+
@pytest.mark.asyncio
333+
async def test_inject_session_state_flat_key_still_works():
334+
"""Flat keys still work even when nested is supported."""
335+
instruction_template = "Value: {simple_key}"
336+
invocation_context = await _create_test_readonly_context(
337+
state={"simple_key": "flat_value"}
338+
)
339+
340+
populated_instruction = await instructions_utils.inject_session_state(
341+
instruction_template, invocation_context
342+
)
343+
assert populated_instruction == "Value: flat_value"

0 commit comments

Comments
 (0)