diff --git a/skyvern/forge/sdk/artifact/storage/s3.py b/skyvern/forge/sdk/artifact/storage/s3.py index 2f3b64f1..466b7663 100644 --- a/skyvern/forge/sdk/artifact/storage/s3.py +++ b/skyvern/forge/sdk/artifact/storage/s3.py @@ -113,7 +113,7 @@ class S3Storage(BaseStorage): uri = uri.replace(f".{file_ext}", f".{file_ext}.zst") artifact.uri = uri - sc = await self._get_storage_class_for_org(artifact.organization_id, self.bucket) + sc = await self._get_storage_class_for_org(artifact.organization_id, self.bucket, len(data)) tags = await self._get_tags_for_org(artifact.organization_id) LOG.debug( "Storing artifact", @@ -125,7 +125,12 @@ class S3Storage(BaseStorage): ) await self.async_client.upload_file(uri, data, storage_class=sc, tags=tags) - async def _get_storage_class_for_org(self, organization_id: str, bucket: str) -> S3StorageClass: + async def _get_storage_class_for_org( + self, + organization_id: str, + bucket: str, + object_size_bytes: int | None = None, + ) -> S3StorageClass: return S3StorageClass.STANDARD async def _get_tags_for_org(self, organization_id: str) -> dict[str, str]: @@ -147,7 +152,7 @@ class S3Storage(BaseStorage): return await self.async_client.create_presigned_urls([artifact.uri for artifact in artifacts]) async def store_artifact_from_path(self, artifact: Artifact, path: str) -> None: - sc = await self._get_storage_class_for_org(artifact.organization_id, self.bucket) + sc = await self._get_storage_class_for_org(artifact.organization_id, self.bucket, os.path.getsize(path)) tags = await self._get_tags_for_org(artifact.organization_id) LOG.debug( "Storing artifact from path", diff --git a/skyvern/forge/sdk/artifact/storage/test_s3_storage.py b/skyvern/forge/sdk/artifact/storage/test_s3_storage.py index 9d1cc724..5984be18 100644 --- a/skyvern/forge/sdk/artifact/storage/test_s3_storage.py +++ b/skyvern/forge/sdk/artifact/storage/test_s3_storage.py @@ -36,7 +36,12 @@ class S3StorageForTests(S3Storage): async def _get_tags_for_org(self, organization_id: str) -> dict[str, str]: return {"dummy": f"org-{organization_id}", "test": "jerry"} - async def _get_storage_class_for_org(self, organization_id: str, bucket: str) -> S3StorageClass: + async def _get_storage_class_for_org( + self, + organization_id: str, + bucket: str, + object_size_bytes: int | None = None, + ) -> S3StorageClass: return S3StorageClass.ONEZONE_IA