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

google_adk_extras/__init__.py

@@ -0,0 +1 @@
+"""Custom ADK Services - Extended implementations of Google ADK services."""

google_adk_extras/artifacts/__init__.py

@@ -0,0 +1,15 @@
+"""Custom artifact service implementations for Google ADK."""
+
+from .base_custom_artifact_service import BaseCustomArtifactService
+from .sql_artifact_service import SQLArtifactService
+from .mongo_artifact_service import MongoArtifactService
+from .local_folder_artifact_service import LocalFolderArtifactService
+from .s3_artifact_service import S3ArtifactService
+
+__all__ = [
+    "BaseCustomArtifactService",
+    "SQLArtifactService",
+    "MongoArtifactService",
+    "LocalFolderArtifactService",
+    "S3ArtifactService",
+]

google_adk_extras/artifacts/base_custom_artifact_service.py

@@ -0,0 +1,207 @@
+"""Base class for custom artifact services."""
+
+import abc
+from typing import Optional, List
+
+from google.adk.artifacts.base_artifact_service import BaseArtifactService
+from google.genai import types
+
+
+class BaseCustomArtifactService(BaseArtifactService, abc.ABC):
+    """Base class for custom artifact services with common functionality."""
+
+    def __init__(self):
+        """Initialize the base custom artifact service."""
+        super().__init__()
+        self._initialized = False
+
+    async def initialize(self) -> None:
+        """Initialize the artifact service.
+        
+        This method should be called before using the service to ensure
+        any required setup (database connections, etc.) is complete.
+        """
+        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 database connections, creating tables, directories, etc.
+        """
+        pass
+
+    async def cleanup(self) -> None:
+        """Clean up resources used by the artifact service.
+        
+        This method should be called when the service is no longer needed
+        to ensure proper cleanup of resources.
+        """
+        if self._initialized:
+            await self._cleanup_impl()
+            self._initialized = False
+
+    @abc.abstractmethod
+    async def _cleanup_impl(self) -> None:
+        """Implementation of service cleanup.
+        
+        This method should handle any cleanup required for the service,
+        such as closing database connections.
+        """
+        pass
+
+    async def save_artifact(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        artifact: types.Part,
+    ) -> int:
+        """Save an artifact."""
+        if not self._initialized:
+            await self.initialize()
+        return await self._save_artifact_impl(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+            artifact=artifact,
+        )
+
+    async def load_artifact(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: Optional[int] = None,
+    ) -> Optional[types.Part]:
+        """Load an artifact."""
+        if not self._initialized:
+            await self.initialize()
+        return await self._load_artifact_impl(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+            version=version,
+        )
+
+    async def list_artifact_keys(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+    ) -> List[str]:
+        """List artifact keys."""
+        if not self._initialized:
+            await self.initialize()
+        return await self._list_artifact_keys_impl(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+        )
+
+    async def delete_artifact(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> None:
+        """Delete an artifact."""
+        if not self._initialized:
+            await self.initialize()
+        await self._delete_artifact_impl(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+        )
+
+    async def list_versions(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> List[int]:
+        """List versions of an artifact."""
+        if not self._initialized:
+            await self.initialize()
+        return await self._list_versions_impl(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+        )
+
+    @abc.abstractmethod
+    async def _save_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        artifact: types.Part,
+    ) -> int:
+        """Implementation of artifact saving."""
+        pass
+
+    @abc.abstractmethod
+    async def _load_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: Optional[int] = None,
+    ) -> Optional[types.Part]:
+        """Implementation of artifact loading."""
+        pass
+
+    @abc.abstractmethod
+    async def _list_artifact_keys_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+    ) -> List[str]:
+        """Implementation of artifact key listing."""
+        pass
+
+    @abc.abstractmethod
+    async def _delete_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> None:
+        """Implementation of artifact deletion."""
+        pass
+
+    @abc.abstractmethod
+    async def _list_versions_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> List[int]:
+        """Implementation of version listing."""
+        pass

google_adk_extras/artifacts/local_folder_artifact_service.py

@@ -0,0 +1,242 @@
+"""Local folder-based artifact service implementation."""
+
+import os
+import json
+import base64
+from typing import Optional, List
+from pathlib import Path
+from datetime import datetime, timezone
+
+from google.genai import types
+from .base_custom_artifact_service import BaseCustomArtifactService
+
+
+class LocalFolderArtifactService(BaseCustomArtifactService):
+    """Local folder-based artifact service implementation."""
+
+    def __init__(self, base_directory: str = "./artifacts"):
+        """Initialize the local folder artifact service.
+        
+        Args:
+            base_directory: Base directory for storing artifacts
+        """
+        super().__init__()
+        self.base_directory = Path(base_directory)
+        # Create base directory if it doesn't exist
+        self.base_directory.mkdir(parents=True, exist_ok=True)
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the file system artifact service."""
+        # Ensure base directory exists
+        self.base_directory.mkdir(parents=True, exist_ok=True)
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up resources (no cleanup needed for file-based service)."""
+        pass
+
+    def _get_artifact_directory(self, app_name: str, user_id: str, session_id: str) -> Path:
+        """Generate directory path for artifacts."""
+        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."""
+        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."""
+        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."""
+        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."""
+        blob = types.Blob(data=data, mime_type=mime_type)
+        return types.Part(inline_data=blob)
+
+    async def _save_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        artifact: types.Part,
+    ) -> int:
+        """Implementation of artifact saving."""
+        try:
+            # Extract blob data
+            data, mime_type = self._serialize_blob(artifact)
+            
+            # Get the next version number
+            metadata_file = self._get_artifact_file_path(app_name, user_id, session_id, filename)
+            
+            if metadata_file.exists():
+                with open(metadata_file, 'r') as f:
+                    metadata = json.load(f)
+                version = len(metadata.get("versions", []))
+            else:
+                metadata = {
+                    "app_name": app_name,
+                    "user_id": user_id,
+                    "session_id": session_id,
+                    "filename": filename,
+                    "versions": []
+                }
+                version = 0
+            
+            # Save data to file
+            data_file = self._get_artifact_data_path(app_name, user_id, session_id, filename, version)
+            with open(data_file, 'wb') as f:
+                f.write(data)
+            
+            # Update metadata
+            version_info = {
+                "version": version,
+                "mime_type": mime_type,
+                "created_at": datetime.now(timezone.utc).isoformat(),
+                "data_file": data_file.name
+            }
+            metadata["versions"].append(version_info)
+            
+            # Save metadata
+            with open(metadata_file, 'w') as f:
+                json.dump(metadata, f, indent=2)
+            
+            return version
+        except Exception as e:
+            raise RuntimeError(f"Failed to save artifact: {e}")
+
+    async def _load_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: Optional[int] = None,
+    ) -> Optional[types.Part]:
+        """Implementation of artifact loading."""
+        try:
+            # Load metadata
+            metadata_file = self._get_artifact_file_path(app_name, user_id, session_id, filename)
+            
+            if not metadata_file.exists():
+                return None
+            
+            with open(metadata_file, 'r') as f:
+                metadata = json.load(f)
+            
+            # Determine version to load
+            versions = metadata.get("versions", [])
+            if not versions:
+                return None
+            
+            if version is not None:
+                # Find specific version
+                version_info = None
+                for v in versions:
+                    if v["version"] == version:
+                        version_info = v
+                        break
+                if not version_info:
+                    return None
+            else:
+                # Load latest version
+                version_info = versions[-1]
+            
+            # Load data
+            data_file = self._get_artifact_directory(app_name, user_id, session_id) / version_info["data_file"]
+            if not data_file.exists():
+                return None
+            
+            with open(data_file, 'rb') as f:
+                data = f.read()
+            
+            # Create Part from blob data
+            return self._deserialize_blob(data, version_info["mime_type"])
+        except Exception as e:
+            raise RuntimeError(f"Failed to load artifact: {e}")
+
+    async def _list_artifact_keys_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+    ) -> List[str]:
+        """Implementation of artifact key listing."""
+        try:
+            directory = self._get_artifact_directory(app_name, user_id, session_id)
+            
+            # Find all metadata files
+            artifact_keys = []
+            if directory.exists():
+                for file_path in directory.glob("*.json"):
+                    # Extract filename from metadata file name (remove .json extension)
+                    filename = file_path.name[:-5]  # Remove .json
+                    artifact_keys.append(filename)
+            
+            return artifact_keys
+        except Exception as e:
+            raise RuntimeError(f"Failed to list artifact keys: {e}")
+
+    async def _delete_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> None:
+        """Implementation of artifact deletion."""
+        try:
+            # Load metadata to find all version files
+            metadata_file = self._get_artifact_file_path(app_name, user_id, session_id, filename)
+            
+            if metadata_file.exists():
+                with open(metadata_file, 'r') as f:
+                    metadata = json.load(f)
+                
+                # Delete all version data files
+                directory = self._get_artifact_directory(app_name, user_id, session_id)
+                for version_info in metadata.get("versions", []):
+                    data_file = directory / version_info["data_file"]
+                    if data_file.exists():
+                        data_file.unlink()
+                
+                # Delete metadata file
+                metadata_file.unlink()
+        except Exception as e:
+            raise RuntimeError(f"Failed to delete artifact: {e}")
+
+    async def _list_versions_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> List[int]:
+        """Implementation of version listing."""
+        try:
+            metadata_file = self._get_artifact_file_path(app_name, user_id, session_id, filename)
+            
+            if not metadata_file.exists():
+                return []
+            
+            with open(metadata_file, 'r') as f:
+                metadata = json.load(f)
+            
+            versions = metadata.get("versions", [])
+            return [v["version"] for v in versions]
+        except Exception as e:
+            raise RuntimeError(f"Failed to list versions: {e}")

google_adk_extras/artifacts/mongo_artifact_service.py

@@ -0,0 +1,221 @@
+"""MongoDB-based artifact service implementation."""
+
+import json
+from typing import Optional, List
+from datetime import datetime
+
+try:
+    from pymongo import MongoClient
+    from pymongo.errors import PyMongoError
+except ImportError:
+    raise ImportError(
+        "PyMongo is required for MongoArtifactService. "
+        "Install it with: pip install pymongo"
+    )
+
+from google.genai import types
+from .base_custom_artifact_service import BaseCustomArtifactService
+
+
+class MongoArtifactService(BaseCustomArtifactService):
+    """MongoDB-based artifact service implementation."""
+
+    def __init__(self, connection_string: str, database_name: str = "adk_artifacts"):
+        """Initialize the MongoDB artifact service.
+        
+        Args:
+            connection_string: MongoDB connection string
+            database_name: Name of the database to use
+        """
+        super().__init__()
+        self.connection_string = connection_string
+        self.database_name = database_name
+        self.client: Optional[MongoClient] = None
+        self.db = None
+        self.collection = None
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the MongoDB connection."""
+        try:
+            self.client = MongoClient(self.connection_string)
+            self.db = self.client[self.database_name]
+            self.collection = self.db.artifacts
+            
+            # Create indexes for better performance
+            self.collection.create_index([
+                ("app_name", 1),
+                ("user_id", 1),
+                ("session_id", 1),
+                ("filename", 1),
+                ("version", 1)
+            ])
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to initialize MongoDB artifact service: {e}")
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up MongoDB connections."""
+        if self.client:
+            self.client.close()
+            self.client = None
+        self.db = None
+        self.collection = None
+
+    def _serialize_blob(self, part: types.Part) -> tuple[bytes, str]:
+        """Extract blob data and mime type from a Part."""
+        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."""
+        blob = types.Blob(data=data, mime_type=mime_type)
+        return types.Part(inline_data=blob)
+
+    async def _save_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        artifact: types.Part,
+    ) -> int:
+        """Implementation of artifact saving."""
+        try:
+            # Extract blob data
+            data, mime_type = self._serialize_blob(artifact)
+            
+            # Get the next version number
+            latest_version_doc = self.collection.find_one(
+                {
+                    "app_name": app_name,
+                    "user_id": user_id,
+                    "session_id": session_id,
+                    "filename": filename
+                },
+                sort=[("version", -1)]
+            )
+            
+            version = (latest_version_doc["version"] + 1) if latest_version_doc else 0
+            
+            # Create document
+            document = {
+                "app_name": app_name,
+                "user_id": user_id,
+                "session_id": session_id,
+                "filename": filename,
+                "version": version,
+                "mime_type": mime_type,
+                "data": data,
+                "created_at": datetime.utcnow()
+            }
+            
+            # Insert into MongoDB
+            self.collection.insert_one(document)
+            
+            return version
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to save artifact: {e}")
+
+    async def _load_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: Optional[int] = None,
+    ) -> Optional[types.Part]:
+        """Implementation of artifact loading."""
+        try:
+            query = {
+                "app_name": app_name,
+                "user_id": user_id,
+                "session_id": session_id,
+                "filename": filename
+            }
+            
+            if version is not None:
+                query["version"] = version
+                sort = None
+            else:
+                # Sort by version descending to get the latest
+                sort = [("version", -1)]
+            
+            document = self.collection.find_one(query, sort=sort)
+            
+            if not document:
+                return None
+            
+            # Create Part from blob data
+            return self._deserialize_blob(document["data"], document["mime_type"])
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to load artifact: {e}")
+
+    async def _list_artifact_keys_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+    ) -> List[str]:
+        """Implementation of artifact key listing."""
+        try:
+            # Get distinct filenames
+            cursor = self.collection.distinct(
+                "filename",
+                {
+                    "app_name": app_name,
+                    "user_id": user_id,
+                    "session_id": session_id
+                }
+            )
+            
+            return list(cursor)
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to list artifact keys: {e}")
+
+    async def _delete_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> None:
+        """Implementation of artifact deletion."""
+        try:
+            # Delete all versions of the artifact
+            self.collection.delete_many({
+                "app_name": app_name,
+                "user_id": user_id,
+                "session_id": session_id,
+                "filename": filename
+            })
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to delete artifact: {e}")
+
+    async def _list_versions_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> List[int]:
+        """Implementation of version listing."""
+        try:
+            cursor = self.collection.find(
+                {
+                    "app_name": app_name,
+                    "user_id": user_id,
+                    "session_id": session_id,
+                    "filename": filename
+                },
+                {"version": 1, "_id": 0}
+            ).sort("version", 1)
+            
+            return [doc["version"] for doc in cursor]
+        except PyMongoError as e:
+            raise RuntimeError(f"Failed to list versions: {e}")