diff --git a/skyvern/forge/sdk/artifact/manager.py b/skyvern/forge/sdk/artifact/manager.py index d7be4e18..e79eb582 100644 --- a/skyvern/forge/sdk/artifact/manager.py +++ b/skyvern/forge/sdk/artifact/manager.py @@ -72,7 +72,9 @@ class ArtifactManager: path: str | None = None, ) -> str: artifact_id = generate_artifact_id() - uri = app.STORAGE.build_uri(artifact_id, step, artifact_type) + uri = app.STORAGE.build_uri( + organization_id=step.organization_id, artifact_id=artifact_id, step=step, artifact_type=artifact_type + ) return await self._create_artifact( aio_task_primary_key=step.task_id, artifact_id=artifact_id, diff --git a/skyvern/forge/sdk/artifact/storage/base.py b/skyvern/forge/sdk/artifact/storage/base.py index 01052b9d..d9087ab5 100644 --- a/skyvern/forge/sdk/artifact/storage/base.py +++ b/skyvern/forge/sdk/artifact/storage/base.py @@ -39,7 +39,7 @@ FILE_EXTENTSION_MAP: dict[ArtifactType, str] = { class BaseStorage(ABC): @abstractmethod - def build_uri(self, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str: + def build_uri(self, *, organization_id: str, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str: pass @abstractmethod diff --git a/skyvern/forge/sdk/artifact/storage/local.py b/skyvern/forge/sdk/artifact/storage/local.py index ef94faf4..15e6b044 100644 --- a/skyvern/forge/sdk/artifact/storage/local.py +++ b/skyvern/forge/sdk/artifact/storage/local.py @@ -28,9 +28,9 @@ class LocalStorage(BaseStorage): def __init__(self, artifact_path: str = settings.ARTIFACT_STORAGE_PATH) -> None: self.artifact_path = artifact_path - def build_uri(self, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str: + def build_uri(self, *, organization_id: str, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str: file_ext = FILE_EXTENTSION_MAP[artifact_type] - return f"file://{self.artifact_path}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}" + return f"file://{self.artifact_path}/{organization_id}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}" async def retrieve_global_workflows(self) -> list[str]: file_path = Path(f"{self.artifact_path}/{settings.ENV}/global_workflows.txt") diff --git a/skyvern/forge/sdk/artifact/storage/s3.py b/skyvern/forge/sdk/artifact/storage/s3.py index 6b0ffa97..fcb93287 100644 --- a/skyvern/forge/sdk/artifact/storage/s3.py +++ b/skyvern/forge/sdk/artifact/storage/s3.py @@ -35,9 +35,9 @@ class S3Storage(BaseStorage): self.async_client = AsyncAWSClient() self.bucket = bucket or settings.AWS_S3_BUCKET_ARTIFACTS - def build_uri(self, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str: + def build_uri(self, *, organization_id: str, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str: file_ext = FILE_EXTENTSION_MAP[artifact_type] - return f"s3://{self.bucket}/{self._PATH_VERSION}/{settings.ENV}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}" + return f"s3://{self.bucket}/{self._PATH_VERSION}/{settings.ENV}/{organization_id}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}" async def retrieve_global_workflows(self) -> list[str]: uri = f"s3://{self.bucket}/{settings.ENV}/global_workflows.txt" diff --git a/skyvern/forge/sdk/artifact/storage/test_local_storage.py b/skyvern/forge/sdk/artifact/storage/test_local_storage.py index ef31301d..d4cd62de 100644 --- a/skyvern/forge/sdk/artifact/storage/test_local_storage.py +++ b/skyvern/forge/sdk/artifact/storage/test_local_storage.py @@ -31,10 +31,15 @@ def local_storage() -> LocalStorage: class TestLocalStorageBuildURIs: def test_build_uri(self, local_storage: LocalStorage) -> None: step = create_fake_step(TEST_STEP_ID) - uri = local_storage.build_uri("artifact123", step, ArtifactType.LLM_PROMPT) + uri = local_storage.build_uri( + organization_id=TEST_ORGANIZATION_ID, + artifact_id="artifact123", + step=step, + artifact_type=ArtifactType.LLM_PROMPT, + ) assert ( uri - == f"file://{local_storage.artifact_path}/{TEST_TASK_ID}/01_0_{TEST_STEP_ID}/2025-06-09T12:00:00_artifact123_llm_prompt.txt" + == f"file://{local_storage.artifact_path}/{TEST_ORGANIZATION_ID}/{TEST_TASK_ID}/01_0_{TEST_STEP_ID}/2025-06-09T12:00:00_artifact123_llm_prompt.txt" ) def test_build_log_uri(self, local_storage: LocalStorage) -> None: diff --git a/skyvern/forge/sdk/artifact/storage/test_s3_storage.py b/skyvern/forge/sdk/artifact/storage/test_s3_storage.py index cd7a80e8..b15e9363 100644 --- a/skyvern/forge/sdk/artifact/storage/test_s3_storage.py +++ b/skyvern/forge/sdk/artifact/storage/test_s3_storage.py @@ -31,10 +31,15 @@ def s3_storage() -> S3Storage: class TestS3StorageBuildURIs: def test_build_uri(self, s3_storage: S3Storage) -> None: step = create_fake_step(TEST_STEP_ID) - uri = s3_storage.build_uri("artifact123", step, ArtifactType.LLM_PROMPT) + uri = s3_storage.build_uri( + organization_id=TEST_ORGANIZATION_ID, + artifact_id="artifact123", + step=step, + artifact_type=ArtifactType.LLM_PROMPT, + ) assert ( uri - == f"s3://{TEST_BUCKET}/v1/{settings.ENV}/{TEST_TASK_ID}/01_0_{TEST_STEP_ID}/2025-06-09T12:00:00_artifact123_llm_prompt.txt" + == f"s3://{TEST_BUCKET}/v1/{settings.ENV}/{TEST_ORGANIZATION_ID}/{TEST_TASK_ID}/01_0_{TEST_STEP_ID}/2025-06-09T12:00:00_artifact123_llm_prompt.txt" ) def test_build_log_uri(self, s3_storage: S3Storage) -> None: