feat: Make session_id optional in BaseArtifactService methods

PiperOrigin-RevId: 816782982

Author: google-genai-bot

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 @@ async def save_artifact(
       *,
       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 @@ async def save_artifact(
     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 @@ async def load_artifact(
       *,
       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 @@ async def load_artifact(
     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 @@ async def load_artifact(
 
   @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 @@ async def list_artifact_keys(
         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.

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 @@ async def save_artifact(
       *,
       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 @@ async def load_artifact(
       *,
       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 @@ async def load_artifact(
 
   @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 @@ async def list_artifact_keys(
 
   @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 @@ async def delete_artifact(
 
   @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 @@ def _get_blob_name(
       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 @@ def _save_artifact(
     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 @@ def _load_artifact(
       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 @@ def _load_artifact(
       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 @@ def _load_artifact(
     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 @@ def _list_artifact_keys(
     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 @@ def _delete_artifact(
     )
     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 @@ def _list_versions(
         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: