Skip to content

Commit bb4cfde

Browse files
seanzhougooglecopybara-github
authored andcommitted
fix: inject artifact into instructions
a. complain when artifact is None b. inject None value as empty string instead of `None` PiperOrigin-RevId: 800930613
1 parent e45c3be commit bb4cfde

File tree

2 files changed

+76
-7
lines changed

2 files changed

+76
-7
lines changed

src/google/adk/utils/instructions_utils.py

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
from __future__ import annotations
1616

17+
import logging
1718
import re
1819

1920
from ..agents.readonly_context import ReadonlyContext
@@ -23,6 +24,8 @@
2324
'inject_session_state',
2425
]
2526

27+
logger = logging.getLogger('google_adk.' + __name__)
28+
2629

2730
async def inject_session_state(
2831
template: str,
@@ -91,16 +94,29 @@ async def _replace_match(match) -> str:
9194
session_id=invocation_context.session.id,
9295
filename=var_name,
9396
)
94-
if not var_name:
95-
raise KeyError(f'Artifact {var_name} not found.')
97+
if artifact is None:
98+
if optional:
99+
logger.debug(
100+
'Artifact %s not found, replacing with empty string', var_name
101+
)
102+
return ''
103+
else:
104+
raise KeyError(f'Artifact {var_name} not found.')
96105
return str(artifact)
97106
else:
98107
if not _is_valid_state_name(var_name):
99108
return match.group()
100109
if var_name in invocation_context.session.state:
101-
return str(invocation_context.session.state[var_name])
110+
value = invocation_context.session.state[var_name]
111+
if value is None:
112+
return ''
113+
return str(value)
102114
else:
103115
if optional:
116+
logger.debug(
117+
'Context variable %s not found, replacing with empty string',
118+
var_name,
119+
)
104120
return ''
105121
else:
106122
raise KeyError(f'Context variable not found: `{var_name}`.')

tests/unittests/utils/test_instructions_utils.py

Lines changed: 57 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,17 @@
1-
from google.adk.agents.invocation_context import InvocationContext
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
215
from google.adk.agents.llm_agent import Agent
316
from google.adk.agents.readonly_context import ReadonlyContext
417
from google.adk.sessions.session import Session
@@ -17,7 +30,7 @@ async def load_artifact(self, app_name, user_id, session_id, filename):
1730
if filename in self.artifacts:
1831
return self.artifacts[filename]
1932
else:
20-
raise KeyError(f"Artifact '{filename}' not found.")
33+
return None
2134

2235

2336
async def _create_test_readonly_context(
@@ -114,7 +127,7 @@ async def test_inject_session_state_with_missing_artifact_raises_key_error():
114127
artifact_service=mock_artifact_service
115128
)
116129

117-
with pytest.raises(KeyError, match="Artifact 'missing_file' not found."):
130+
with pytest.raises(KeyError, match="Artifact missing_file not found."):
118131
await instructions_utils.inject_session_state(
119132
instruction_template, invocation_context
120133
)
@@ -200,7 +213,7 @@ async def test_inject_session_state_with_empty_artifact_name_raises_key_error():
200213
artifact_service=mock_artifact_service
201214
)
202215

203-
with pytest.raises(KeyError, match="Artifact '' not found."):
216+
with pytest.raises(KeyError, match="Artifact not found."):
204217
await instructions_utils.inject_session_state(
205218
instruction_template, invocation_context
206219
)
@@ -214,3 +227,43 @@ async def test_inject_session_state_artifact_service_not_initialized_raises_valu
214227
await instructions_utils.inject_session_state(
215228
instruction_template, invocation_context
216229
)
230+
231+
232+
@pytest.mark.asyncio
233+
async def test_inject_session_state_with_optional_missing_artifact_returns_empty():
234+
instruction_template = "Optional artifact: {artifact.missing_file?}"
235+
mock_artifact_service = MockArtifactService(
236+
{"my_file": "This is my artifact content."}
237+
)
238+
invocation_context = await _create_test_readonly_context(
239+
artifact_service=mock_artifact_service
240+
)
241+
242+
populated_instruction = await instructions_utils.inject_session_state(
243+
instruction_template, invocation_context
244+
)
245+
assert populated_instruction == "Optional artifact: "
246+
247+
248+
@pytest.mark.asyncio
249+
async def test_inject_session_state_with_none_state_value_returns_empty():
250+
instruction_template = "Value: {test_key}"
251+
invocation_context = await _create_test_readonly_context(
252+
state={"test_key": None}
253+
)
254+
255+
populated_instruction = await instructions_utils.inject_session_state(
256+
instruction_template, invocation_context
257+
)
258+
assert populated_instruction == "Value: "
259+
260+
261+
@pytest.mark.asyncio
262+
async def test_inject_session_state_with_optional_missing_state_returns_empty():
263+
instruction_template = "Optional value: {missing_key?}"
264+
invocation_context = await _create_test_readonly_context()
265+
266+
populated_instruction = await instructions_utils.inject_session_state(
267+
instruction_template, invocation_context
268+
)
269+
assert populated_instruction == "Optional value: "

0 commit comments

Comments
 (0)