agno 0.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/cloud/aws/base.py

@@ -0,0 +1,214 @@
+from typing import Any, Optional
+
+from pydantic import BaseModel, ConfigDict
+
+from agno.cloud.aws.s3.api_client import AwsApiClient
+from agno.utils.log import logger
+
+
+class AwsResource(BaseModel):
+    """Base class for AWS Resources."""
+
+    # Resource name (required)
+    name: str
+
+    # Resource type
+    resource_type: Optional[str] = None
+    active_resource: Optional[Any] = None
+
+    skip_delete: bool = False
+    skip_create: bool = False
+    skip_read: bool = False
+    skip_update: bool = False
+
+    wait_for_create: bool = True
+    wait_for_update: bool = True
+    wait_for_delete: bool = True
+    waiter_delay: int = 30
+    waiter_max_attempts: int = 50
+    save_output: bool = False
+    use_cache: bool = True
+
+    service_name: str
+    service_client: Optional[Any] = None
+    service_resource: Optional[Any] = None
+
+    aws_region: Optional[str] = None
+    aws_profile: Optional[str] = None
+
+    aws_client: Optional[AwsApiClient] = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True)
+
+    def get_resource_name(self) -> str:
+        return self.name or self.__class__.__name__
+
+    def get_resource_type(self) -> str:
+        if self.resource_type is None:
+            return self.__class__.__name__
+        return self.resource_type
+
+    def get_service_client(self, aws_client: AwsApiClient):
+        from boto3 import session
+
+        if self.service_client is None:
+            boto3_session: session = aws_client.boto3_session
+            self.service_client = boto3_session.client(service_name=self.service_name)
+        return self.service_client
+
+    def get_service_resource(self, aws_client: AwsApiClient):
+        from boto3 import session
+
+        if self.service_resource is None:
+            boto3_session: session = aws_client.boto3_session
+            self.service_resource = boto3_session.resource(service_name=self.service_name)
+        return self.service_resource
+
+    def get_aws_region(self) -> Optional[str]:
+        if self.aws_region:
+            return self.aws_region
+
+        from os import getenv
+
+        aws_region_env = getenv("AWS_REGION")
+        if aws_region_env is not None:
+            logger.debug(f"{'AWS_REGION'}: {aws_region_env}")
+            self.aws_region = aws_region_env
+        return self.aws_region
+
+    def get_aws_profile(self) -> Optional[str]:
+        if self.aws_profile:
+            return self.aws_profile
+
+        from os import getenv
+
+        aws_profile_env = getenv("AWS_PROFILE")
+        if aws_profile_env is not None:
+            logger.debug(f"{'AWS_PROFILE'}: {aws_profile_env}")
+            self.aws_profile = aws_profile_env
+        return self.aws_profile
+
+    def get_aws_client(self) -> AwsApiClient:
+        if self.aws_client is not None:
+            return self.aws_client
+        self.aws_client = AwsApiClient(aws_region=self.get_aws_region(), aws_profile=self.get_aws_profile())
+        return self.aws_client
+
+    def _read(self, aws_client: AwsApiClient) -> Any:
+        logger.warning(f"@_read method not defined for {self.get_resource_name()}")
+        return True
+
+    def read(self, aws_client: Optional[AwsApiClient] = None) -> Any:
+        """Reads the resource from Aws"""
+        # Step 1: Use cached value if available
+        if self.use_cache and self.active_resource is not None:
+            return self.active_resource
+
+        # Step 2: Skip resource creation if skip_read = True
+        if self.skip_read:
+            print(f"Skipping read: {self.get_resource_name()}")
+            return True
+
+        # Step 3: Read resource
+        client: AwsApiClient = aws_client or self.get_aws_client()
+        return self._read(client)
+
+    def is_active(self, aws_client: AwsApiClient) -> bool:
+        """Returns True if the resource is active on Aws"""
+        _resource = self.read(aws_client=aws_client)
+        return True if _resource is not None else False
+
+    def _create(self, aws_client: AwsApiClient) -> bool:
+        logger.warning(f"@_create method not defined for {self.get_resource_name()}")
+        return True
+
+    def create(self, aws_client: Optional[AwsApiClient] = None) -> bool:
+        """Creates the resource on Aws"""
+
+        # Step 1: Skip resource creation if skip_create = True
+        if self.skip_create:
+            print(f"Skipping create: {self.get_resource_name()}")
+            return True
+
+        # Step 2: Check if resource is active and use_cache = True
+        client: AwsApiClient = aws_client or self.get_aws_client()
+        if self.use_cache and self.is_active(client):
+            self.resource_created = True
+            print(f"{self.get_resource_type()}: {self.get_resource_name()} already exists")
+        # Step 3: Create the resource
+        else:
+            self.resource_created = self._create(client)
+            if self.resource_created:
+                print(f"{self.get_resource_type()}: {self.get_resource_name()} created")
+
+        # Step 4: Run post create steps
+        if self.resource_created:
+            logger.debug(f"Running post-create for {self.get_resource_type()}: {self.get_resource_name()}")
+            return self.post_create(client)
+        logger.error(f"Failed to create {self.get_resource_type()}: {self.get_resource_name()}")
+        return self.resource_created
+
+    def post_create(self, aws_client: AwsApiClient) -> bool:
+        return True
+
+    def _update(self, aws_client: AwsApiClient) -> Any:
+        logger.warning(f"@_update method not defined for {self.get_resource_name()}")
+        return True
+
+    def update(self, aws_client: Optional[AwsApiClient] = None) -> bool:
+        """Updates the resource on Aws"""
+
+        # Step 1: Skip resource update if skip_update = True
+        if self.skip_update:
+            print(f"Skipping update: {self.get_resource_name()}")
+            return True
+
+        # Step 2: Update the resource
+        client: AwsApiClient = aws_client or self.get_aws_client()
+        if self.is_active(client):
+            self.resource_updated = self._update(client)
+        else:
+            print(f"{self.get_resource_type()}: {self.get_resource_name()} does not exist")
+            return True
+
+        # Step 3: Run post update steps
+        if self.resource_updated:
+            print(f"{self.get_resource_type()}: {self.get_resource_name()} updated")
+            logger.debug(f"Running post-update for {self.get_resource_type()}: {self.get_resource_name()}")
+            return self.post_update(client)
+        logger.error(f"Failed to update {self.get_resource_type()}: {self.get_resource_name()}")
+        return self.resource_updated
+
+    def post_update(self, aws_client: AwsApiClient) -> bool:
+        return True
+
+    def _delete(self, aws_client: AwsApiClient) -> Any:
+        logger.warning(f"@_delete method not defined for {self.get_resource_name()}")
+        return True
+
+    def delete(self, aws_client: Optional[AwsApiClient] = None) -> bool:
+        """Deletes the resource from Aws"""
+
+        # Step 1: Skip resource deletion if skip_delete = True
+        if self.skip_delete:
+            print(f"Skipping delete: {self.get_resource_name()}")
+            return True
+
+        # Step 2: Delete the resource
+        client: AwsApiClient = aws_client or self.get_aws_client()
+        if self.is_active(client):
+            self.resource_deleted = self._delete(client)
+        else:
+            print(f"{self.get_resource_type()}: {self.get_resource_name()} does not exist")
+            return True
+
+        # Step 3: Run post delete steps
+        if self.resource_deleted:
+            print(f"{self.get_resource_type()}: {self.get_resource_name()} deleted")
+            logger.debug(f"Running post-delete for {self.get_resource_type()}: {self.get_resource_name()}.")
+            return self.post_delete(client)
+        logger.error(f"Failed to delete {self.get_resource_type()}: {self.get_resource_name()}")
+        return self.resource_deleted
+
+    def post_delete(self, aws_client: AwsApiClient) -> bool:
+        return True

agno/cloud/aws/s3/__init__.py

@@ -0,0 +1,2 @@
+from agno.cloud.aws.s3.bucket import S3Bucket
+from agno.cloud.aws.s3.object import S3Object

agno/cloud/aws/s3/api_client.py

@@ -0,0 +1,43 @@
+from typing import Any, Optional
+
+from agno.utils.log import logger
+
+
+class AwsApiClient:
+    def __init__(
+        self,
+        aws_region: Optional[str] = None,
+        aws_profile: Optional[str] = None,
+    ):
+        super().__init__()
+        self.aws_region: Optional[str] = aws_region
+        self.aws_profile: Optional[str] = aws_profile
+
+        # aws boto3 session
+        self._boto3_session: Optional[Any] = None
+        logger.debug("**-+-** AwsApiClient created")
+
+    def create_boto3_session(self) -> Optional[Any]:
+        """Create a boto3 session"""
+        import boto3
+
+        logger.debug("Creating boto3.Session")
+        try:
+            self._boto3_session = boto3.Session(
+                region_name=self.aws_region,
+                profile_name=self.aws_profile,
+            )
+            logger.debug("**-+-** boto3.Session created")
+            logger.debug(f"\taws_region: {self._boto3_session.region_name}")
+            logger.debug(f"\taws_profile: {self._boto3_session.profile_name}")
+        except Exception as e:
+            logger.error("Could not connect to aws. Please confirm aws cli is installed and configured")
+            logger.error(e)
+            exit(0)
+        return self._boto3_session
+
+    @property
+    def boto3_session(self) -> Optional[Any]:
+        if self._boto3_session is None:
+            self._boto3_session = self.create_boto3_session()
+        return self._boto3_session

agno/cloud/aws/s3/bucket.py

@@ -0,0 +1,195 @@
+from typing import Any, Dict, List, Optional
+
+from typing_extensions import Literal
+
+from agno.cloud.aws.base import AwsResource
+from agno.cloud.aws.s3.api_client import AwsApiClient
+from agno.cloud.aws.s3.object import S3Object
+from agno.utils.log import logger
+
+
+class S3Bucket(AwsResource):
+    """
+    Reference:
+    - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#service-resource
+    """
+
+    resource_type: str = "s3"
+    service_name: str = "s3"
+
+    # Name of the bucket
+    name: str
+    # The canned ACL to apply to the bucket.
+    acl: Optional[Literal["private", "public-read", "public-read-write", "authenticated-read"]] = None
+    grant_full_control: Optional[str] = None
+    grant_read: Optional[str] = None
+    grant_read_ACP: Optional[str] = None
+    grant_write: Optional[str] = None
+    grant_write_ACP: Optional[str] = None
+    object_lock_enabled_for_bucket: Optional[bool] = None
+    object_ownership: Optional[Literal["BucketOwnerPreferred", "ObjectWriter", "BucketOwnerEnforced"]] = None
+
+    @property
+    def uri(self) -> str:
+        """Returns the URI of the s3.Bucket
+
+        Returns:
+            str: The URI of the s3.Bucket
+        """
+        return f"s3://{self.name}"
+
+    def get_resource(self, aws_client: Optional[AwsApiClient] = None) -> Optional[Any]:
+        """Returns the s3.Bucket
+
+        Args:
+            aws_client: The AwsApiClient for the current cluster
+        """
+        client: AwsApiClient = aws_client or self.get_aws_client()
+        service_resource = self.get_service_resource(client)
+        return service_resource.Bucket(name=self.name)
+
+    def _create(self, aws_client: AwsApiClient) -> bool:
+        """Creates the s3.Bucket
+
+        Args:
+            aws_client: The AwsApiClient for the current cluster
+        """
+        print(f"Creating {self.get_resource_type()}: {self.get_resource_name()}")
+
+        # Step 1: Build bucket configuration
+        # Bucket names are GLOBALLY unique!
+        # AWS will give you the IllegalLocationConstraintException if you collide
+        # with an already existing bucket if you've specified a region different than
+        # the region of the already existing bucket. If you happen to guess the correct region of the
+        # existing bucket it will give you the BucketAlreadyExists exception.
+        bucket_configuration = None
+        if aws_client.aws_region is not None and aws_client.aws_region != "us-east-1":
+            bucket_configuration = {"LocationConstraint": aws_client.aws_region}
+
+        # create a dict of args which are not null, otherwise aws type validation fails
+        not_null_args: Dict[str, Any] = {}
+        if bucket_configuration:
+            not_null_args["CreateBucketConfiguration"] = bucket_configuration
+        if self.acl:
+            not_null_args["ACL"] = self.acl
+        if self.grant_full_control:
+            not_null_args["GrantFullControl"] = self.grant_full_control
+        if self.grant_read:
+            not_null_args["GrantRead"] = self.grant_read
+        if self.grant_read_ACP:
+            not_null_args["GrantReadACP"] = self.grant_read_ACP
+        if self.grant_write:
+            not_null_args["GrantWrite"] = self.grant_write
+        if self.grant_write_ACP:
+            not_null_args["GrantWriteACP"] = self.grant_write_ACP
+        if self.object_lock_enabled_for_bucket:
+            not_null_args["ObjectLockEnabledForBucket"] = self.object_lock_enabled_for_bucket
+        if self.object_ownership:
+            not_null_args["ObjectOwnership"] = self.object_ownership
+
+        # Step 2: Create Bucket
+        service_client = self.get_service_client(aws_client)
+        try:
+            response = service_client.create_bucket(
+                Bucket=self.name,
+                **not_null_args,
+            )
+            logger.debug(f"Response: {response}")
+            bucket_location = response.get("Location")
+            if bucket_location is not None:
+                logger.debug(f"Bucket created: {bucket_location}")
+                self.active_resource = response
+                return True
+        except Exception as e:
+            logger.error(f"{self.get_resource_type()} could not be created.")
+            logger.error(e)
+        return False
+
+    def post_create(self, aws_client: AwsApiClient) -> bool:
+        # Wait for Bucket to be created
+        if self.wait_for_create:
+            try:
+                print(f"Waiting for {self.get_resource_type()} to be created.")
+                waiter = self.get_service_client(aws_client).get_waiter("bucket_exists")
+                waiter.wait(
+                    Bucket=self.name,
+                    WaiterConfig={
+                        "Delay": self.waiter_delay,
+                        "MaxAttempts": self.waiter_max_attempts,
+                    },
+                )
+            except Exception as e:
+                logger.error("Waiter failed.")
+                logger.error(e)
+        return True
+
+    def _read(self, aws_client: AwsApiClient) -> Optional[Any]:
+        """Returns the s3.Bucket
+
+        Args:
+            aws_client: The AwsApiClient for the current cluster
+        """
+        logger.debug(f"Reading {self.get_resource_type()}: {self.get_resource_name()}")
+
+        from botocore.exceptions import ClientError
+
+        try:
+            service_resource = self.get_service_resource(aws_client)
+            bucket = service_resource.Bucket(name=self.name)
+            bucket.load()
+            creation_date = bucket.creation_date
+            logger.debug(f"Bucket creation_date: {creation_date}")
+            if creation_date is not None:
+                logger.debug(f"Bucket found: {bucket.name}")
+                self.active_resource = {
+                    "name": bucket.name,
+                    "creation_date": creation_date,
+                }
+        except ClientError as ce:
+            logger.debug(f"ClientError: {ce}")
+        except Exception as e:
+            logger.error(f"Error reading {self.get_resource_type()}.")
+            logger.error(e)
+        return self.active_resource
+
+    def _delete(self, aws_client: AwsApiClient) -> bool:
+        """Deletes the s3.Bucket
+
+        Args:
+            aws_client: The AwsApiClient for the current cluster
+        """
+        print(f"Deleting {self.get_resource_type()}: {self.get_resource_name()}")
+
+        service_client = self.get_service_client(aws_client)
+        self.active_resource = None
+        try:
+            response = service_client.delete_bucket(Bucket=self.name)
+            logger.debug(f"Response: {response}")
+            return True
+        except Exception as e:
+            logger.error(f"{self.get_resource_type()} could not be deleted.")
+            logger.error("Please try again or delete resources manually.")
+            logger.error(e)
+        return False
+
+    def get_objects(self, aws_client: Optional[AwsApiClient] = None, prefix: Optional[str] = None) -> List[Any]:
+        """Returns a list of s3.Object objects for the s3.Bucket
+
+        Args:
+            aws_client: The AwsApiClient for the current cluster
+            prefix: Prefix to filter objects by
+        """
+        bucket = self.get_resource(aws_client)
+        if bucket is None:
+            logger.warning(f"Could not get bucket: {self.name}")
+            return []
+
+        logger.debug(f"Getting objects for bucket: {bucket.name}")
+        # Get all objects in bucket
+        object_summaries = bucket.objects.all()
+        all_objects: List[S3Object] = []
+        for object_summary in object_summaries:
+            if prefix is not None and not object_summary.key.startswith(prefix):
+                continue
+            all_objects.append(S3Object(bucket_name=bucket.name, name=object_summary.key))
+        return all_objects

agno/cloud/aws/s3/object.py

@@ -0,0 +1,57 @@
+from pathlib import Path
+from typing import Any, Optional
+
+from agno.cloud.aws.base import AwsResource
+from agno.cloud.aws.s3.api_client import AwsApiClient
+from agno.utils.log import logger
+
+
+class S3Object(AwsResource):
+    """
+    Reference:
+    - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3/object/index.html
+    """
+
+    resource_type: str = "s3"
+    service_name: str = "s3"
+
+    # The Object’s bucket_name identifier. This must be set.
+    bucket_name: str
+
+    @property
+    def uri(self) -> str:
+        """Returns the URI of the s3.Object
+
+        Returns:
+            str: The URI of the s3.Object
+        """
+        return f"s3://{self.bucket_name}/{self.name}"
+
+    def get_resource(self, aws_client: Optional[AwsApiClient] = None) -> Any:
+        """Returns the s3.Object
+
+        Args:
+            aws_client: The AwsApiClient for the current cluster
+
+        Returns:
+            The s3.Object
+        """
+        client: AwsApiClient = aws_client or self.get_aws_client()
+        service_resource = self.get_service_resource(client)
+        return service_resource.Object(
+            bucket_name=self.bucket_name,
+            key=self.name,
+        )
+
+    def download(self, path: Path, aws_client: Optional[AwsApiClient] = None) -> None:
+        """Downloads the s3.Object to the specified path
+
+        Args:
+            path: The path to download the s3.Object to
+            aws_client: The AwsApiClient for the current cluster
+        """
+        logger.info(f"Downloading {self.uri} to {path}")
+        object_resource = self.get_resource(aws_client=aws_client)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open(mode="wb") as f:
+            object_resource.download_fileobj(f)

agno/compression/__init__.py

@@ -0,0 +1,3 @@
+from agno.compression.manager import CompressionManager
+
+__all__ = ["CompressionManager"]