From 0e3c0f78f5ac6e3a6bd946b257a854e66774cbd7 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Wed, 8 Oct 2025 11:08:07 -0700 Subject: [PATCH] feat: Make `session_id` optional in `BaseArtifactService` methods PiperOrigin-RevId: 816782982 --- .../adk/artifacts/base_artifact_service.py | 37 ++++++--- .../adk/artifacts/gcs_artifact_service.py | 77 ++++++++++++------- .../artifacts/in_memory_artifact_service.py | 47 +++++++---- .../adk/tools/_forwarding_artifact_service.py | 20 +++-- 4 files changed, 126 insertions(+), 55 deletions(-) diff --git a/src/google/adk/artifacts/base_artifact_service.py b/src/google/adk/artifacts/base_artifact_service.py index 249df966..5fc67ae7 100644 --- a/src/google/adk/artifacts/base_artifact_service.py +++ b/src/google/adk/artifacts/base_artifact_service.py @@ -11,7 +11,7 @@ # 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 abc import ABC from abc import abstractmethod @@ -29,9 +29,9 @@ class BaseArtifactService(ABC): *, app_name: str, user_id: str, - session_id: str, filename: str, artifact: types.Part, + session_id: Optional[str] = None, ) -> int: """Saves an artifact to the artifact service storage. @@ -42,9 +42,9 @@ class BaseArtifactService(ABC): Args: app_name: The app name. user_id: The user ID. - session_id: The session ID. filename: The filename of the artifact. artifact: The artifact to save. + session_id: The session ID. If `None`, the artifact is user-scoped. Returns: The revision ID. The first version of the artifact has a revision ID of 0. @@ -57,8 +57,8 @@ class BaseArtifactService(ABC): *, app_name: str, user_id: str, - session_id: str, filename: str, + session_id: Optional[str] = None, version: Optional[int] = None, ) -> Optional[types.Part]: """Gets an artifact from the artifact service storage. @@ -69,8 +69,8 @@ class BaseArtifactService(ABC): Args: app_name: The app name. user_id: The user ID. - session_id: The session ID. filename: The filename of the artifact. + session_id: The session ID. If `None`, load the user-scoped artifact. version: The version of the artifact. If None, the latest version will be returned. @@ -80,7 +80,7 @@ class BaseArtifactService(ABC): @abstractmethod async def list_artifact_keys( - self, *, app_name: str, user_id: str, session_id: str + self, *, app_name: str, user_id: str, session_id: Optional[str] = None ) -> list[str]: """Lists all the artifact filenames within a session. @@ -90,33 +90,48 @@ class BaseArtifactService(ABC): session_id: The ID of the session. Returns: - A list of all artifact filenames within a session. + A list of artifact filenames. If `session_id` is provided, returns + both session-scoped and user-scoped artifact filenames. If `session_id` + is `None`, returns + user-scoped artifact filenames. """ @abstractmethod async def delete_artifact( - self, *, app_name: str, user_id: str, session_id: str, filename: str + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, ) -> None: """Deletes an artifact. Args: app_name: The name of the application. user_id: The ID of the user. - session_id: The ID of the session. filename: The name of the artifact file. + session_id: The ID of the session. If `None`, delete the user-scoped + artifact. """ @abstractmethod async def list_versions( - self, *, app_name: str, user_id: str, session_id: str, filename: str + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, ) -> list[int]: """Lists all versions of an artifact. Args: app_name: The name of the application. user_id: The ID of the user. - session_id: The ID of the session. filename: The name of the artifact file. + session_id: The ID of the session. If `None`, only list the user-scoped + artifacts versions. Returns: A list of all available versions of the artifact. diff --git a/src/google/adk/artifacts/gcs_artifact_service.py b/src/google/adk/artifacts/gcs_artifact_service.py index a49d170e..7c40ef0c 100644 --- a/src/google/adk/artifacts/gcs_artifact_service.py +++ b/src/google/adk/artifacts/gcs_artifact_service.py @@ -41,7 +41,6 @@ class GcsArtifactService(BaseArtifactService): def __init__(self, bucket_name: str, **kwargs): """Initializes the GcsArtifactService. - Args: bucket_name: The name of the bucket to use. **kwargs: Keyword arguments to pass to the Google Cloud Storage client. @@ -56,9 +55,9 @@ class GcsArtifactService(BaseArtifactService): *, app_name: str, user_id: str, - session_id: str, filename: str, artifact: types.Part, + session_id: Optional[str] = None, ) -> int: return await asyncio.to_thread( self._save_artifact, @@ -75,8 +74,8 @@ class GcsArtifactService(BaseArtifactService): *, app_name: str, user_id: str, - session_id: str, filename: str, + session_id: Optional[str] = None, version: Optional[int] = None, ) -> Optional[types.Part]: return await asyncio.to_thread( @@ -90,7 +89,7 @@ class GcsArtifactService(BaseArtifactService): @override async def list_artifact_keys( - self, *, app_name: str, user_id: str, session_id: str + self, *, app_name: str, user_id: str, session_id: Optional[str] = None ) -> list[str]: return await asyncio.to_thread( self._list_artifact_keys, @@ -101,7 +100,12 @@ class GcsArtifactService(BaseArtifactService): @override async def delete_artifact( - self, *, app_name: str, user_id: str, session_id: str, filename: str + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, ) -> None: return await asyncio.to_thread( self._delete_artifact, @@ -113,7 +117,12 @@ class GcsArtifactService(BaseArtifactService): @override async def list_versions( - self, *, app_name: str, user_id: str, session_id: str, filename: str + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, ) -> list[int]: return await asyncio.to_thread( self._list_versions, @@ -139,31 +148,36 @@ class GcsArtifactService(BaseArtifactService): self, app_name: str, user_id: str, - session_id: str, filename: str, version: int, + session_id: Optional[str] = None, ) -> str: """Constructs the blob name in GCS. Args: app_name: The name of the application. user_id: The ID of the user. - session_id: The ID of the session. filename: The name of the artifact file. version: The version of the artifact. + session_id: The ID of the session. Returns: The constructed blob name in GCS. """ if self._file_has_user_namespace(filename): return f"{app_name}/{user_id}/user/{filename}/{version}" + + if session_id is None: + raise ValueError( + "Session ID must be provided for session-scoped artifacts." + ) return f"{app_name}/{user_id}/{session_id}/{filename}/{version}" def _save_artifact( self, app_name: str, user_id: str, - session_id: str, + session_id: Optional[str], filename: str, artifact: types.Part, ) -> int: @@ -176,7 +190,7 @@ class GcsArtifactService(BaseArtifactService): version = 0 if not versions else max(versions) + 1 blob_name = self._get_blob_name( - app_name, user_id, session_id, filename, version + app_name, user_id, filename, version, session_id ) blob = self.bucket.blob(blob_name) @@ -198,7 +212,7 @@ class GcsArtifactService(BaseArtifactService): self, app_name: str, user_id: str, - session_id: str, + session_id: Optional[str], filename: str, version: Optional[int] = None, ) -> Optional[types.Part]: @@ -214,7 +228,7 @@ class GcsArtifactService(BaseArtifactService): version = max(versions) blob_name = self._get_blob_name( - app_name, user_id, session_id, filename, version + app_name, user_id, filename, version, session_id ) blob = self.bucket.blob(blob_name) @@ -227,17 +241,18 @@ class GcsArtifactService(BaseArtifactService): return artifact def _list_artifact_keys( - self, app_name: str, user_id: str, session_id: str + self, app_name: str, user_id: str, session_id: Optional[str] ) -> list[str]: filenames = set() - session_prefix = f"{app_name}/{user_id}/{session_id}/" - session_blobs = self.storage_client.list_blobs( - self.bucket, prefix=session_prefix - ) - for blob in session_blobs: - *_, filename, _ = blob.name.split("/") - filenames.add(filename) + if session_id: + session_prefix = f"{app_name}/{user_id}/{session_id}/" + session_blobs = self.storage_client.list_blobs( + self.bucket, prefix=session_prefix + ) + for blob in session_blobs: + *_, filename, _ = blob.name.split("/") + filenames.add(filename) user_namespace_prefix = f"{app_name}/{user_id}/user/" user_namespace_blobs = self.storage_client.list_blobs( @@ -250,7 +265,11 @@ class GcsArtifactService(BaseArtifactService): return sorted(list(filenames)) def _delete_artifact( - self, app_name: str, user_id: str, session_id: str, filename: str + self, + app_name: str, + user_id: str, + session_id: Optional[str], + filename: str, ) -> None: versions = self._list_versions( app_name=app_name, @@ -260,18 +279,23 @@ class GcsArtifactService(BaseArtifactService): ) for version in versions: blob_name = self._get_blob_name( - app_name, user_id, session_id, filename, version + app_name, user_id, filename, version, session_id ) blob = self.bucket.blob(blob_name) blob.delete() return def _list_versions( - self, app_name: str, user_id: str, session_id: str, filename: str + self, + app_name: str, + user_id: str, + session_id: Optional[str], + filename: str, ) -> list[int]: """Lists all available versions of an artifact. - This method retrieves all versions of a specific artifact by querying GCS blobs + This method retrieves all versions of a specific artifact by querying GCS + blobs that match the constructed blob name prefix. Args: @@ -281,10 +305,11 @@ class GcsArtifactService(BaseArtifactService): filename: The name of the artifact file. Returns: - A list of version numbers (integers) available for the specified artifact. + A list of version numbers (integers) available for the specified + artifact. Returns an empty list if no versions are found. """ - prefix = self._get_blob_name(app_name, user_id, session_id, filename, "") + prefix = self._get_blob_name(app_name, user_id, filename, "", session_id) blobs = self.storage_client.list_blobs(self.bucket, prefix=prefix) versions = [] for blob in blobs: diff --git a/src/google/adk/artifacts/in_memory_artifact_service.py b/src/google/adk/artifacts/in_memory_artifact_service.py index e4837695..6b994130 100644 --- a/src/google/adk/artifacts/in_memory_artifact_service.py +++ b/src/google/adk/artifacts/in_memory_artifact_service.py @@ -48,21 +48,30 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): return filename.startswith("user:") def _artifact_path( - self, app_name: str, user_id: str, session_id: str, filename: str + self, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str], ) -> str: """Constructs the artifact path. Args: app_name: The name of the application. user_id: The ID of the user. - session_id: The ID of the session. filename: The name of the artifact file. + session_id: The ID of the session. Returns: The constructed artifact path. """ if self._file_has_user_namespace(filename): return f"{app_name}/{user_id}/user/{filename}" + + if session_id is None: + raise ValueError( + "Session ID must be provided for session-scoped artifacts." + ) return f"{app_name}/{user_id}/{session_id}/{filename}" @override @@ -71,11 +80,11 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): *, app_name: str, user_id: str, - session_id: str, filename: str, artifact: types.Part, + session_id: Optional[str] = None, ) -> int: - path = self._artifact_path(app_name, user_id, session_id, filename) + path = self._artifact_path(app_name, user_id, filename, session_id) if path not in self.artifacts: self.artifacts[path] = [] version = len(self.artifacts[path]) @@ -88,11 +97,11 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): *, app_name: str, user_id: str, - session_id: str, filename: str, + session_id: Optional[str] = None, version: Optional[int] = None, ) -> Optional[types.Part]: - path = self._artifact_path(app_name, user_id, session_id, filename) + path = self._artifact_path(app_name, user_id, filename, session_id) versions = self.artifacts.get(path) if not versions: return None @@ -102,13 +111,15 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): @override async def list_artifact_keys( - self, *, app_name: str, user_id: str, session_id: str + self, *, app_name: str, user_id: str, session_id: Optional[str] = None ) -> list[str]: - session_prefix = f"{app_name}/{user_id}/{session_id}/" usernamespace_prefix = f"{app_name}/{user_id}/user/" + session_prefix = ( + f"{app_name}/{user_id}/{session_id}/" if session_id else None + ) filenames = [] for path in self.artifacts: - if path.startswith(session_prefix): + if session_prefix and path.startswith(session_prefix): filename = path.removeprefix(session_prefix) filenames.append(filename) elif path.startswith(usernamespace_prefix): @@ -118,18 +129,28 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel): @override async def delete_artifact( - self, *, app_name: str, user_id: str, session_id: str, filename: str + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, ) -> None: - path = self._artifact_path(app_name, user_id, session_id, filename) + path = self._artifact_path(app_name, user_id, filename, session_id) if not self.artifacts.get(path): return None self.artifacts.pop(path, None) @override async def list_versions( - self, *, app_name: str, user_id: str, session_id: str, filename: str + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, ) -> list[int]: - path = self._artifact_path(app_name, user_id, session_id, filename) + path = self._artifact_path(app_name, user_id, filename, session_id) versions = self.artifacts.get(path) if not versions: return [] diff --git a/src/google/adk/tools/_forwarding_artifact_service.py b/src/google/adk/tools/_forwarding_artifact_service.py index 44607cd1..df59933c 100644 --- a/src/google/adk/tools/_forwarding_artifact_service.py +++ b/src/google/adk/tools/_forwarding_artifact_service.py @@ -39,9 +39,9 @@ class ForwardingArtifactService(BaseArtifactService): *, app_name: str, user_id: str, - session_id: str, filename: str, artifact: types.Part, + session_id: Optional[str] = None, ) -> int: return await self.tool_context.save_artifact( filename=filename, artifact=artifact @@ -53,8 +53,8 @@ class ForwardingArtifactService(BaseArtifactService): *, app_name: str, user_id: str, - session_id: str, filename: str, + session_id: Optional[str] = None, version: Optional[int] = None, ) -> Optional[types.Part]: return await self.tool_context.load_artifact( @@ -63,13 +63,18 @@ class ForwardingArtifactService(BaseArtifactService): @override async def list_artifact_keys( - self, *, app_name: str, user_id: str, session_id: str + self, *, app_name: str, user_id: str, session_id: Optional[str] = None ) -> 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 + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, ) -> None: del app_name, user_id, session_id if self._invocation_context.artifact_service is None: @@ -83,7 +88,12 @@ class ForwardingArtifactService(BaseArtifactService): @override async def list_versions( - self, *, app_name: str, user_id: str, session_id: str, filename: str + self, + *, + app_name: str, + user_id: str, + filename: str, + session_id: Optional[str] = None, ) -> list[int]: del app_name, user_id, session_id if self._invocation_context.artifact_service is None: