Merge branch 'main' into sid/evaluate-full-response

Author: ankursharmas

src/google/adk/plugins/save_files_as_artifacts_plugin.py

@@ -47,13 +47,23 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
   tool to the agent, or load the artifacts in your own tool to use the files.
   """
 
-  def __init__(self, name: str = 'save_files_as_artifacts_plugin'):
+  def __init__(
+      self,
+      name: str = 'save_files_as_artifacts_plugin',
+      *,
+      attach_file_reference: bool = True,
+  ):
     """Initialize the save files as artifacts plugin.
 
     Args:
       name: The name of the plugin instance.
+      attach_file_reference: Whether to attach a file reference to the
+        user message. If False, only save the files as artifacts without
+        adding a file reference, and the files will not be directly
+        accessible to the model.
     """
     super().__init__(name)
+    self._attach_file_reference = attach_file_reference
 
   async def on_user_message_callback(
       self,
@@ -108,15 +118,16 @@ async def on_user_message_callback(
         )
         new_parts.append(placeholder_part)
 
-        file_part = await self._build_file_reference_part(
-            invocation_context=invocation_context,
-            filename=file_name,
-            version=version,
-            mime_type=inline_data.mime_type,
-            display_name=display_name,
-        )
-        if file_part:
-          new_parts.append(file_part)
+        if self._attach_file_reference:
+          file_part = await self._build_file_reference_part(
+              invocation_context=invocation_context,
+              filename=file_name,
+              version=version,
+              mime_type=inline_data.mime_type,
+              display_name=display_name,
+          )
+          if file_part:
+            new_parts.append(file_part)
         pending_delta[file_name] = version
 
         modified = True

tests/unittests/plugins/test_save_files_as_artifacts.py

@@ -91,6 +91,36 @@ async def test_save_files_with_display_name(self):
     assert result.parts[1].file_data.display_name == "test_document.pdf"
     assert result.parts[1].file_data.mime_type == "application/pdf"
 
+  @pytest.mark.asyncio
+  async def test_attach_file_reference_false(self):
+    """Test that file reference is not attached when attach_file_reference is False."""
+    plugin = SaveFilesAsArtifactsPlugin(attach_file_reference=False)
+
+    inline_data = types.Blob(
+        display_name="test_document.pdf",
+        data=b"test data",
+        mime_type="application/pdf",
+    )
+
+    original_part = types.Part(inline_data=inline_data)
+    user_message = types.Content(parts=[original_part])
+
+    result = await plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
+
+    self.mock_context.artifact_service.save_artifact.assert_called_once_with(
+        app_name="test_app",
+        user_id="test_user",
+        session_id="test_session",
+        filename="test_document.pdf",
+        artifact=original_part,
+    )
+
+    assert result
+    assert len(result.parts) == 1
+    assert result.parts[0].text == '[Uploaded Artifact: "test_document.pdf"]'
+
   @pytest.mark.asyncio
   async def test_save_files_without_display_name(self):
     """Test saving files when inline_data has no display_name."""