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

google_adk_extras/artifacts/s3_artifact_service.py

@@ -0,0 +1,323 @@
+"""S3-compatible artifact service implementation."""
+
+import json
+import base64
+from typing import Optional, List
+from datetime import datetime
+
+try:
+    import boto3
+    from botocore.exceptions import ClientError, NoCredentialsError
+except ImportError:
+    raise ImportError(
+        "Boto3 is required for S3ArtifactService. "
+        "Install it with: pip install boto3"
+    )
+
+from google.genai import types
+from .base_custom_artifact_service import BaseCustomArtifactService
+
+
+class S3ArtifactService(BaseCustomArtifactService):
+    """S3-compatible artifact service implementation."""
+
+    def __init__(
+        self,
+        bucket_name: str,
+        endpoint_url: Optional[str] = None,
+        region_name: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        prefix: str = "adk-artifacts",
+    ):
+        """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
+        """
+        super().__init__()
+        self.bucket_name = bucket_name
+        self.endpoint_url = endpoint_url
+        self.region_name = region_name
+        self.aws_access_key_id = aws_access_key_id
+        self.aws_secret_access_key = aws_secret_access_key
+        self.prefix = prefix
+        self.s3_client = None
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the S3 client."""
+        try:
+            # Create S3 client
+            self.s3_client = boto3.client(
+                's3',
+                endpoint_url=self.endpoint_url,
+                region_name=self.region_name,
+                aws_access_key_id=self.aws_access_key_id,
+                aws_secret_access_key=self.aws_secret_access_key,
+            )
+            
+            # Verify bucket exists or create it
+            try:
+                self.s3_client.head_bucket(Bucket=self.bucket_name)
+            except ClientError as e:
+                error_code = int(e.response['Error']['Code'])
+                if error_code == 404:
+                    # Bucket doesn't exist, create it
+                    if self.region_name:
+                        self.s3_client.create_bucket(
+                            Bucket=self.bucket_name,
+                            CreateBucketConfiguration={'LocationConstraint': self.region_name}
+                        )
+                    else:
+                        self.s3_client.create_bucket(Bucket=self.bucket_name)
+                else:
+                    raise
+        except NoCredentialsError:
+            raise RuntimeError("AWS credentials not found. Please provide credentials.")
+        except ClientError as e:
+            raise RuntimeError(f"Failed to initialize S3 artifact service: {e}")
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up S3 client."""
+        if self.s3_client:
+            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."""
+        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."""
+        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."""
+        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_key = self._get_artifact_key(app_name, user_id, session_id, filename)
+            
+            try:
+                # Try to load existing metadata
+                response = self.s3_client.get_object(Bucket=self.bucket_name, Key=metadata_key)
+                metadata = json.loads(response['Body'].read().decode('utf-8'))
+                version = len(metadata.get("versions", []))
+            except ClientError as e:
+                if e.response['Error']['Code'] == 'NoSuchKey':
+                    # Metadata doesn't exist, create new
+                    metadata = {
+                        "app_name": app_name,
+                        "user_id": user_id,
+                        "session_id": session_id,
+                        "filename": filename,
+                        "versions": []
+                    }
+                    version = 0
+                else:
+                    raise
+            
+            # Save data to S3
+            data_key = self._get_artifact_data_key(app_name, user_id, session_id, filename, version)
+            self.s3_client.put_object(
+                Bucket=self.bucket_name,
+                Key=data_key,
+                Body=data
+            )
+            
+            # Update metadata
+            version_info = {
+                "version": version,
+                "mime_type": mime_type,
+                "created_at": datetime.utcnow().isoformat(),
+                "data_key": data_key
+            }
+            metadata["versions"].append(version_info)
+            
+            # Save metadata to S3
+            self.s3_client.put_object(
+                Bucket=self.bucket_name,
+                Key=metadata_key,
+                Body=json.dumps(metadata, indent=2).encode('utf-8')
+            )
+            
+            return version
+        except ClientError 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_key = self._get_artifact_key(app_name, user_id, session_id, filename)
+            
+            try:
+                response = self.s3_client.get_object(Bucket=self.bucket_name, Key=metadata_key)
+                metadata = json.loads(response['Body'].read().decode('utf-8'))
+            except ClientError as e:
+                if e.response['Error']['Code'] == 'NoSuchKey':
+                    return None
+                else:
+                    raise
+            
+            # 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
+            try:
+                data_response = self.s3_client.get_object(
+                    Bucket=self.bucket_name,
+                    Key=version_info["data_key"]
+                )
+                data = data_response['Body'].read()
+            except ClientError as e:
+                if e.response['Error']['Code'] == 'NoSuchKey':
+                    return None
+                else:
+                    raise
+            
+            # Create Part from blob data
+            return self._deserialize_blob(data, version_info["mime_type"])
+        except ClientError 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:
+            # List objects with the prefix
+            prefix = f"{self.prefix}/{app_name}/{user_id}/{session_id}/"
+            response = self.s3_client.list_objects_v2(
+                Bucket=self.bucket_name,
+                Prefix=prefix,
+                Delimiter='/'
+            )
+            
+            artifact_keys = []
+            if 'Contents' in response:
+                for obj in response['Contents']:
+                    key = obj['Key']
+                    # Check if it's a metadata file
+                    if key.endswith('.json') and key.startswith(prefix):
+                        # Extract filename from metadata key
+                        filename = key[len(prefix):-5]  # Remove prefix and .json extension
+                        artifact_keys.append(filename)
+            
+            return artifact_keys
+        except ClientError 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_key = self._get_artifact_key(app_name, user_id, session_id, filename)
+            
+            try:
+                response = self.s3_client.get_object(Bucket=self.bucket_name, Key=metadata_key)
+                metadata = json.loads(response['Body'].read().decode('utf-8'))
+                
+                # Delete all version data files
+                for version_info in metadata.get("versions", []):
+                    self.s3_client.delete_object(
+                        Bucket=self.bucket_name,
+                        Key=version_info["data_key"]
+                    )
+                
+                # Delete metadata file
+                self.s3_client.delete_object(
+                    Bucket=self.bucket_name,
+                    Key=metadata_key
+                )
+            except ClientError as e:
+                if e.response['Error']['Code'] != 'NoSuchKey':
+                    raise
+        except ClientError 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_key = self._get_artifact_key(app_name, user_id, session_id, filename)
+            
+            try:
+                response = self.s3_client.get_object(Bucket=self.bucket_name, Key=metadata_key)
+                metadata = json.loads(response['Body'].read().decode('utf-8'))
+            except ClientError as e:
+                if e.response['Error']['Code'] == 'NoSuchKey':
+                    return []
+                else:
+                    raise
+            
+            versions = metadata.get("versions", [])
+            return [v["version"] for v in versions]
+        except ClientError as e:
+            raise RuntimeError(f"Failed to list versions: {e}")

google_adk_extras/artifacts/sql_artifact_service.py

@@ -0,0 +1,255 @@
+"""SQL-based artifact service implementation using SQLAlchemy."""
+
+import json
+from typing import Optional, List
+from datetime import datetime, timezone
+
+try:
+    from sqlalchemy import create_engine, Column, String, Text, Integer, DateTime, LargeBinary
+    from sqlalchemy.orm import declarative_base, sessionmaker
+    from sqlalchemy.exc import SQLAlchemyError
+except ImportError:
+    raise ImportError(
+        "SQLAlchemy is required for SQLArtifactService. "
+        "Install it with: pip install sqlalchemy"
+    )
+
+from google.genai import types
+from .base_custom_artifact_service import BaseCustomArtifactService
+
+
+# Use the modern declarative_base import
+Base = declarative_base()
+
+
+class SQLArtifactModel(Base):
+    """SQLAlchemy model for storing artifacts."""
+    __tablename__ = 'adk_artifacts'
+
+    # Composite primary key
+    app_name = Column(String, primary_key=True)
+    user_id = Column(String, primary_key=True)
+    session_id = Column(String, primary_key=True)
+    filename = Column(String, primary_key=True)
+    version = Column(Integer, primary_key=True)
+    
+    # Artifact data
+    mime_type = Column(String, nullable=False)
+    data = Column(LargeBinary, nullable=False)  # Blob data
+    created_at = Column(DateTime(timezone=True), nullable=False, default=lambda: datetime.now(timezone.utc))
+    
+    # Metadata
+    metadata_json = Column(Text, nullable=True)  # Additional metadata as JSON
+
+
+class SQLArtifactService(BaseCustomArtifactService):
+    """SQL-based artifact service implementation."""
+
+    def __init__(self, database_url: str):
+        """Initialize the SQL artifact service.
+        
+        Args:
+            database_url: Database connection string (e.g., 'sqlite:///artifacts.db')
+        """
+        super().__init__()
+        self.database_url = database_url
+        self.engine: Optional[object] = None
+        self.session_local: Optional[object] = None
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the database connection and create tables."""
+        try:
+            self.engine = create_engine(self.database_url)
+            Base.metadata.create_all(self.engine)
+            self.session_local = sessionmaker(
+                autocommit=False, 
+                autoflush=False, 
+                bind=self.engine
+            )
+        except SQLAlchemyError as e:
+            raise RuntimeError(f"Failed to initialize SQL artifact service: {e}")
+
+    async def _cleanup_impl(self) -> None:
+        """Clean up database connections."""
+        if self.engine:
+            self.engine.dispose()
+            self.engine = None
+        self.session_local = None
+
+    def _get_db_session(self):
+        """Get a database session."""
+        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."""
+        if part.inline_data:
+            return part.inline_data.data, part.inline_data.mime_type or "application/octet-stream"
+        else:
+            # If it's not inline data, we need to handle other types
+            # For now, we'll raise an error - in a full implementation,
+            # we'd need to handle other Part types
+            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."""
+        db_session = self._get_db_session()
+        try:
+            # Extract blob data
+            data, mime_type = self._serialize_blob(artifact)
+            
+            # Get the next version number
+            latest_version_result = db_session.query(SQLArtifactModel).filter(
+                SQLArtifactModel.app_name == app_name,
+                SQLArtifactModel.user_id == user_id,
+                SQLArtifactModel.session_id == session_id,
+                SQLArtifactModel.filename == filename
+            ).order_by(SQLArtifactModel.version.desc()).first()
+            
+            version = (latest_version_result.version + 1) if latest_version_result else 0
+            
+            # Create artifact model
+            db_artifact = SQLArtifactModel(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename,
+                version=version,
+                mime_type=mime_type,
+                data=data
+            )
+            
+            # Save to database
+            db_session.add(db_artifact)
+            db_session.commit()
+            
+            return version
+        except SQLAlchemyError as e:
+            db_session.rollback()
+            raise RuntimeError(f"Failed to save artifact: {e}")
+        finally:
+            db_session.close()
+
+    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."""
+        db_session = self._get_db_session()
+        try:
+            query = db_session.query(SQLArtifactModel).filter(
+                SQLArtifactModel.app_name == app_name,
+                SQLArtifactModel.user_id == user_id,
+                SQLArtifactModel.session_id == session_id,
+                SQLArtifactModel.filename == filename
+            )
+            
+            if version is not None:
+                query = query.filter(SQLArtifactModel.version == version)
+            else:
+                # Get the latest version
+                query = query.order_by(SQLArtifactModel.version.desc())
+            
+            db_artifact = query.first()
+            
+            if not db_artifact:
+                return None
+            
+            # Create Part from blob data
+            return self._deserialize_blob(db_artifact.data, db_artifact.mime_type)
+        except SQLAlchemyError as e:
+            raise RuntimeError(f"Failed to load artifact: {e}")
+        finally:
+            db_session.close()
+
+    async def _list_artifact_keys_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+    ) -> List[str]:
+        """Implementation of artifact key listing."""
+        db_session = self._get_db_session()
+        try:
+            # Get distinct filenames
+            filenames = db_session.query(SQLArtifactModel.filename).filter(
+                SQLArtifactModel.app_name == app_name,
+                SQLArtifactModel.user_id == user_id,
+                SQLArtifactModel.session_id == session_id
+            ).distinct().all()
+            
+            return [filename[0] for filename in filenames]
+        except SQLAlchemyError as e:
+            raise RuntimeError(f"Failed to list artifact keys: {e}")
+        finally:
+            db_session.close()
+
+    async def _delete_artifact_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> None:
+        """Implementation of artifact deletion."""
+        db_session = self._get_db_session()
+        try:
+            # Delete all versions of the artifact
+            db_session.query(SQLArtifactModel).filter(
+                SQLArtifactModel.app_name == app_name,
+                SQLArtifactModel.user_id == user_id,
+                SQLArtifactModel.session_id == session_id,
+                SQLArtifactModel.filename == filename
+            ).delete()
+            
+            db_session.commit()
+        except SQLAlchemyError as e:
+            db_session.rollback()
+            raise RuntimeError(f"Failed to delete artifact: {e}")
+        finally:
+            db_session.close()
+
+    async def _list_versions_impl(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+    ) -> List[int]:
+        """Implementation of version listing."""
+        db_session = self._get_db_session()
+        try:
+            versions = db_session.query(SQLArtifactModel.version).filter(
+                SQLArtifactModel.app_name == app_name,
+                SQLArtifactModel.user_id == user_id,
+                SQLArtifactModel.session_id == session_id,
+                SQLArtifactModel.filename == filename
+            ).order_by(SQLArtifactModel.version.asc()).all()
+            
+            return [version[0] for version in versions]
+        except SQLAlchemyError as e:
+            raise RuntimeError(f"Failed to list versions: {e}")
+        finally:
+            db_session.close()

google_adk_extras/memory/__init__.py

@@ -0,0 +1,15 @@
+"""Custom ADK memory services package."""
+
+from .base_custom_memory_service import BaseCustomMemoryService
+from .sql_memory_service import SQLMemoryService
+from .mongo_memory_service import MongoMemoryService
+from .redis_memory_service import RedisMemoryService
+from .yaml_file_memory_service import YamlFileMemoryService
+
+__all__ = [
+    "BaseCustomMemoryService",
+    "SQLMemoryService",
+    "MongoMemoryService",
+    "RedisMemoryService",
+    "YamlFileMemoryService",
+]

google_adk_extras/memory/base_custom_memory_service.py

@@ -0,0 +1,90 @@
+"""Base class for custom memory services."""
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+from google.adk.memory.base_memory_service import BaseMemoryService
+
+if TYPE_CHECKING:
+    from google.adk.sessions.session import Session
+
+
+class BaseCustomMemoryService(BaseMemoryService):
+    """Base class for custom memory services with common functionality."""
+
+    def __init__(self):
+        """Initialize the base custom memory service."""
+        super().__init__()
+        self._initialized = False
+
+    @abstractmethod
+    async def _add_session_to_memory_impl(self, session: "Session") -> None:
+        """Implementation of adding a session to memory.
+        
+        Args:
+            session: The session to add to memory.
+        """
+
+    @abstractmethod
+    async def _search_memory_impl(
+        self, *, app_name: str, user_id: str, query: str
+    ) -> "SearchMemoryResponse":
+        """Implementation of searching memory.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The id of the user.
+            query: The query to search for.
+            
+        Returns:
+            A SearchMemoryResponse containing the matching memories.
+        """
+
+    async def add_session_to_memory(self, session: "Session") -> None:
+        """Add a session to the memory service.
+        
+        Args:
+            session: The session to add.
+        """
+        if not self._initialized:
+            await self.initialize()
+        await self._add_session_to_memory_impl(session)
+
+    async def search_memory(
+        self, *, app_name: str, user_id: str, query: str
+    ) -> "SearchMemoryResponse":
+        """Search for sessions that match the query.
+        
+        Args:
+            app_name: The name of the application.
+            user_id: The id of the user.
+            query: The query to search for.
+            
+        Returns:
+            A SearchMemoryResponse containing the matching memories.
+        """
+        if not self._initialized:
+            await self.initialize()
+        return await self._search_memory_impl(
+            app_name=app_name, user_id=user_id, query=query
+        )
+
+    async def initialize(self) -> None:
+        """Initialize the memory service."""
+        if not self._initialized:
+            await self._initialize_impl()
+            self._initialized = True
+
+    async def cleanup(self) -> None:
+        """Clean up the memory service."""
+        if self._initialized:
+            await self._cleanup_impl()
+            self._initialized = False
+
+    @abstractmethod
+    async def _initialize_impl(self) -> None:
+        """Implementation of initialization."""
+
+    @abstractmethod
+    async def _cleanup_impl(self) -> None:
+        """Implementation of cleanup."""