chore: move ArtifactStorage to engine/storage/ module by nabinchha · Pull Request #321 · NVIDIA-NeMo/DataDesigner

nabinchha · 2026-02-10T21:55:56Z

📋 Summary

Consolidate all storage logic under engine/storage/ by moving ArtifactStorage alongside MediaStorage for better code organization.

🔄 Changes

🔧 Changed

Moved artifact_storage.py from engine/dataset_builders/ to engine/storage/
Moved test_artifact_storage.py from tests/engine/dataset_builders/ to tests/engine/storage/
Updated engine/storage/__init__.py to export ArtifactStorage, BatchStage, and related constants
Updated imports across 18 source and test files

#317 needs to merge first

…eInferenceParameters, EmbeddingInferenceParameters

…th BaseInferenceParameters

…lved based on the type of InferenceParameters

…eneration

Consolidate all storage logic (ArtifactStorage, MediaStorage) under engine/storage/ for better organization. Update all imports across source and test files.

greptile-apps · 2026-02-10T21:59:49Z

Greptile Overview

Greptile Summary

This PR consolidates storage-related code by moving ArtifactStorage from engine/dataset_builders/ to engine/storage/ alongside MediaStorage. All 18 import statements across source and test files have been updated to reference the new location.

Key changes:

Moved artifact_storage.py to engine/storage/ module with no content changes
Moved test_artifact_storage.py to tests/engine/storage/ with updated mock patch path
Updated imports in 18 files from dataset_builders.artifact_storage to storage.artifact_storage
Removed duplicate _has_image_columns method in column_wise_builder.py

Note: The PR description mentions updating engine/storage/__init__.py to export ArtifactStorage, BatchStage, and related constants, but this file was not modified in the changeset. The __init__.py currently only exports MediaStorage and StorageMode.

Confidence Score: 5/5

This PR is safe to merge with minimal risk
This is a straightforward refactoring that moves files and updates import paths without changing any logic. All imports have been systematically updated, including TYPE_CHECKING blocks and mock patch paths. The removal of the duplicate method is a minor cleanup improvement.
No files require special attention

Important Files Changed

Filename	Overview
packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py	File moved from `dataset_builders/` to `storage/` without content changes
packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py	Test file moved to match new location, imports updated correctly, patch path updated
packages/data-designer-engine/src/data_designer/engine/dataset_builders/column_wise_builder.py	Import updated and duplicate `_has_image_columns` method removed
packages/data-designer/src/data_designer/integrations/huggingface/client.py	Import path updated to new storage location for constants
packages/data-designer/src/data_designer/interface/data_designer.py	Import path updated to new storage location

_{Last reviewed commit: f0ea307}

- Reduce num_records to 2 for image generation in tutorial notebook - Add tests for different image response formats (dict and plain string) - Parametrize PNG/JPG media storage tests for better maintainability

…pati/chore/move-artifact-storage

greptile-apps

_{20 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-11T18:31:15Z

+    "MediaStorage",
+    "StorageMode",
+]


Missing trailing newline

packages/data-designer-engine/src/data_designer/engine/storage/__init__.py appears to be missing a final newline at EOF. Some tooling (e.g., strict linters / formatting checks) treats this as a failure. Adding a newline after the closing ] avoids CI/style check issues.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

Prompt To Fix With AI

This is a comment left during a code review. Path: packages/data-designer-engine/src/data_designer/engine/storage/__init__.py Line: 23:25 Comment: **Missing trailing newline** `packages/data-designer-engine/src/data_designer/engine/storage/__init__.py` appears to be missing a final newline at EOF. Some tooling (e.g., strict linters / formatting checks) treats this as a failure. Adding a newline after the closing `]` avoids CI/style check issues. <sub>Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!</sub> How can I resolve this? If you propose a fix, please make it concise.

…pati/chore/move-artifact-storage

The base branch was changed.

greptile-apps

