chore: Make ArtifactService transparent to AgentTools

PiperOrigin-RevId: 767225493

Author: google-genai-bot

src/google/adk/tools/_forwarding_artifact_service.py

@@ -0,0 +1,96 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from google.genai import types
+from typing_extensions import override
+
+from ..artifacts.base_artifact_service import BaseArtifactService
+
+if TYPE_CHECKING:
+  from .tool_context import ToolContext
+
+
+class ForwardingArtifactService(BaseArtifactService):
+  """Artifact service that forwards to the parent tool context."""
+
+  def __init__(self, tool_context: ToolContext):
+    self.tool_context = tool_context
+    self._invocation_context = tool_context._invocation_context
+
+  @override
+  async def save_artifact(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      filename: str,
+      artifact: types.Part,
+  ) -> int:
+    return await self.tool_context.save_artifact(
+        filename=filename, artifact=artifact
+    )
+
+  @override
+  async def load_artifact(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      filename: str,
+      version: Optional[int] = None,
+  ) -> Optional[types.Part]:
+    return await self.tool_context.load_artifact(
+        filename=filename, version=version
+    )
+
+  @override
+  async def list_artifact_keys(
+      self, *, app_name: str, user_id: str, session_id: str
+  ) -> list[str]:
+    return await self.tool_context.list_artifacts()
+
+  @override
+  async def delete_artifact(
+      self, *, app_name: str, user_id: str, session_id: str, filename: str
+  ) -> None:
+    del app_name, user_id, session_id
+    if self._invocation_context.artifact_service is None:
+      raise ValueError("Artifact service is not initialized.")
+    await self._invocation_context.artifact_service.delete_artifact(
+        app_name=self._invocation_context.app_name,
+        user_id=self._invocation_context.user_id,
+        session_id=self._invocation_context.session.id,
+        filename=filename,
+    )
+
+  @override
+  async def list_versions(
+      self, *, app_name: str, user_id: str, session_id: str, filename: str
+  ) -> list[int]:
+    del app_name, user_id, session_id
+    if self._invocation_context.artifact_service is None:
+      raise ValueError("Artifact service is not initialized.")
+    return await self._invocation_context.artifact_service.list_versions(
+        app_name=self._invocation_context.app_name,
+        user_id=self._invocation_context.user_id,
+        session_id=self._invocation_context.session.id,
+        filename=filename,
+    )

src/google/adk/tools/agent_tool.py

@@ -25,6 +25,7 @@
 from ..memory.in_memory_memory_service import InMemoryMemoryService
 from ..runners import Runner
 from ..sessions.in_memory_session_service import InMemorySessionService
+from ._forwarding_artifact_service import ForwardingArtifactService
 from .base_tool import BaseTool
 from .tool_context import ToolContext
 
@@ -123,9 +124,7 @@ async def run_async(
     runner = Runner(
         app_name=self.agent.name,
         agent=self.agent,
-        # TODO(kech): Remove the access to the invocation context.
-        #   It seems we don't need re-use artifact_service if we forward below.
-        artifact_service=tool_context._invocation_context.artifact_service,
+        artifact_service=ForwardingArtifactService(tool_context),
         session_service=InMemorySessionService(),
         memory_service=InMemoryMemoryService(),
     )
@@ -144,24 +143,6 @@ async def run_async(
         tool_context.state.update(event.actions.state_delta)
       last_event = event
 
-    if runner.artifact_service:
-      # Forward all artifacts to parent session.
-      artifact_names = await runner.artifact_service.list_artifact_keys(
-          app_name=session.app_name,
-          user_id=session.user_id,
-          session_id=session.id,
-      )
-      for artifact_name in artifact_names:
-        if artifact := await runner.artifact_service.load_artifact(
-            app_name=session.app_name,
-            user_id=session.user_id,
-            session_id=session.id,
-            filename=artifact_name,
-        ):
-          await tool_context.save_artifact(
-              filename=artifact_name, artifact=artifact
-          )
-
     if not last_event or not last_event.content or not last_event.content.parts:
       return ''
     if isinstance(self.agent, LlmAgent) and self.agent.output_schema:

tests/unittests/tools/test_agent_tool.py

@@ -13,20 +13,15 @@
 # limitations under the License.
 
 from google.adk.agents import Agent
+from google.adk.agents import SequentialAgent
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.tools.agent_tool import AgentTool
 from google.genai.types import Part
 from pydantic import BaseModel
-import pytest
 from pytest import mark
 
 from .. import testing_utils
 
-pytestmark = pytest.mark.skip(
-    reason='Skipping until tool.func evaluations are fixed (async)'
-)
-
-
 function_call_custom = Part.from_function_call(
     name='tool_agent', args={'custom_input': 'test1'}
 )
@@ -112,6 +107,55 @@ def test_update_state():
   assert runner.session.state['state_1'] == 'changed_value'
 
 
+def test_update_artifacts():
+  """The agent tool can read and write artifacts."""
+
+  async def before_tool_agent(callback_context: CallbackContext):
+    # Artifact 1 should be available in the tool agent.
+    artifact = await callback_context.load_artifact('artifact_1')
+    await callback_context.save_artifact(
+        'artifact_2', Part.from_text(text=artifact.text + ' 2')
+    )
+
+  tool_agent = SequentialAgent(
+      name='tool_agent',
+      before_agent_callback=before_tool_agent,
+  )
+
+  async def before_main_agent(callback_context: CallbackContext):
+    await callback_context.save_artifact(
+        'artifact_1', Part.from_text(text='test')
+    )
+
+  async def after_main_agent(callback_context: CallbackContext):
+    # Artifact 2 should be available after the tool agent.
+    artifact_2 = await callback_context.load_artifact('artifact_2')
+    await callback_context.save_artifact(
+        'artifact_3', Part.from_text(text=artifact_2.text + ' 3')
+    )
+
+  mock_model = testing_utils.MockModel.create(
+      responses=[function_call_no_schema, 'response2']
+  )
+  root_agent = Agent(
+      name='root_agent',
+      before_agent_callback=before_main_agent,
+      after_agent_callback=after_main_agent,
+      tools=[AgentTool(agent=tool_agent)],
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent)
+  runner.run('test1')
+
+  artifacts_path = f'test_app/test_user/{runner.session_id}'
+  assert runner.runner.artifact_service.artifacts == {
+      f'{artifacts_path}/artifact_1': [Part.from_text(text='test')],
+      f'{artifacts_path}/artifact_2': [Part.from_text(text='test 2')],
+      f'{artifacts_path}/artifact_3': [Part.from_text(text='test 2 3')],
+  }
+
+
 @mark.parametrize(
     'env_variables',
     [