solace-agent-mesh 1.0.6__py3-none-any.whl → 1.0.7__py3-none-any.whl

solace_agent_mesh/agent/adk/artifacts/__init__.py

@@ -0,0 +1 @@
+"""Artifact service implementations for Solace Agent Mesh."""

solace_agent_mesh/agent/adk/{filesystem_artifact_service.py → artifacts/filesystem_artifact_service.py}

@@ -2,13 +2,12 @@
 An ADK ArtifactService implementation using the local filesystem for storage.
 """
 
-import os
+import asyncio
 import json
-import shutil
 import logging
-import asyncio
+import os
+import shutil
 import unicodedata
-from typing import Optional, List
 
 from google.adk.artifacts import BaseArtifactService
 from google.genai import types as adk_types
@@ -124,7 +123,7 @@ class FilesystemArtifactService(BaseArtifactService):
                 artifact_dir,
                 e,
             )
-            raise IOError(f"Could not create artifact directory: {e}") from e
+            raise OSError(f"Could not create artifact directory: {e}") from e
 
         versions = await self.list_versions(
             app_name=app_name,
@@ -164,7 +163,7 @@ class FilesystemArtifactService(BaseArtifactService):
                 version,
             )
             return version
-        except (IOError, OSError, ValueError, TypeError) as e:
+        except (OSError, ValueError, TypeError) as e:
             logger.error(
                 "%sFailed to save artifact '%s' version %d: %s",
                 log_prefix,
@@ -176,7 +175,7 @@ class FilesystemArtifactService(BaseArtifactService):
                 await asyncio.to_thread(os.remove, version_path)
             if await asyncio.to_thread(os.path.exists, metadata_path):
                 await asyncio.to_thread(os.remove, metadata_path)
-            raise IOError(f"Failed to save artifact version {version}: {e}") from e
+            raise OSError(f"Failed to save artifact version {version}: {e}") from e
 
     @override
     async def load_artifact(
@@ -186,9 +185,9 @@ class FilesystemArtifactService(BaseArtifactService):
         user_id: str,
         session_id: str,
         filename: str,
-        version: Optional[int] = None,
-    ) -> Optional[adk_types.Part]:
-        log_prefix = "[FSArtifact:Load] "
+        version: int | None = None,
+    ) -> adk_types.Part | None:
+        log_prefix = f"[FSArtifact:Load:{filename}] "
         filename = self._normalize_filename_unicode(filename)
         artifact_dir = self._get_artifact_dir(app_name, user_id, session_id, filename)
 
@@ -228,7 +227,7 @@ class FilesystemArtifactService(BaseArtifactService):
         try:
 
             def _read_metadata_file():
-                with open(metadata_path, "r", encoding="utf-8") as f:
+                with open(metadata_path, encoding="utf-8") as f:
                     return json.load(f)
 
             metadata = await asyncio.to_thread(_read_metadata_file)
@@ -253,7 +252,7 @@ class FilesystemArtifactService(BaseArtifactService):
             )
             return artifact_part
 
-        except (IOError, OSError, json.JSONDecodeError) as e:
+        except (OSError, json.JSONDecodeError) as e:
             logger.error(
                 "%sFailed to load artifact '%s' version %d: %s",
                 log_prefix,
@@ -266,7 +265,7 @@ class FilesystemArtifactService(BaseArtifactService):
     @override
     async def list_artifact_keys(
         self, *, app_name: str, user_id: str, session_id: str
-    ) -> List[str]:
+    ) -> list[str]:
         log_prefix = "[FSArtifact:ListKeys] "
         filenames = set()
         app_name_sanitized = os.path.basename(app_name)
@@ -339,8 +338,8 @@ class FilesystemArtifactService(BaseArtifactService):
     @override
     async def list_versions(
         self, *, app_name: str, user_id: str, session_id: str, filename: str
-    ) -> List[int]:
-        log_prefix = "[FSArtifact:ListVersions] "
+    ) -> list[int]:
+        log_prefix = f"[FSArtifact:ListVersions:{filename}] "
         artifact_dir = self._get_artifact_dir(app_name, user_id, session_id, filename)
         versions = []
 

solace_agent_mesh/agent/adk/artifacts/s3_artifact_service.py

@@ -0,0 +1,440 @@
+"""
+An ADK ArtifactService implementation using Amazon S3 compatible storage.
+"""
+
+import asyncio
+import logging
+import unicodedata
+
+import boto3
+from botocore.client import BaseClient
+from botocore.exceptions import BotoCoreError, ClientError, NoCredentialsError
+from google.adk.artifacts import BaseArtifactService
+from google.genai import types as adk_types
+from typing_extensions import override
+
+logger = logging.getLogger(__name__)
+
+
+class S3ArtifactService(BaseArtifactService):
+    """
+    An artifact service implementation using Amazon S3 compatible storage.
+
+    Stores artifacts in an S3-compatible bucket with a structured key format:
+    {app_name}/{user_id}/{session_id_or_user}/{filename}/{version}
+
+    Supports AWS S3 and S3-compatible APIs like MinIO.
+
+    Required S3 Permissions:
+    The IAM user or role must have the following minimum permissions for the specific bucket:
+    - s3:GetObject: Read artifacts from the bucket
+    - s3:PutObject: Store new artifacts to the bucket
+    - s3:DeleteObject: Delete artifacts from the bucket
+
+    Example IAM Policy (replace 'your-bucket-name' with actual bucket):
+    {
+        "Version": "2012-10-17",
+        "Statement": [
+            {
+                "Effect": "Allow",
+                "Action": [
+                    "s3:GetObject",
+                    "s3:PutObject",
+                    "s3:DeleteObject"
+                ],
+                "Resource": "arn:aws:s3:::your-bucket-name/*"
+            },
+            {
+                "Effect": "Allow",
+                "Action": "s3:ListBucket",
+                "Resource": "arn:aws:s3:::your-bucket-name"
+            }
+        ]
+    }
+    """
+
+    def __init__(
+        self,
+        bucket_name: str,
+        s3_client: BaseClient | None = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            bucket_name: The name of the S3 bucket to use.
+            s3_client: Optional pre-configured S3 client. If None, creates a new client.
+            **kwargs: Optional parameters for boto3 client configuration.
+
+        Raises:
+            ValueError: If bucket_name is not provided.
+            NoCredentialsError: If AWS credentials are not available.
+        """
+        if not bucket_name:
+            raise ValueError("bucket_name cannot be empty for S3ArtifactService")
+
+        self.bucket_name = bucket_name
+
+        if s3_client is None:
+            try:
+                self.s3 = boto3.client("s3", **kwargs)
+            except NoCredentialsError as e:
+                logger.error("AWS credentials not found. Please configure credentials.")
+                raise ValueError(
+                    "AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables or configure AWS credentials."
+                ) from e
+        else:
+            self.s3 = s3_client
+
+        try:
+            self.s3.head_bucket(Bucket=self.bucket_name)
+            logger.info(
+                "S3ArtifactService initialized successfully. Bucket: %s",
+                self.bucket_name,
+            )
+        except ClientError as e:
+            error_code = e.response.get("Error", {}).get("Code", "Unknown")
+            if error_code == "404":
+                logger.error("S3 bucket '%s' does not exist", self.bucket_name)
+                raise ValueError(
+                    f"S3 bucket '{self.bucket_name}' does not exist"
+                ) from e
+            elif error_code == "403":
+                logger.error("Access denied to S3 bucket '%s'", self.bucket_name)
+                raise ValueError(
+                    f"Access denied to S3 bucket '{self.bucket_name}'"
+                ) from e
+            else:
+                logger.error("Failed to access S3 bucket '%s': %s", self.bucket_name, e)
+                raise ValueError(
+                    f"Failed to access S3 bucket '{self.bucket_name}': {e}"
+                ) from e
+
+    def _file_has_user_namespace(self, filename: str) -> bool:
+        return filename.startswith("user:")
+
+    def _get_object_key(
+        self,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: int | str,
+    ) -> str:
+        """Constructs the S3 object key for an artifact."""
+        filename = self._normalize_filename_unicode(filename)
+
+        if self._file_has_user_namespace(filename):
+            filename_clean = filename.split(":", 1)[1]
+            return f"{app_name}/{user_id}/user/{filename_clean}/{version}"
+        return f"{app_name}/{user_id}/{session_id}/{filename}/{version}"
+
+    def _normalize_filename_unicode(self, filename: str) -> str:
+        """Normalizes Unicode characters in a filename to their standard form."""
+        return unicodedata.normalize("NFKC", filename)
+
+    @override
+    async def save_artifact(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        artifact: adk_types.Part,
+    ) -> int:
+        log_prefix = f"[S3Artifact:Save:{filename}] "
+
+        if not artifact.inline_data or artifact.inline_data.data is None:
+            raise ValueError("Artifact Part has no inline_data to save.")
+
+        filename = self._normalize_filename_unicode(filename)
+
+        # Get existing versions to determine next version number
+        versions = await self.list_versions(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+        )
+        version = 0 if not versions else max(versions) + 1
+
+        object_key = self._get_object_key(
+            app_name, user_id, session_id, filename, version
+        )
+
+        try:
+
+            def _put_object():
+                return self.s3.put_object(
+                    Bucket=self.bucket_name,
+                    Key=object_key,
+                    Body=artifact.inline_data.data,
+                    ContentType=artifact.inline_data.mime_type,
+                    Metadata={
+                        "original_filename": filename,
+                        "user_id": user_id,
+                        "session_id": session_id,
+                        "version": str(version),
+                    },
+                )
+
+            await asyncio.to_thread(_put_object)
+
+            logger.info(
+                "%sSaved artifact '%s' version %d successfully to S3 key: %s",
+                log_prefix,
+                filename,
+                version,
+                object_key,
+            )
+            return version
+
+        except ClientError as e:
+            logger.error(
+                "%sFailed to save artifact '%s' version %d to S3: %s",
+                log_prefix,
+                filename,
+                version,
+                e,
+            )
+            raise OSError(
+                f"Failed to save artifact version {version} to S3: {e}"
+            ) from e
+        except BotoCoreError as e:
+            logger.error(
+                "%sBotoCore error saving artifact '%s' version %d: %s",
+                log_prefix,
+                filename,
+                version,
+                e,
+            )
+            raise OSError(
+                f"BotoCore error saving artifact version {version}: {e}"
+            ) from e
+
+    @override
+    async def load_artifact(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: int | None = None,
+    ) -> adk_types.Part | None:
+        log_prefix = f"[S3Artifact:Load:{filename}] "
+        filename = self._normalize_filename_unicode(filename)
+
+        load_version = version
+        if load_version is None:
+            versions = await self.list_versions(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename,
+            )
+            if not versions:
+                logger.debug("%sNo versions found for artifact.", log_prefix)
+                return None
+            load_version = max(versions)
+            logger.debug("%sLoading latest version: %d", log_prefix, load_version)
+        else:
+            logger.debug("%sLoading specified version: %d", log_prefix, load_version)
+
+        object_key = self._get_object_key(
+            app_name, user_id, session_id, filename, load_version
+        )
+
+        try:
+
+            def _get_object():
+                return self.s3.get_object(Bucket=self.bucket_name, Key=object_key)
+
+            response = await asyncio.to_thread(_get_object)
+            data = response["Body"].read()
+            mime_type = response.get("ContentType", "application/octet-stream")
+
+            artifact_part = adk_types.Part.from_bytes(data=data, mime_type=mime_type)
+
+            logger.info(
+                "%sLoaded artifact '%s' version %d successfully (%d bytes, %s)",
+                log_prefix,
+                filename,
+                load_version,
+                len(data),
+                mime_type,
+            )
+            return artifact_part
+
+        except ClientError as e:
+            error_code = e.response.get("Error", {}).get("Code", "Unknown")
+            if error_code == "NoSuchKey":
+                logger.debug("%sArtifact not found: %s", log_prefix, object_key)
+                return None
+            else:
+                logger.error(
+                    "%sFailed to load artifact '%s' version %d from S3: %s",
+                    log_prefix,
+                    filename,
+                    load_version,
+                    e,
+                )
+                return None
+        except BotoCoreError as e:
+            logger.error(
+                "%sBotoCore error loading artifact '%s' version %d: %s",
+                log_prefix,
+                filename,
+                load_version,
+                e,
+            )
+            return None
+
+    @override
+    async def list_artifact_keys(
+        self, *, app_name: str, user_id: str, session_id: str
+    ) -> list[str]:
+        log_prefix = "[S3Artifact:ListKeys] "
+        filenames = set()
+
+        session_prefix = f"{app_name}/{user_id}/{session_id}/"
+        try:
+
+            def _list_session_objects():
+                paginator = self.s3.get_paginator("list_objects_v2")
+                return paginator.paginate(
+                    Bucket=self.bucket_name, Prefix=session_prefix
+                )
+
+            session_pages = await asyncio.to_thread(_list_session_objects)
+            for page in session_pages:
+                for obj in page.get("Contents", []):
+                    parts = obj["Key"].split("/")
+                    if len(parts) >= 5:  # scope/user/session/filename/version
+                        filename = parts[3]
+                        filenames.add(filename)
+        except ClientError as e:
+            logger.warning(
+                "%sError listing session objects with prefix '%s': %s",
+                log_prefix,
+                session_prefix,
+                e,
+            )
+
+        user_prefix = f"{app_name}/{user_id}/user/"
+        try:
+
+            def _list_user_objects():
+                paginator = self.s3.get_paginator("list_objects_v2")
+                return paginator.paginate(Bucket=self.bucket_name, Prefix=user_prefix)
+
+            user_pages = await asyncio.to_thread(_list_user_objects)
+            for page in user_pages:
+                for obj in page.get("Contents", []):
+                    parts = obj["Key"].split("/")
+                    if len(parts) >= 5:  # scope/user/user/filename/version
+                        filename = parts[3]
+                        filenames.add(f"user:{filename}")
+        except ClientError as e:
+            logger.warning(
+                "%sError listing user objects with prefix '%s': %s",
+                log_prefix,
+                user_prefix,
+                e,
+            )
+
+        sorted_filenames = sorted(list(filenames))
+        logger.debug("%sFound %d artifact keys.", log_prefix, len(sorted_filenames))
+        return sorted_filenames
+
+    @override
+    async def delete_artifact(
+        self, *, app_name: str, user_id: str, session_id: str, filename: str
+    ) -> None:
+        log_prefix = f"[S3Artifact:Delete:{filename}] "
+        filename = self._normalize_filename_unicode(filename)
+
+        # Get all versions to delete
+        versions = await self.list_versions(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+        )
+
+        if not versions:
+            logger.debug("%sNo versions found to delete for artifact.", log_prefix)
+            return
+
+        # Delete all versions
+        for version in versions:
+            object_key = self._get_object_key(
+                app_name, user_id, session_id, filename, version
+            )
+            try:
+
+                def _delete_object():
+                    return self.s3.delete_object(
+                        Bucket=self.bucket_name, Key=object_key
+                    )
+
+                await asyncio.to_thread(_delete_object)
+                logger.debug(
+                    "%sDeleted version %d: %s", log_prefix, version, object_key
+                )
+            except ClientError as e:
+                logger.warning(
+                    "%sFailed to delete version %d (%s): %s",
+                    log_prefix,
+                    version,
+                    object_key,
+                    e,
+                )
+
+        logger.info(
+            "%sDeleted artifact '%s' (%d versions)",
+            log_prefix,
+            filename,
+            len(versions),
+        )
+
+    @override
+    async def list_versions(
+        self, *, app_name: str, user_id: str, session_id: str, filename: str
+    ) -> list[int]:
+        log_prefix = f"[S3Artifact:ListVersions:{filename}] "
+        filename = self._normalize_filename_unicode(filename)
+
+        # Get the prefix for this specific artifact (without version)
+        prefix = self._get_object_key(app_name, user_id, session_id, filename, "")
+        versions = []
+
+        try:
+
+            def _list_objects():
+                paginator = self.s3.get_paginator("list_objects_v2")
+                return paginator.paginate(Bucket=self.bucket_name, Prefix=prefix)
+
+            pages = await asyncio.to_thread(_list_objects)
+            for page in pages:
+                for obj in page.get("Contents", []):
+                    parts = obj["Key"].split("/")
+                    if len(parts) >= 5:  # scope/user/session_or_user/filename/version
+                        try:
+                            version = int(parts[4])
+                            versions.append(version)
+                        except ValueError:
+                            continue  # Skip non-integer versions
+
+        except ClientError as e:
+            logger.error(
+                "%sError listing versions with prefix '%s': %s",
+                log_prefix,
+                prefix,
+                e,
+            )
+            return []
+
+        sorted_versions = sorted(versions)
+        logger.debug("%sFound versions: %s", log_prefix, sorted_versions)
+        return sorted_versions