_{20 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-12T21:53:28Z

Additional Comments (1)

packages/data-designer-engine/src/data_designer/engine/storage/__init__.py
PR description states this file should export ArtifactStorage, BatchStage, and related constants, but they're not in __all__. Consider adding them for consistency with the PR description.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

Prompt To Fix With AI

This is a comment left during a code review.
Path: packages/data-designer-engine/src/data_designer/engine/storage/__init__.py
Line: 6:6

Comment:
PR description states this file should export `ArtifactStorage`, `BatchStage`, and related constants, but they're not in `__all__`. Consider adding them for consistency with the PR description.

<sub>Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!</sub>

How can I resolve this? If you propose a fix, please make it concise.

nabinchha added 30 commits November 25, 2025 12:16

Add generation type to ModelConfig

dc041f7

pass tests

0d6b830

added generate_text_embeddings

254fd8a

tests

1126ea1

remove sensitive=True old artifact no longer needed

744bc8f

Slight refactor

b913f8d

slight refactor

052db7a

Added embedding generator

5504c8d

chunk_separator -> chunk_pattern

4b6f877

update tests

04fc0f3

rename for consistency

26d6da1

Restructure InferenceParameters -> CompletionInferenceParameters, Bas…

6facbd2

…eInferenceParameters, EmbeddingInferenceParameters

Remove purpose from consolidated kwargs

2c1b267

WithModelConfiguration.inference_parameters should should be typed wi…

4b1492b

…th BaseInferenceParameters

Type as WithModelGeneration

c445caf

Add image generation modality

4b8aa2b

update return type for generate_kwargs

2c5933f

make generation_type a field of ModelConfig as opposed to a prop reso…

c6c29d4

…lved based on the type of InferenceParameters

remove regex based chunking from embedding generator

06a724b

Merge branch 'main' into nmulepati/feat/support-embedding-and-image-g…

6b9733f

…eneration

Merge branch 'main' into nmulepati/feat/support-embedding-and-image-g…

81949e6

…eneration

save progress

f291033

Merge branch 'main' into nmulepati/feat/125-support-image-generation

e0a4657

Simplify to ImageInferenceParams. Persist images in create mode to disk

1506ab5

support generation of multiple images

ed9787b

clean up visualization

7dea87a

clean up some util methods + add tests

31cc24e

Streamline integration for image generation

0f07f7b

streamline generation

2aae6cc

track images generated in usage

1677f06

Move ArtifactStorage to engine/storage/ module

a5fab8a

Consolidate all storage logic (ArtifactStorage, MediaStorage) under engine/storage/ for better organization. Update all imports across source and test files.

nabinchha requested a review from a team as a code owner February 10, 2026 21:55

nabinchha mentioned this pull request Feb 10, 2026

feat: add image generation support with multi-modal context #317

Merged

nabinchha added 9 commits February 10, 2026 15:12

Address PR review comments

cf2b364

- Reduce num_records to 2 for image generation in tutorial notebook - Add tests for different image response formats (dict and plain string) - Parametrize PNG/JPG media storage tests for better maintainability

Use regex for base64 character validation in is_base64_image

5aa7e10

move to a constant

ecaeb72

fix pyproject.toml

622b1c4

regen colab notebooks

400e97b

raise a ValueError if we fail to detect image format

469a3d2

Fix diffusion image gen

1e43394

Add requests to config pyproject.toml

8f6be9b

Merge branch 'nmulepati/feat/125-support-image-generation' into nmule…

a277dc1

…pati/chore/move-artifact-storage

greptile-apps Bot reviewed Feb 11, 2026

View reviewed changes

nabinchha added 5 commits February 11, 2026 15:15

Merge branch 'main' into nmulepati/feat/125-support-image-generation

d85ccb3

address pr feedback from andre

87dcab1

Merge branch 'main' into nmulepati/feat/125-support-image-generation

12cc1fe

Merge branch 'main' into nmulepati/feat/125-support-image-generation

a0ea92b

Merge branch 'nmulepati/feat/125-support-image-generation' into nmule…

7590778

…pati/chore/move-artifact-storage

johnnygreco previously approved these changes Feb 12, 2026

View reviewed changes

Base automatically changed from nmulepati/feat/125-support-image-generation to main February 12, 2026 21:00

nabinchha requested a review from johnnygreco February 12, 2026 21:01

Merge branch 'main' into nmulepati/chore/move-artifact-storage

f0ea307

greptile-apps Bot reviewed Feb 12, 2026

View reviewed changes

andreatgretel approved these changes Feb 12, 2026

View reviewed changes

johnnygreco approved these changes Feb 12, 2026

View reviewed changes

nabinchha merged commit 1f172ea into main Feb 12, 2026
47 checks passed

nabinchha deleted the nmulepati/chore/move-artifact-storage branch February 12, 2026 21:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

chore: move ArtifactStorage to engine/storage/ module#321

chore: move ArtifactStorage to engine/storage/ module#321
nabinchha merged 73 commits into
mainfrom
nmulepati/chore/move-artifact-storage

nabinchha commented Feb 10, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Feb 10, 2026 •

edited

Loading

Confidence Score: 5/5

Uh oh!

greptile-apps Bot left a comment

Uh oh!

greptile-apps Bot Feb 11, 2026

Uh oh!

greptile-apps Bot left a comment

Uh oh!

greptile-apps Bot commented Feb 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

nabinchha commented Feb 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📋 Summary

🔄 Changes

🔧 Changed

Uh oh!

greptile-apps Bot commented Feb 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Overview

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot commented Feb 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nabinchha commented Feb 10, 2026 •

edited

Loading

greptile-apps Bot commented Feb 10, 2026 •

edited

Loading