Missed a spot - add org id to build_uri (#2674)
This commit is contained in:
@@ -72,7 +72,9 @@ class ArtifactManager:
|
|||||||
path: str | None = None,
|
path: str | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
artifact_id = generate_artifact_id()
|
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(
|
return await self._create_artifact(
|
||||||
aio_task_primary_key=step.task_id,
|
aio_task_primary_key=step.task_id,
|
||||||
artifact_id=artifact_id,
|
artifact_id=artifact_id,
|
||||||
|
|||||||
@@ -39,7 +39,7 @@ FILE_EXTENTSION_MAP: dict[ArtifactType, str] = {
|
|||||||
|
|
||||||
class BaseStorage(ABC):
|
class BaseStorage(ABC):
|
||||||
@abstractmethod
|
@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
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
|
|||||||
@@ -28,9 +28,9 @@ class LocalStorage(BaseStorage):
|
|||||||
def __init__(self, artifact_path: str = settings.ARTIFACT_STORAGE_PATH) -> None:
|
def __init__(self, artifact_path: str = settings.ARTIFACT_STORAGE_PATH) -> None:
|
||||||
self.artifact_path = artifact_path
|
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]
|
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]:
|
async def retrieve_global_workflows(self) -> list[str]:
|
||||||
file_path = Path(f"{self.artifact_path}/{settings.ENV}/global_workflows.txt")
|
file_path = Path(f"{self.artifact_path}/{settings.ENV}/global_workflows.txt")
|
||||||
|
|||||||
@@ -35,9 +35,9 @@ class S3Storage(BaseStorage):
|
|||||||
self.async_client = AsyncAWSClient()
|
self.async_client = AsyncAWSClient()
|
||||||
self.bucket = bucket or settings.AWS_S3_BUCKET_ARTIFACTS
|
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]
|
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]:
|
async def retrieve_global_workflows(self) -> list[str]:
|
||||||
uri = f"s3://{self.bucket}/{settings.ENV}/global_workflows.txt"
|
uri = f"s3://{self.bucket}/{settings.ENV}/global_workflows.txt"
|
||||||
|
|||||||
@@ -31,10 +31,15 @@ def local_storage() -> LocalStorage:
|
|||||||
class TestLocalStorageBuildURIs:
|
class TestLocalStorageBuildURIs:
|
||||||
def test_build_uri(self, local_storage: LocalStorage) -> None:
|
def test_build_uri(self, local_storage: LocalStorage) -> None:
|
||||||
step = create_fake_step(TEST_STEP_ID)
|
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 (
|
assert (
|
||||||
uri
|
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:
|
def test_build_log_uri(self, local_storage: LocalStorage) -> None:
|
||||||
|
|||||||
@@ -31,10 +31,15 @@ def s3_storage() -> S3Storage:
|
|||||||
class TestS3StorageBuildURIs:
|
class TestS3StorageBuildURIs:
|
||||||
def test_build_uri(self, s3_storage: S3Storage) -> None:
|
def test_build_uri(self, s3_storage: S3Storage) -> None:
|
||||||
step = create_fake_step(TEST_STEP_ID)
|
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 (
|
assert (
|
||||||
uri
|
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:
|
def test_build_log_uri(self, s3_storage: S3Storage) -> None:
|
||||||
|
|||||||
Reference in New Issue
Block a user