From 34e8ee7340d10fbed2d5349b065b7a4f0088e8b0 Mon Sep 17 00:00:00 2001 From: HughChaw <146055770+Hughhhhcoder@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:18:21 +0800 Subject: [PATCH] fix: filter hidden artifacts from LLM context --- src/google/adk/tools/load_artifacts_tool.py | 49 ++++++++++-- .../tools/test_load_artifacts_tool.py | 76 +++++++++++++++++++ 2 files changed, 120 insertions(+), 5 deletions(-) diff --git a/src/google/adk/tools/load_artifacts_tool.py b/src/google/adk/tools/load_artifacts_tool.py index c193aa280f..18b9db46bb 100644 --- a/src/google/adk/tools/load_artifacts_tool.py +++ b/src/google/adk/tools/load_artifacts_tool.py @@ -295,6 +295,10 @@ def as_safe_part_for_llm( [types.Part, str], types.Part | None | Awaitable[types.Part | None], ] +ProcessArtifactNamesCallback: TypeAlias = Callable[ + [list[str]], + list[str] | Awaitable[list[str]], +] class LoadArtifactsTool(BaseTool): @@ -304,6 +308,7 @@ def __init__( self, *, process_artifact: ProcessArtifactCallback | None = None, + process_artifact_names: ProcessArtifactNamesCallback | None = None, enable_spreadsheet_parsing: bool = False, ): """Initializes the tool. @@ -315,11 +320,17 @@ def __init__( filtered before being added to the LLM request. If `None` (default), the built-in safety conversion (`as_safe_part_for_llm`) is used to convert unsupported formats (e.g., extracting text from DOCX/CSV/JSON/plain text - or replacing binary data with safe placeholder descriptions). If a - custom function is supplied, it bypasses default safety conversions; + or replacing binary data with safe placeholder descriptions). If a custom + function is supplied, it bypasses default safety conversions; returning `None` skips the artifact so it is omitted from the request. If a custom callback raises an exception, the error is logged and the artifact is skipped. + process_artifact_names: An optional sync or async callable with signature + `(artifact_names: list[str]) -> list[str]`. The returned names are + exposed to the LLM and are the only names the tool will load. This can + be used to hide artifacts that are intended only for control code. If + the callback raises an exception, the error is logged and no artifacts + are exposed for the request. enable_spreadsheet_parsing: Whether to enable spreadsheet parsing files (e.g., .xlsx, .xls) into text. Defaults to False. """ @@ -331,6 +342,9 @@ def __init__( web UI)."""), ) self._process_artifact: ProcessArtifactCallback | None = process_artifact + self._process_artifact_names: ProcessArtifactNamesCallback | None = ( + process_artifact_names + ) self._enable_spreadsheet_parsing: bool = enable_spreadsheet_parsing def _get_declaration(self) -> types.FunctionDeclaration | None: @@ -369,6 +383,11 @@ async def run_async( self, *, args: dict[str, Any], tool_context: ToolContext ) -> Any: artifact_names: list[str] = args.get('artifact_names', []) + if self._process_artifact_names is not None: + allowed_artifact_names = await self._get_artifact_names(tool_context) + artifact_names = [ + name for name in artifact_names if name in allowed_artifact_names + ] return { 'artifact_names': artifact_names, 'status': ( @@ -377,6 +396,20 @@ async def run_async( ), } + async def _get_artifact_names(self, tool_context: ToolContext) -> list[str]: + artifact_names = await tool_context.list_artifacts() + if self._process_artifact_names is None: + return artifact_names + + try: + processed_names = self._process_artifact_names(list(artifact_names)) + if inspect.isawaitable(processed_names): + processed_names = await processed_names + return processed_names + except Exception: # pylint: disable=broad-exception-caught + logger.exception('Failed to process artifact names, skipping.') + return [] + @override async def process_llm_request( self, *, tool_context: ToolContext, llm_request: LlmRequest @@ -392,7 +425,7 @@ async def process_llm_request( async def _append_artifacts_to_llm_request( self, *, tool_context: ToolContext, llm_request: LlmRequest ): - artifact_names = await tool_context.list_artifacts() + artifact_names = await self._get_artifact_names(tool_context) if not artifact_names: return @@ -413,8 +446,14 @@ async def _append_artifacts_to_llm_request( function_response = llm_request.contents[-1].parts[0].function_response if function_response and function_response.name == 'load_artifacts': response = function_response.response or {} - artifact_names = response.get('artifact_names', []) - for artifact_name in artifact_names: + requested_artifact_names = response.get('artifact_names', []) + if self._process_artifact_names is not None: + requested_artifact_names = [ + name + for name in requested_artifact_names + if name in artifact_names + ] + for artifact_name in requested_artifact_names: # Try session-scoped first (default behavior) artifact = await tool_context.load_artifact(artifact_name) diff --git a/tests/unittests/tools/test_load_artifacts_tool.py b/tests/unittests/tools/test_load_artifacts_tool.py index a12675dd4e..fa1e25fe09 100644 --- a/tests/unittests/tools/test_load_artifacts_tool.py +++ b/tests/unittests/tools/test_load_artifacts_tool.py @@ -15,6 +15,7 @@ import base64 import io from typing import Any +from typing import cast from unittest import mock import zipfile @@ -24,6 +25,7 @@ from google.adk.tools.load_artifacts_tool import _maybe_base64_to_bytes from google.adk.tools.load_artifacts_tool import load_artifacts_tool from google.adk.tools.load_artifacts_tool import LoadArtifactsTool +from google.adk.tools.tool_context import ToolContext from google.genai import types import pandas as pd import pytest @@ -476,6 +478,80 @@ async def test_load_artifacts_registers_dynamic_instructions(): assert len(llm_request.contents) == 0 +@pytest.mark.asyncio +async def test_load_artifacts_hides_filtered_names_and_contents(): + """Filtered artifacts are omitted from both instructions and LLM contents.""" + visible_name = 'visible.txt' + hidden_name = 'internal.txt' + + def filter_artifact_names(artifact_names: list[str]) -> list[str]: + return [name for name in artifact_names if name == visible_name] + + tool = LoadArtifactsTool(process_artifact_names=filter_artifact_names) + tool_context = _StubToolContext({ + visible_name: types.Part.from_text(text='visible content'), + hidden_name: types.Part.from_text(text='hidden content'), + }) + tool_response = await tool.run_async( + args={'artifact_names': [visible_name, hidden_name]}, + tool_context=cast(ToolContext, tool_context), + ) + assert tool_response['artifact_names'] == [visible_name] + llm_request = LlmRequest( + contents=[ + types.Content( + role='user', + parts=[ + types.Part( + function_response=types.FunctionResponse( + name='load_artifacts', + response=tool_response, + ) + ) + ], + ) + ] + ) + + await tool.process_llm_request( + tool_context=tool_context, llm_request=llm_request + ) + + instruction = llm_request._dynamic_instructions[0] + assert visible_name in instruction + assert hidden_name not in instruction + assert len(llm_request.contents) == 2 + assert llm_request.contents[-1].parts[0].text == ( + f'Artifact {visible_name} is:' + ) + assert llm_request.contents[-1].parts[1].text == 'visible content' + + +@pytest.mark.asyncio +async def test_load_artifacts_accepts_async_process_artifact_names(): + """Async artifact-name callbacks filter names before LLM context assembly.""" + artifact_names_seen = [] + + async def filter_artifact_names(artifact_names: list[str]) -> list[str]: + artifact_names_seen.extend(artifact_names) + return ['visible.txt'] + + tool = LoadArtifactsTool(process_artifact_names=filter_artifact_names) + tool_context = _StubToolContext({ + 'visible.txt': types.Part.from_text(text='visible content'), + 'internal.txt': types.Part.from_text(text='hidden content'), + }) + llm_request = LlmRequest() + + await tool.process_llm_request( + tool_context=tool_context, llm_request=llm_request + ) + + assert artifact_names_seen == ['visible.txt', 'internal.txt'] + assert 'visible.txt' in llm_request._dynamic_instructions[0] + assert 'internal.txt' not in llm_request._dynamic_instructions[0] + + def test_load_artifacts_tool_keyword_only(): """process_artifact must be passed as keyword argument.""" with pytest.raises(TypeError):