-
Notifications
You must be signed in to change notification settings - Fork 3.3k
feat: support file_data URI references in GcsArtifactService #5322
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
f004631
e865bac
05151d4
5425cab
5b315e7
6ee4a32
cd7a755
35a4e48
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 | ||
|
|
@@ -230,9 +231,21 @@ def _save_artifact( | |
| content_type="text/plain", | ||
| ) | ||
| elif artifact.file_data: | ||
| raise NotImplementedError( | ||
| "Saving artifact with file_data is not supported yet in" | ||
| " GcsArtifactService." | ||
| if not artifact.file_data.file_uri: | ||
| raise InputValidationError("Artifact file_data must have a file_uri.") | ||
| if artifact_util.is_artifact_ref(artifact): | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It is good practice to leverage the utility functions in artifact_util to validate internal artifact:// URIs before attempting to store them |
||
| if not artifact_util.parse_artifact_uri(artifact.file_data.file_uri): | ||
| raise InputValidationError( | ||
| f"Invalid artifact reference URI: {artifact.file_data.file_uri}" | ||
| ) | ||
| # Store the URI as blob metadata; no content to upload. | ||
| blob.metadata = { | ||
| **(blob.metadata or {}), | ||
| "file_uri": artifact.file_data.file_uri, | ||
| } | ||
| blob.upload_from_string( | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Uploading an empty string as content for URI references is a good approach. It allows the blob to exist as a container for metadata in GCS without incurring significant storage costs |
||
| b"", | ||
| content_type=artifact.file_data.mime_type or None, | ||
| ) | ||
| else: | ||
| raise InputValidationError( | ||
|
|
@@ -263,15 +276,25 @@ def _load_artifact( | |
| blob_name = self._get_blob_name( | ||
| app_name, user_id, filename, version, session_id | ||
| ) | ||
| blob = self.bucket.blob(blob_name) | ||
| blob = self.bucket.get_blob(blob_name) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Changing from bucket.blob() to bucket.get_blob() is necessary here to fetch the blob's metadata (including the file_uri). While this adds an extra RPC per load operation, it is a justified trade-off for supporting the new URI reference feature |
||
| if blob is None: | ||
| return None | ||
|
|
||
| # If the artifact was saved as a file_data URI reference, restore it. | ||
| if blob.metadata and "file_uri" in blob.metadata: | ||
| return types.Part( | ||
| file_data=types.FileData( | ||
| file_uri=blob.metadata["file_uri"], | ||
| mime_type=blob.content_type or None, | ||
| ) | ||
| ) | ||
|
|
||
| artifact_bytes = blob.download_as_bytes() | ||
| if not artifact_bytes: | ||
| return None | ||
| artifact = types.Part.from_bytes( | ||
| return types.Part.from_bytes( | ||
| data=artifact_bytes, mime_type=blob.content_type | ||
| ) | ||
| return artifact | ||
|
|
||
| def _list_artifact_keys( | ||
| self, app_name: str, user_id: str, session_id: Optional[str] | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You are accessing artifact.file_data.file_uri multiple times here (for the null check, validation, and metadata storage). Consider assigning it to a local variable at the start of the elif artifact.file_data: block to improve readability and avoid repeated attribute access