feat: support file_data URI references in GcsArtifactService

Merge #5322

`GcsArtifactService` previously raised `NotImplementedError` when saving a `types.Part` with `file_data`. This PR adds full `file_data` URI support to match the existing `InMemoryArtifactService` behavior.

### Link to Issue or Description of Change

- Closes: #5230

**Problem:**
`GcsArtifactService._save_artifact` raised `NotImplementedError` for `file_data` parts, while `InMemoryArtifactService` already handled them. This made `GcsArtifactService` unusable for any workflow that passes file URI references (e.g. `gs://` or `artifact://` URIs) as artifacts.

**Solution:**
Store `file_data` URI references as zero-byte GCS blobs with a `file_uri` custom metadata key — no content to upload, just the URI pointer. On load, check for that metadata key before downloading bytes and restore the original `types.Part(file_data=...)`. Also switched `_load_artifact` to use `bucket.get_blob()` (which populates metadata) instead of `bucket.blob()` (which does not). `artifact://` URIs are validated via `artifact_util.parse_artifact_uri` before saving.

### Testing Plan

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

Added 4 new tests in `tests/unittests/artifacts/test_artifact_service.py`:

- `test_gcs_save_artifact_with_external_gcs_uri` — roundtrip save/load with a `gs://` URI
- `test_gcs_save_artifact_with_artifact_ref_uri` — roundtrip save/load with an `artifact://` URI
- `test_gcs_save_artifact_file_data_without_mime_type` — `file_data` with no mime_type
- `test_gcs_save_artifact_file_data_missing_uri_raises` — empty `file_uri` raises `InputValidationError`

Co-authored-by: Shangjie Chen <deanchen@google.com>
COPYBARA_INTEGRATE_REVIEW=#5322 from blueeye-040:feat/gcs-file-data-support 99a932d
PiperOrigin-RevId: 934629799

Author: 2 people

src/google/adk/artifacts/gcs_artifact_service.py

@@ -32,6 +32,7 @@
 from google.genai import types
 from typing_extensions import override
 
+from . import artifact_util
 from ..errors.input_validation_error import InputValidationError
 from .base_artifact_service import ArtifactVersion
 from .base_artifact_service import BaseArtifactService
@@ -41,6 +42,8 @@
 
 _GCS_DISPLAY_NAME_METADATA_KEY = "adkDisplayName"
 _GCS_IS_TEXT_METADATA_KEY = "adkIsText"
+_GCS_FILE_URI_METADATA_KEY = "adkFileUri"
+_GCS_FILE_MIME_TYPE_METADATA_KEY = "adkFileMimeType"
 
 
 class GcsArtifactService(BaseArtifactService):
@@ -243,9 +246,27 @@ def _save_artifact(
           content_type="text/plain",
       )
     elif artifact.file_data:
-      raise NotImplementedError(
-          "Saving artifact with file_data is not supported yet in"
-          " GcsArtifactService."
+      file_data = artifact.file_data
+      assert file_data is not None
+      file_uri = file_data.file_uri
+      if not file_uri:
+        raise InputValidationError("Artifact file_data must have a file_uri.")
+      if artifact_util.is_artifact_ref(artifact):
+        if not artifact_util.parse_artifact_uri(file_uri):
+          raise InputValidationError(
+              f"Invalid artifact reference URI: {file_uri}"
+          )
+      # Store the URI and mime_type (if any) as blob metadata; no content to upload.
+      metadata = {
+          **(blob.metadata or {}),
+          _GCS_FILE_URI_METADATA_KEY: file_uri,
+      }
+      if file_data.mime_type:
+        metadata[_GCS_FILE_MIME_TYPE_METADATA_KEY] = file_data.mime_type
+      blob.metadata = metadata
+      blob.upload_from_string(
+          b"",
+          content_type=file_data.mime_type or None,
       )
     else:
       raise InputValidationError(
@@ -280,6 +301,39 @@ def _load_artifact(
     if not blob:
       return None
 
+    # If the artifact was saved as a file_data URI reference, restore or resolve it.
+    file_uri = None
+    if blob.metadata:
+      file_uri = blob.metadata.get(
+          _GCS_FILE_URI_METADATA_KEY
+      ) or blob.metadata.get("file_uri")
+
+    if file_uri:
+      if file_uri.startswith("artifact://"):
+        parsed_uri = artifact_util.parse_artifact_uri(file_uri)
+        if not parsed_uri:
+          raise InputValidationError(
+              f"Invalid artifact reference URI: {file_uri}"
+          )
+        return self._load_artifact(
+            app_name=parsed_uri.app_name,
+            user_id=parsed_uri.user_id,
+            session_id=parsed_uri.session_id,
+            filename=parsed_uri.filename,
+            version=parsed_uri.version,
+        )
+      mime_type = None
+      if blob.metadata:
+        mime_type = blob.metadata.get(_GCS_FILE_MIME_TYPE_METADATA_KEY)
+      if mime_type is None:
+        mime_type = blob.content_type or None
+      return types.Part(
+          file_data=types.FileData(
+              file_uri=file_uri,
+              mime_type=mime_type,
+          )
+      )
+
     artifact_bytes = blob.download_as_bytes()
     if blob.metadata and blob.metadata.get(_GCS_IS_TEXT_METADATA_KEY) == "true":
       return types.Part(text=artifact_bytes.decode("utf-8"))

src/google/adk/errors/input_validation_error.py

@@ -18,7 +18,7 @@
 class InputValidationError(ValueError):
   """Represents an error raised when user input fails validation."""
 
-  def __init__(self, message="Invalid input."):
+  def __init__(self, message: str = "Invalid input.") -> None:
     """Initializes the InputValidationError exception.
 
     Args:

tests/unittests/artifacts/test_artifact_service.py

@@ -938,6 +938,201 @@ def test_converts_text_dict(self):
     assert result.text == "hello world"
 
 
+# ---------------------------------------------------------------------------
+# GCS file_data (URI reference) tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio  # type: ignore[untyped-decorator]
+async def test_gcs_save_artifact_with_external_gcs_uri() -> None:
+  """GcsArtifactService saves and loads a gs:// file_data URI reference."""
+  service = mock_gcs_artifact_service()  # type: ignore[no-untyped-call]
+  artifact = types.Part(
+      file_data=types.FileData(
+          file_uri="gs://my-bucket/report.pdf",
+          mime_type="application/pdf",
+      )
+  )
+
+  version = await service.save_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="report.pdf",
+      artifact=artifact,
+  )
+  assert version == 0
+
+  loaded = await service.load_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="report.pdf",
+  )
+  assert loaded is not None
+  assert loaded.file_data is not None
+  assert loaded.file_data.file_uri == "gs://my-bucket/report.pdf"
+  assert loaded.file_data.mime_type == "application/pdf"
+
+
+@pytest.mark.asyncio  # type: ignore[untyped-decorator]
+async def test_gcs_save_artifact_with_artifact_ref_uri() -> None:
+  """GcsArtifactService saves and recursively loads an internal artifact:// URI reference."""
+  service = mock_gcs_artifact_service()  # type: ignore[no-untyped-call]
+
+  # Save the referenced (source) artifact first.
+  source_artifact = types.Part(text="source content")
+  await service.save_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="source.txt",
+      artifact=source_artifact,
+  )
+
+  artifact_ref_uri = "artifact://apps/app/users/user1/sessions/sess1/artifacts/source.txt/versions/0"
+  artifact = types.Part(
+      file_data=types.FileData(
+          file_uri=artifact_ref_uri,
+          mime_type="text/plain",
+      )
+  )
+
+  version = await service.save_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="ref.txt",
+      artifact=artifact,
+  )
+  assert version == 0
+
+  loaded = await service.load_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="ref.txt",
+  )
+  assert loaded is not None
+  assert loaded.text == "source content"
+
+
+@pytest.mark.asyncio  # type: ignore[untyped-decorator]
+async def test_gcs_save_artifact_file_data_without_mime_type() -> None:
+  """GcsArtifactService handles file_data with no mime_type."""
+  service = mock_gcs_artifact_service()  # type: ignore[no-untyped-call]
+  artifact = types.Part(
+      file_data=types.FileData(file_uri="gs://my-bucket/data.bin")
+  )
+
+  version = await service.save_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="data.bin",
+      artifact=artifact,
+  )
+  assert version == 0
+
+  loaded = await service.load_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="data.bin",
+  )
+  assert loaded is not None
+  assert loaded.file_data is not None
+  assert loaded.file_data.file_uri == "gs://my-bucket/data.bin"
+
+
+@pytest.mark.asyncio  # type: ignore[untyped-decorator]
+async def test_gcs_save_artifact_file_data_missing_uri_raises() -> None:
+  """GcsArtifactService raises InputValidationError when file_uri is empty."""
+  service = mock_gcs_artifact_service()  # type: ignore[no-untyped-call]
+  artifact = types.Part(file_data=types.FileData(file_uri=""))
+
+  with pytest.raises(InputValidationError):
+    await service.save_artifact(
+        app_name="app",
+        user_id="user1",
+        session_id="sess1",
+        filename="empty.bin",
+        artifact=artifact,
+    )
+
+
+@pytest.mark.asyncio  # type: ignore[untyped-decorator]
+async def test_gcs_save_artifact_file_data_invalid_uri_raises() -> None:
+  """GcsArtifactService raises InputValidationError when file_uri is an invalid artifact:// URI template."""
+  service = mock_gcs_artifact_service()  # type: ignore[no-untyped-call]
+  artifact = types.Part(
+      file_data=types.FileData(
+          file_uri="artifact://apps/app/invalid",
+          mime_type="text/plain",
+      )
+  )
+
+  with pytest.raises(InputValidationError):
+    await service.save_artifact(
+        app_name="app",
+        user_id="user1",
+        session_id="sess1",
+        filename="invalid_ref.txt",
+        artifact=artifact,
+    )
+
+
+@pytest.mark.asyncio  # type: ignore[untyped-decorator]
+async def test_gcs_save_artifact_metadata_namespacing_and_mime() -> None:
+  """GcsArtifactService saves file_data using namespaced metadata keys."""
+  service = mock_gcs_artifact_service()  # type: ignore[no-untyped-call]
+  artifact = types.Part(
+      file_data=types.FileData(
+          file_uri="gs://my-bucket/report.pdf",
+          mime_type="application/pdf",
+      )
+  )
+
+  await service.save_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="report.pdf",
+      artifact=artifact,
+  )
+
+  blob_name = service._get_blob_name("app", "user1", "report.pdf", 0, "sess1")
+  blob = service.bucket.get_blob(blob_name)
+  assert blob is not None
+  assert blob.metadata.get("adkFileUri") == "gs://my-bucket/report.pdf"
+  assert blob.metadata.get("adkFileMimeType") == "application/pdf"
+  assert "file_uri" not in blob.metadata
+
+
+@pytest.mark.asyncio  # type: ignore[untyped-decorator]
+async def test_gcs_load_artifact_file_data_fallback_compatibility() -> None:
+  """GcsArtifactService loads file_data with old file_uri metadata key for backward compatibility."""
+  service = mock_gcs_artifact_service()  # type: ignore[no-untyped-call]
+  blob_name = service._get_blob_name(
+      "app", "user1", "old_report.pdf", 0, "sess1"
+  )
+  blob = service.bucket.blob(blob_name)
+  # Manually setup metadata with old key
+  blob.metadata = {"file_uri": "gs://my-bucket/old_report.pdf"}
+  blob.upload_from_string(b"", content_type="application/pdf")
+
+  loaded = await service.load_artifact(
+      app_name="app",
+      user_id="user1",
+      session_id="sess1",
+      filename="old_report.pdf",
+  )
+  assert loaded is not None
+  assert loaded.file_data is not None
+  assert loaded.file_data.file_uri == "gs://my-bucket/old_report.pdf"
+  assert loaded.file_data.mime_type == "application/pdf"
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "service_type",