feat: Make session_id optional in BaseArtifactService methods

PiperOrigin-RevId: 816782982

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.

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,10 +241,11 @@ 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()
 
+    if session_id:
       session_prefix = f"{app_name}/{user_id}/{session_id}/"
       session_blobs = self.storage_client.list_blobs(
           self.bucket, prefix=session_prefix
@@ -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:

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 []

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: