Revert "Patch: Fix file size limit, MIME type uses library (google#3781)"

This reverts commit aa5d6a4.

Author: AakashSuresh2003

src/google/adk/plugins/save_files_as_artifacts_plugin.py

@@ -26,7 +26,6 @@
 
 from ..agents.invocation_context import InvocationContext
 from .base_plugin import BasePlugin
-import mimetypes
 
 logger = logging.getLogger('google_adk.' + __name__)
 
@@ -109,10 +108,9 @@ async def on_user_message_callback(
         # Check if file exceeds Files API limit (2GB)
         if file_size > _MAX_FILES_API_SIZE_BYTES:
           file_size_gb = file_size / (1024 * 1024 * 1024)
-          max_size_gb = _MAX_FILES_API_SIZE_BYTES / (1024 * 1024 * 1024)
           error_message = (
               f'File {display_name} ({file_size_gb:.2f} GB) exceeds the'
-              f' maximum supported size of {max_size_gb:.2f} GB. Please upload a smaller file.'
+              ' maximum supported size of 2GB. Please upload a smaller file.'
           )
           logger.warning(error_message)
           new_parts.append(types.Part(text=f'[Upload Error: {error_message}]'))
@@ -217,7 +215,15 @@ async def _upload_to_files_api(
         file_extension = os.path.splitext(inline_data.display_name)[1]
       elif inline_data.mime_type:
         # Simple mime type to extension mapping
-        file_extension = mimetypes.guess_extension(inline_data.mime_type) or ''
+        mime_to_ext = {
+            'application/pdf': '.pdf',
+            'image/png': '.png',
+            'image/jpeg': '.jpg',
+            'image/gif': '.gif',
+            'text/plain': '.txt',
+            'application/json': '.json',
+        }
+        file_extension = mime_to_ext.get(inline_data.mime_type, '')
 
       # Create temporary file
       with tempfile.NamedTemporaryFile(

tests/unittests/plugins/test_save_files_as_artifacts.py

@@ -497,33 +497,60 @@ async def test_files_api_upload_failure(self):
       assert "huge_file.pdf" in result.parts[0].text
       assert "API quota exceeded" in result.parts[0].text
 
-    @pytest.mark.asyncio
-    async def test_file_exceeds_files_api_limit(self):
-        """Test that files exceeding 2GB limit are rejected with clear error."""
-        # Use a small file for the test
-        large_data = b"x" * 1000
-        inline_data = types.Blob(
-            display_name="huge_video.mp4",
-            data=large_data,
-            mime_type="video/mp4",
-        )
-        user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
-
-        # Patch the size limit to be smaller than the file data
-        with patch(
-            "google.adk.plugins.save_files_as_artifacts_plugin._MAX_FILES_API_SIZE_BYTES", 500
-        ):
-            result = await self.plugin.on_user_message_callback(
-                invocation_context=self.mock_context, user_message=user_message
+  @pytest.mark.asyncio
+  async def test_file_exceeds_files_api_limit(self):
+    """Test that files exceeding 2GB limit are rejected with clear error."""
+    # Create a file larger than 2GB (simulated with a descriptor that reports large size)
+    # Create a mock object that behaves like bytes but reports 2GB+ size
+    large_data = b"x" * 1000  # Small actual data for testing
+
+    # Create inline_data with the small data
+    inline_data = types.Blob(
+        display_name="huge_video.mp4",
+        data=large_data,
+        mime_type="video/mp4",
+    )
+
+    user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
+
+    # Patch the file size check to simulate a 2GB+ file
+    original_callback = self.plugin.on_user_message_callback
+
+    async def patched_callback(*, invocation_context, user_message):
+      # Temporarily replace the data length check
+      for part in user_message.parts:
+        if part.inline_data:
+          # Simulate 2GB + 1 byte size
+          file_size_over_limit = (2 * 1024 * 1024 * 1024) + 1
+          # Manually inject the check that would happen in the real code
+          if file_size_over_limit > (2 * 1024 * 1024 * 1024):
+            file_size_gb = file_size_over_limit / (1024 * 1024 * 1024)
+            display_name = part.inline_data.display_name or "unknown"
+            error_message = (
+                f"File {display_name} ({file_size_gb:.2f} GB) exceeds the"
+                " maximum supported size of 2GB. Please upload a smaller file."
+            )
+            return types.Content(
+                role="user",
+                parts=[types.Part(text=f"[Upload Error: {error_message}]")],
             )
+      return await original_callback(
+          invocation_context=invocation_context, user_message=user_message
+      )
+
+    self.plugin.on_user_message_callback = patched_callback
+
+    result = await self.plugin.on_user_message_callback(
+        invocation_context=self.mock_context, user_message=user_message
+    )
 
-        # Should not attempt any upload
-        self.mock_context.artifact_service.save_artifact.assert_not_called()
+    # Should not attempt any upload
+    self.mock_context.artifact_service.save_artifact.assert_not_called()
 
-        # Should return error message about the limit
-        assert result is not None
-        assert len(result.parts) == 1
-        assert "[Upload Error:" in result.parts[0].text
-        assert "huge_video.mp4" in result.parts[0].text
-        # Note: This assertion will depend on fixing the hardcoded "2GB" in the error message.
-        assert "exceeds the maximum supported size" in result.parts[0].text
+    # Should return error message about 2GB limit
+    assert result is not None
+    assert len(result.parts) == 1
+    assert "[Upload Error:" in result.parts[0].text
+    assert "huge_video.mp4" in result.parts[0].text
+    assert "2.00 GB" in result.parts[0].text
+    assert "exceeds the maximum supported size" in result.parts[0].text