google-adk-extras 0.1.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

google_adk_extras/artifacts/__init__.py

@@ -1,15 +1,28 @@
-"""Custom artifact service implementations for Google ADK."""
+"""Custom artifact service implementations for Google ADK.
+
+Optional backends are imported lazily based on installed dependencies.
+"""
 
 from .base_custom_artifact_service import BaseCustomArtifactService
-from .sql_artifact_service import SQLArtifactService
-from .mongo_artifact_service import MongoArtifactService
+
+try:
+    from .sql_artifact_service import SQLArtifactService  # type: ignore
+except Exception:
+    SQLArtifactService = None  # type: ignore
+
+try:
+    from .mongo_artifact_service import MongoArtifactService  # type: ignore
+except Exception:
+    MongoArtifactService = None  # type: ignore
+
 from .local_folder_artifact_service import LocalFolderArtifactService
-from .s3_artifact_service import S3ArtifactService
-
-__all__ = [
-    "BaseCustomArtifactService",
-    "SQLArtifactService",
-    "MongoArtifactService",
-    "LocalFolderArtifactService",
-    "S3ArtifactService",
-]
+
+try:
+    from .s3_artifact_service import S3ArtifactService  # type: ignore
+except Exception:
+    S3ArtifactService = None  # type: ignore
+
+__all__ = ["BaseCustomArtifactService", "LocalFolderArtifactService"]
+for _name in ("SQLArtifactService", "MongoArtifactService", "S3ArtifactService"):
+    if globals().get(_name) is not None:
+        __all__.append(_name)

google_adk_extras/artifacts/base_custom_artifact_service.py

@@ -8,7 +8,11 @@ from google.genai import types
 
 
 class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
-    """Base class for custom artifact services with common functionality."""
+    """Base class for custom artifact services with common functionality.
+
+    This abstract base class provides a foundation for implementing custom
+    artifact services with automatic initialization and cleanup handling.
+    """
 
     def __init__(self):
         """Initialize the base custom artifact service."""
@@ -20,6 +24,9 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         
         This method should be called before using the service to ensure
         any required setup (database connections, etc.) is complete.
+        
+        Raises:
+            RuntimeError: If initialization fails.
         """
         if not self._initialized:
             await self._initialize_impl()
@@ -31,6 +38,9 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         
         This method should handle any setup required for the service to function,
         such as database connections, creating tables, directories, etc.
+        
+        Raises:
+            RuntimeError: If initialization fails.
         """
         pass
 
@@ -62,7 +72,21 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         filename: str,
         artifact: types.Part,
     ) -> int:
-        """Save an artifact."""
+        """Save 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 file to save.
+            artifact: The artifact to save.
+            
+        Returns:
+            The version number of the saved artifact.
+            
+        Raises:
+            RuntimeError: If saving the artifact fails.
+        """
         if not self._initialized:
             await self.initialize()
         return await self._save_artifact_impl(
@@ -82,7 +106,22 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         filename: str,
         version: Optional[int] = None,
     ) -> Optional[types.Part]:
-        """Load an artifact."""
+        """Load 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 file to load.
+            version: Optional version number to load. If not provided,
+                the latest version will be loaded.
+                
+        Returns:
+            The loaded artifact if found, None otherwise.
+            
+        Raises:
+            RuntimeError: If loading the artifact fails.
+        """
         if not self._initialized:
             await self.initialize()
         return await self._load_artifact_impl(
@@ -100,7 +139,19 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         user_id: str,
         session_id: str,
     ) -> List[str]:
-        """List artifact keys."""
+        """List artifact keys.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session.
+            
+        Returns:
+            A list of artifact keys (filenames).
+            
+        Raises:
+            RuntimeError: If listing artifact keys fails.
+        """
         if not self._initialized:
             await self.initialize()
         return await self._list_artifact_keys_impl(
@@ -117,7 +168,17 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         session_id: str,
         filename: str,
     ) -> None:
-        """Delete an artifact."""
+        """Delete 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 file to delete.
+            
+        Raises:
+            RuntimeError: If deleting the artifact fails.
+        """
         if not self._initialized:
             await self.initialize()
         await self._delete_artifact_impl(
@@ -135,7 +196,20 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         session_id: str,
         filename: str,
     ) -> List[int]:
-        """List versions of an artifact."""
+        """List 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 file to list versions for.
+            
+        Returns:
+            A list of version numbers.
+            
+        Raises:
+            RuntimeError: If listing versions fails.
+        """
         if not self._initialized:
             await self.initialize()
         return await self._list_versions_impl(
@@ -155,7 +229,21 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         filename: str,
         artifact: types.Part,
     ) -> int:
-        """Implementation of artifact saving."""
+        """Implementation of artifact saving.
+        
+        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 file to save.
+            artifact: The artifact to save.
+            
+        Returns:
+            The version number of the saved artifact.
+            
+        Raises:
+            RuntimeError: If saving the artifact fails.
+        """
         pass
 
     @abc.abstractmethod
@@ -168,7 +256,21 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         filename: str,
         version: Optional[int] = None,
     ) -> Optional[types.Part]:
-        """Implementation of artifact loading."""
+        """Implementation of artifact loading.
+        
+        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 file to load.
+            version: Optional version number to load.
+            
+        Returns:
+            The loaded artifact if found, None otherwise.
+            
+        Raises:
+            RuntimeError: If loading the artifact fails.
+        """
         pass
 
     @abc.abstractmethod
@@ -179,7 +281,19 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         user_id: str,
         session_id: str,
     ) -> List[str]:
-        """Implementation of artifact key listing."""
+        """Implementation of artifact key listing.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session.
+            
+        Returns:
+            A list of artifact keys (filenames).
+            
+        Raises:
+            RuntimeError: If listing artifact keys fails.
+        """
         pass
 
     @abc.abstractmethod
@@ -191,7 +305,17 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         session_id: str,
         filename: str,
     ) -> None:
-        """Implementation of artifact deletion."""
+        """Implementation of artifact deletion.
+        
+        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 file to delete.
+            
+        Raises:
+            RuntimeError: If deleting the artifact fails.
+        """
         pass
 
     @abc.abstractmethod
@@ -203,5 +327,18 @@ class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
         session_id: str,
         filename: str,
     ) -> List[int]:
-        """Implementation of version listing."""
+        """Implementation of version listing.
+        
+        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 file to list versions for.
+            
+        Returns:
+            A list of version numbers.
+            
+        Raises:
+            RuntimeError: If listing versions fails.
+        """
         pass

google_adk_extras/artifacts/local_folder_artifact_service.py

@@ -12,13 +12,17 @@ from .base_custom_artifact_service import BaseCustomArtifactService
 
 
 class LocalFolderArtifactService(BaseCustomArtifactService):
-    """Local folder-based artifact service implementation."""
+    """Local folder-based artifact service implementation.
+
+    This service stores artifacts in the local file system with full versioning support.
+    Each artifact is stored with its metadata in JSON format and binary data in separate files.
+    """
 
     def __init__(self, base_directory: str = "./artifacts"):
         """Initialize the local folder artifact service.
         
         Args:
-            base_directory: Base directory for storing artifacts
+            base_directory: Base directory for storing artifacts. Defaults to "./artifacts".
         """
         super().__init__()
         self.base_directory = Path(base_directory)
@@ -26,7 +30,10 @@ class LocalFolderArtifactService(BaseCustomArtifactService):
         self.base_directory.mkdir(parents=True, exist_ok=True)
 
     async def _initialize_impl(self) -> None:
-        """Initialize the file system artifact service."""
+        """Initialize the file system artifact service.
+        
+        Ensures the base directory exists.
+        """
         # Ensure base directory exists
         self.base_directory.mkdir(parents=True, exist_ok=True)
 
@@ -35,30 +42,78 @@ class LocalFolderArtifactService(BaseCustomArtifactService):
         pass
 
     def _get_artifact_directory(self, app_name: str, user_id: str, session_id: str) -> Path:
-        """Generate directory path for artifacts."""
+        """Generate directory path for artifacts.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session.
+            
+        Returns:
+            Path to the artifact directory.
+        """
         directory = self.base_directory / app_name / user_id / session_id
         directory.mkdir(parents=True, exist_ok=True)
         return directory
 
     def _get_artifact_file_path(self, app_name: str, user_id: str, session_id: str, filename: str) -> Path:
-        """Generate file path for artifact metadata."""
+        """Generate file path for artifact metadata.
+        
+        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 file.
+            
+        Returns:
+            Path to the metadata file.
+        """
         directory = self._get_artifact_directory(app_name, user_id, session_id)
         return directory / f"{filename}.json"
 
     def _get_artifact_data_path(self, app_name: str, user_id: str, session_id: str, filename: str, version: int) -> Path:
-        """Generate file path for artifact data."""
+        """Generate file path for artifact data.
+        
+        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 file.
+            version: The version number.
+            
+        Returns:
+            Path to the data file.
+        """
         directory = self._get_artifact_directory(app_name, user_id, session_id)
         return directory / f"{filename}.v{version}.data"
 
     def _serialize_blob(self, part: types.Part) -> tuple[bytes, str]:
-        """Extract blob data and mime type from a Part."""
+        """Extract blob data and mime type from a Part.
+        
+        Args:
+            part: The Part object containing the blob data.
+            
+        Returns:
+            A tuple of (data, mime_type).
+            
+        Raises:
+            ValueError: If the part type is not supported.
+        """
         if part.inline_data:
             return part.inline_data.data, part.inline_data.mime_type or "application/octet-stream"
         else:
             raise ValueError("Only inline_data parts are supported")
 
     def _deserialize_blob(self, data: bytes, mime_type: str) -> types.Part:
-        """Create a Part from blob data and mime type."""
+        """Create a Part from blob data and mime type.
+        
+        Args:
+            data: The binary data.
+            mime_type: The MIME type of the data.
+            
+        Returns:
+            A Part object containing the blob data.
+        """
         blob = types.Blob(data=data, mime_type=mime_type)
         return types.Part(inline_data=blob)
 
@@ -71,7 +126,22 @@ class LocalFolderArtifactService(BaseCustomArtifactService):
         filename: str,
         artifact: types.Part,
     ) -> int:
-        """Implementation of artifact saving."""
+        """Implementation of artifact saving.
+        
+        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 file to save.
+            artifact: The artifact to save.
+            
+        Returns:
+            The version number of the saved artifact.
+            
+        Raises:
+            RuntimeError: If saving the artifact fails.
+            ValueError: If the artifact type is not supported.
+        """
         try:
             # Extract blob data
             data, mime_type = self._serialize_blob(artifact)
@@ -124,7 +194,22 @@ class LocalFolderArtifactService(BaseCustomArtifactService):
         filename: str,
         version: Optional[int] = None,
     ) -> Optional[types.Part]:
-        """Implementation of artifact loading."""
+        """Implementation of artifact loading.
+        
+        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 file to load.
+            version: Optional version number to load. If not provided,
+                the latest version will be loaded.
+                
+        Returns:
+            The loaded artifact if found, None otherwise.
+            
+        Raises:
+            RuntimeError: If loading the artifact fails.
+        """
         try:
             # Load metadata
             metadata_file = self._get_artifact_file_path(app_name, user_id, session_id, filename)
@@ -173,7 +258,19 @@ class LocalFolderArtifactService(BaseCustomArtifactService):
         user_id: str,
         session_id: str,
     ) -> List[str]:
-        """Implementation of artifact key listing."""
+        """Implementation of artifact key listing.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session.
+            
+        Returns:
+            A list of artifact keys (filenames).
+            
+        Raises:
+            RuntimeError: If listing artifact keys fails.
+        """
         try:
             directory = self._get_artifact_directory(app_name, user_id, session_id)
             
@@ -197,7 +294,17 @@ class LocalFolderArtifactService(BaseCustomArtifactService):
         session_id: str,
         filename: str,
     ) -> None:
-        """Implementation of artifact deletion."""
+        """Implementation of artifact deletion.
+        
+        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 file to delete.
+            
+        Raises:
+            RuntimeError: If deleting the artifact fails.
+        """
         try:
             # Load metadata to find all version files
             metadata_file = self._get_artifact_file_path(app_name, user_id, session_id, filename)
@@ -226,7 +333,20 @@ class LocalFolderArtifactService(BaseCustomArtifactService):
         session_id: str,
         filename: str,
     ) -> List[int]:
-        """Implementation of version listing."""
+        """Implementation of version listing.
+        
+        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 file to list versions for.
+            
+        Returns:
+            A list of version numbers.
+            
+        Raises:
+            RuntimeError: If listing versions fails.
+        """
         try:
             metadata_file = self._get_artifact_file_path(app_name, user_id, session_id, filename)