Missed a spot - add org id to build_uri (#2674)

This commit is contained in:
Asher Foa
2025-06-10 16:47:53 -04:00
committed by GitHub
parent 7b27ac9bdf
commit f05170df61
6 changed files with 22 additions and 10 deletions

View File

@@ -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,

View File

@@ -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

View File

@@ -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")

View File

@@ -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"

View File

@@ -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:

View File

@@ -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: