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

google_adk_extras/artifacts/s3_artifact_service.py

@@ -3,7 +3,7 @@
 import json
 import base64
 from typing import Optional, List
-from datetime import datetime
+from datetime import datetime, timezone
 
 try:
     import boto3
@@ -19,7 +19,12 @@ from .base_custom_artifact_service import BaseCustomArtifactService
 
 
 class S3ArtifactService(BaseCustomArtifactService):
-    """S3-compatible artifact service implementation."""
+    """S3-compatible artifact service implementation.
+
+    This service stores artifacts in AWS S3 or S3-compatible storage services.
+    It supports versioning and works with any S3-compatible service including
+    AWS S3, MinIO, Google Cloud Storage, etc.
+    """
 
     def __init__(
         self,
@@ -33,12 +38,12 @@ class S3ArtifactService(BaseCustomArtifactService):
         """Initialize the S3 artifact service.
         
         Args:
-            bucket_name: S3 bucket name
-            endpoint_url: S3 endpoint URL (for non-AWS S3 services)
-            region_name: AWS region name
-            aws_access_key_id: AWS access key ID
-            aws_secret_access_key: AWS secret access key
-            prefix: Prefix for artifact storage paths
+            bucket_name: S3 bucket name.
+            endpoint_url: S3 endpoint URL (for non-AWS S3 services like MinIO).
+            region_name: AWS region name.
+            aws_access_key_id: AWS access key ID.
+            aws_secret_access_key: AWS secret access key.
+            prefix: Prefix for artifact storage paths. Defaults to "adk-artifacts".
         """
         super().__init__()
         self.bucket_name = bucket_name
@@ -50,7 +55,14 @@ class S3ArtifactService(BaseCustomArtifactService):
         self.s3_client = None
 
     async def _initialize_impl(self) -> None:
-        """Initialize the S3 client."""
+        """Initialize the S3 client.
+        
+        Creates the S3 client and verifies the bucket exists (or creates it).
+        
+        Raises:
+            RuntimeError: If S3 initialization fails.
+            NoCredentialsError: If AWS credentials are not found.
+        """
         try:
             # Create S3 client
             self.s3_client = boto3.client(
@@ -88,22 +100,61 @@ class S3ArtifactService(BaseCustomArtifactService):
             self.s3_client = None
 
     def _get_artifact_key(self, app_name: str, user_id: str, session_id: str, filename: str) -> str:
-        """Generate S3 key for artifact metadata."""
+        """Generate S3 key 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:
+            S3 key for the metadata file.
+        """
         return f"{self.prefix}/{app_name}/{user_id}/{session_id}/{filename}.json"
 
     def _get_artifact_data_key(self, app_name: str, user_id: str, session_id: str, filename: str, version: int) -> str:
-        """Generate S3 key for artifact data."""
+        """Generate S3 key 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:
+            S3 key for the data file.
+        """
         return f"{self.prefix}/{app_name}/{user_id}/{session_id}/{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)
 
@@ -116,7 +167,22 @@ class S3ArtifactService(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)
@@ -155,7 +221,7 @@ class S3ArtifactService(BaseCustomArtifactService):
             version_info = {
                 "version": version,
                 "mime_type": mime_type,
-                "created_at": datetime.utcnow().isoformat(),
+                "created_at": datetime.now(timezone.utc).isoformat(),
                 "data_key": data_key
             }
             metadata["versions"].append(version_info)
@@ -180,7 +246,22 @@ class S3ArtifactService(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_key = self._get_artifact_key(app_name, user_id, session_id, filename)
@@ -237,7 +318,19 @@ class S3ArtifactService(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:
             # List objects with the prefix
             prefix = f"{self.prefix}/{app_name}/{user_id}/{session_id}/"
@@ -269,7 +362,17 @@ class S3ArtifactService(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_key = self._get_artifact_key(app_name, user_id, session_id, filename)
@@ -304,7 +407,20 @@ class S3ArtifactService(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_key = self._get_artifact_key(app_name, user_id, session_id, filename)
             

google_adk_extras/artifacts/sql_artifact_service.py

@@ -43,7 +43,12 @@ class SQLArtifactModel(Base):
 
 
 class SQLArtifactService(BaseCustomArtifactService):
-    """SQL-based artifact service implementation."""
+    """SQL-based artifact service implementation.
+
+    This service stores artifacts in a SQL database using SQLAlchemy.
+    It supports various SQL databases including SQLite, PostgreSQL, and MySQL.
+    Artifacts are stored with full versioning support.
+    """
 
     def __init__(self, database_url: str):
         """Initialize the SQL artifact service.
@@ -57,7 +62,11 @@ class SQLArtifactService(BaseCustomArtifactService):
         self.session_local: Optional[object] = None
 
     async def _initialize_impl(self) -> None:
-        """Initialize the database connection and create tables."""
+        """Initialize the database connection and create tables.
+        
+        Raises:
+            RuntimeError: If database initialization fails.
+        """
         try:
             self.engine = create_engine(self.database_url)
             Base.metadata.create_all(self.engine)
@@ -77,13 +86,30 @@ class SQLArtifactService(BaseCustomArtifactService):
         self.session_local = None
 
     def _get_db_session(self):
-        """Get a database session."""
+        """Get a database session.
+        
+        Returns:
+            A database session object.
+            
+        Raises:
+            RuntimeError: If the service is not initialized.
+        """
         if not self.session_local:
             raise RuntimeError("Service not initialized")
         return self.session_local()
 
     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:
@@ -93,7 +119,15 @@ class SQLArtifactService(BaseCustomArtifactService):
             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)
 
@@ -106,7 +140,22 @@ class SQLArtifactService(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.
+        """
         db_session = self._get_db_session()
         try:
             # Extract blob data
@@ -153,7 +202,22 @@ class SQLArtifactService(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.
+        """
         db_session = self._get_db_session()
         try:
             query = db_session.query(SQLArtifactModel).filter(
@@ -188,7 +252,19 @@ class SQLArtifactService(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.
+        """
         db_session = self._get_db_session()
         try:
             # Get distinct filenames
@@ -212,7 +288,17 @@ class SQLArtifactService(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.
+        """
         db_session = self._get_db_session()
         try:
             # Delete all versions of the artifact
@@ -238,7 +324,20 @@ class SQLArtifactService(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.
+        """
         db_session = self._get_db_session()
         try:
             versions = db_session.query(SQLArtifactModel.version).filter(

google_adk_extras/credentials/__init__.py

@@ -0,0 +1,34 @@
+"""Custom credential service implementations for Google ADK.
+
+Optional services are imported lazily to avoid import-time failures when
+their third-party dependencies are not installed.
+"""
+
+from .base_custom_credential_service import BaseCustomCredentialService
+from .google_oauth2_credential_service import GoogleOAuth2CredentialService
+from .github_oauth2_credential_service import GitHubOAuth2CredentialService
+from .microsoft_oauth2_credential_service import MicrosoftOAuth2CredentialService
+from .x_oauth2_credential_service import XOAuth2CredentialService
+from .http_basic_auth_credential_service import (
+    HTTPBasicAuthCredentialService,
+    HTTPBasicAuthWithCredentialsService,
+)
+
+# Optional: JWT (requires PyJWT)
+try:
+    from .jwt_credential_service import JWTCredentialService  # type: ignore
+except Exception:  # ImportError or transitive import errors
+    JWTCredentialService = None  # type: ignore
+
+__all__ = [
+    "BaseCustomCredentialService",
+    "GoogleOAuth2CredentialService",
+    "GitHubOAuth2CredentialService",
+    "MicrosoftOAuth2CredentialService",
+    "XOAuth2CredentialService",
+    "HTTPBasicAuthCredentialService",
+    "HTTPBasicAuthWithCredentialsService",
+]
+
+if JWTCredentialService is not None:
+    __all__.append("JWTCredentialService")

google_adk_extras/credentials/base_custom_credential_service.py

@@ -0,0 +1,113 @@
+"""Base class for custom credential services."""
+
+import abc
+from typing import Optional
+
+from google.adk.auth.credential_service.base_credential_service import BaseCredentialService, CallbackContext
+from google.adk.auth import AuthConfig, AuthCredential
+
+
+class BaseCustomCredentialService(BaseCredentialService, abc.ABC):
+    """Base class for custom credential services with common functionality.
+
+    This abstract base class provides a foundation for implementing custom
+    credential services with automatic initialization and cleanup handling.
+    """
+
+    def __init__(self):
+        """Initialize the base custom credential service."""
+        super().__init__()
+        self._initialized = False
+
+    async def initialize(self) -> None:
+        """Initialize the credential service.
+        
+        This method should be called before using the service to ensure
+        any required setup (connections, validations, etc.) is complete.
+        
+        Raises:
+            RuntimeError: If initialization fails.
+        """
+        if not self._initialized:
+            await self._initialize_impl()
+            self._initialized = True
+
+    @abc.abstractmethod
+    async def _initialize_impl(self) -> None:
+        """Implementation of service initialization.
+        
+        This method should handle any setup required for the service to function,
+        such as validating credentials, establishing connections, etc.
+        
+        Raises:
+            RuntimeError: If initialization fails.
+        """
+        pass
+
+    async def cleanup(self) -> None:
+        """Clean up resources used by the credential service.
+        
+        This method should be called when the service is no longer needed
+        to ensure proper cleanup of any resources.
+        """
+        if self._initialized:
+            await self._cleanup_impl()
+            self._initialized = False
+
+    async def _cleanup_impl(self) -> None:
+        """Implementation of service cleanup.
+        
+        This method should handle cleanup of any resources used by the service.
+        The default implementation does nothing, but subclasses can override
+        to perform specific cleanup operations.
+        """
+        pass
+
+    def _check_initialized(self) -> None:
+        """Check if the service has been initialized.
+        
+        Raises:
+            RuntimeError: If the service has not been initialized.
+        """
+        if not self._initialized:
+            raise RuntimeError(
+                f"{self.__class__.__name__} must be initialized before use. "
+                "Call await service.initialize() first."
+            )
+
+    @abc.abstractmethod
+    async def load_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> Optional[AuthCredential]:
+        """Load the credential by auth config and current callback context.
+
+        Args:
+            auth_config: The auth config which contains the auth scheme and auth
+                credential information. auth_config.get_credential_key will be used to
+                build the key to load the credential.
+            callback_context: The context of the current invocation when the tool is
+                trying to load the credential.
+
+        Returns:
+            Optional[AuthCredential]: the credential saved in the store, or None if not found.
+        """
+        pass
+
+    @abc.abstractmethod
+    async def save_credential(
+        self,
+        auth_config: AuthConfig,
+        callback_context: CallbackContext,
+    ) -> None:
+        """Save the exchanged_auth_credential in auth config.
+
+        Args:
+            auth_config: The auth config which contains the auth scheme and auth
+                credential information. auth_config.get_credential_key will be used to
+                build the key to save the credential.
+            callback_context: The context of the current invocation when the tool is
+                trying to save the credential.
+        """
+        pass