refactor: remove task metadata property (#216)

* remove metadata

* docs and tests

* don't need that test

* use static method for generation strategy

* update docs

* add docstring

Author: johnnygreco

docs/plugins/example.md

@@ -64,24 +64,12 @@ The implementation class defines the actual business logic of the plugin. For co
 import logging
 import pandas as pd
 
-from data_designer.engine.column_generators.generators.base import (
-    ColumnGenerator,
-    GenerationStrategy,
-    GeneratorMetadata,
-)
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorFullColumn, GenerationStrategy
 
 # Data Designer uses the standard Python logging module for logging
 logger = logging.getLogger(__name__)
 
-class IndexMultiplierColumnGenerator(ColumnGenerator[IndexMultiplierColumnConfig]):
-    @staticmethod
-    def metadata() -> GeneratorMetadata:
-        """Define metadata about this generator."""
-        return GeneratorMetadata(
-            name="index-multiplier",
-            description="Generates values by multiplying the row index by a user-specified multiplier",
-            generation_strategy=GenerationStrategy.FULL_COLUMN,
-        )
+class IndexMultiplierColumnGenerator(ColumnGeneratorFullColumn[IndexMultiplierColumnConfig]):
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:
         """Generate the column data.
@@ -105,20 +93,20 @@ class IndexMultiplierColumnGenerator(ColumnGenerator[IndexMultiplierColumnConfig
 
 **Key points:**
 
-- Generic type `ColumnGenerator[IndexMultiplierColumnConfig]` connects the task to its config
-- `metadata()` describes your generator and its requirements
-- `generation_strategy` can be `FULL_COLUMN`, `CELL_BY_CELL`
+- Generic type `ColumnGeneratorFullColumn[IndexMultiplierColumnConfig]` connects the task to its config
 - You have access to the configuration parameters via `self.config`
 
 !!! info "Understanding generation_strategy"
     The `generation_strategy` specifies how the column generator will generate data.
 
     - **`FULL_COLUMN`**: Generates the full column (at the batch level) in a single call to `generate`
-        - `generate` must take as input a `pd.DataFrame` with all previous columns and return a `pd.DataFrame` with the generated column appended
+        - `generate` must take as input a `pd.DataFrame` with all previous columns and return a `pd.DataFrame` with the generated column appended.
+        - Inherit from `ColumnGeneratorFullColumn` for this strategy, as we do in the example above.
 
     - **`CELL_BY_CELL`**: Generates one cell at a time
         - `generate` must take as input a `dict` with key/value pairs for all previous columns and return a `dict` with an additional key/value for the generated cell
         - Supports concurrent workers via a `max_parallel_requests` parameter on the configuration
+        - Inherit from `ColumnGeneratorCellByCell` for this strategy.
 
 ## Step 4: Create the plugin object
 
@@ -147,11 +135,8 @@ from typing import Literal
 import pandas as pd
 
 from data_designer.config.column_configs import SingleColumnConfig
-from data_designer.engine.column_generators.generators.base import (
-    ColumnGenerator,
-    GenerationStrategy,
-    GeneratorMetadata,
-)
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorFullColumn
+
 from data_designer.plugins import Plugin, PluginType
 
 # Data Designer uses the standard Python logging module for logging
@@ -169,15 +154,7 @@ class IndexMultiplierColumnConfig(SingleColumnConfig):
     column_type: Literal["index-multiplier"] = "index-multiplier"
 
 
-class IndexMultiplierColumnGenerator(ColumnGenerator[IndexMultiplierColumnConfig]):
-    @staticmethod
-    def metadata() -> GeneratorMetadata:
-        """Define metadata about this generator."""
-        return GeneratorMetadata(
-            name="index-multiplier",
-            description="Generates values by multiplying the row index by a user-specified multiplier",
-            generation_strategy=GenerationStrategy.FULL_COLUMN,
-        )
+class IndexMultiplierColumnGenerator(ColumnGeneratorFullColumn[IndexMultiplierColumnConfig]):
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:
         """Generate the column data.

src/data_designer/engine/analysis/column_profilers/base.py

@@ -13,7 +13,7 @@
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import SingleColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType
-from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata, TaskConfigT
+from data_designer.engine.configurable_task import ConfigurableTask, TaskConfigT
 
 logger = logging.getLogger(__name__)
 
@@ -32,17 +32,14 @@ def as_tuple(self) -> tuple[SingleColumnConfig, pd.DataFrame]:
         return (self.column_config, self.df)
 
 
-class ColumnProfilerMetadata(ConfigurableTaskMetadata):
-    applicable_column_types: list[DataDesignerColumnType]
-
-
 class ColumnProfiler(ConfigurableTask[TaskConfigT], ABC):
     @staticmethod
     @abstractmethod
-    def metadata() -> ColumnProfilerMetadata: ...
+    def get_applicable_column_types() -> list[DataDesignerColumnType]:
+        """Returns a list of column types that this profiler can be applied to during dataset profiling."""
 
     @abstractmethod
     def profile(self, column_config_with_df: ColumnConfigWithDataFrame) -> BaseModel: ...
 
     def _initialize(self) -> None:
-        logger.info(f"💫 Initializing column profiler: '{self.metadata().name}'")
+        logger.info(f"💫 Initializing column profiler: '{self.name}'")

src/data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -20,11 +20,7 @@
     NumericalDistribution,
 )
 from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP, DataDesignerColumnType
-from data_designer.engine.analysis.column_profilers.base import (
-    ColumnConfigWithDataFrame,
-    ColumnProfiler,
-    ColumnProfilerMetadata,
-)
+from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
 from data_designer.engine.analysis.utils.judge_score_processing import (
     extract_judge_score_distributions,
     sample_scores_and_reasoning,
@@ -37,12 +33,8 @@
 
 class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
     @staticmethod
-    def metadata() -> ColumnProfilerMetadata:
-        return ColumnProfilerMetadata(
-            name="judge_score_profiler",
-            description="Analyzes LLM-as-judge score distributions in a Data Designer dataset.",
-            applicable_column_types=[DataDesignerColumnType.LLM_JUDGE],
-        )
+    def get_applicable_column_types() -> list[DataDesignerColumnType]:
+        return [DataDesignerColumnType.LLM_JUDGE]
 
     def get_model(self, model_alias: str) -> ModelFacade:
         return self.resource_provider.model_registry.get_model(model_alias=model_alias)

src/data_designer/engine/analysis/dataset_profiler.py

@@ -13,10 +13,7 @@
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import SingleColumnConfig
-from data_designer.config.column_types import (
-    COLUMN_TYPE_EMOJI_MAP,
-    ColumnConfigT,
-)
+from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP, ColumnConfigT
 from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
 from data_designer.engine.analysis.column_statistics import get_column_statistics_calculator
 from data_designer.engine.analysis.errors import DatasetProfilerConfigurationError
@@ -81,14 +78,14 @@ def profile_dataset(
         column_profiles = []
         for profiler_config in self.config.column_profiler_configs or []:
             profiler = self._create_column_profiler(profiler_config)
-            applicable_column_types = profiler.metadata().applicable_column_types
+            applicable_column_types = profiler.get_applicable_column_types()
             for c in self.config.column_configs:
                 if c.column_type in applicable_column_types:
                     params = ColumnConfigWithDataFrame(column_config=c, df=dataset)
                     column_profiles.append(profiler.profile(params))
             if len(column_profiles) == 0:
                 logger.warning(
-                    f"⚠️ No applicable column types found for the '{profiler.metadata().name}' profiler. "
+                    f"⚠️ No applicable column types found for the '{profiler.name}' profiler. "
                     f"This profiler is applicable to the following column types: {applicable_column_types}"
                 )
 

src/data_designer/engine/column_generators/generators/base.py

@@ -11,7 +11,7 @@
 
 import pandas as pd
 
-from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata, DataT, TaskConfigT
+from data_designer.engine.configurable_task import ConfigurableTask, DataT, TaskConfigT
 
 if TYPE_CHECKING:
     from data_designer.config.models import BaseInferenceParams, ModelConfig
@@ -27,22 +27,14 @@ class GenerationStrategy(str, Enum):
     FULL_COLUMN = "full_column"
 
 
-class GeneratorMetadata(ConfigurableTaskMetadata):
-    generation_strategy: GenerationStrategy
-
-
 class ColumnGenerator(ConfigurableTask[TaskConfigT], ABC):
     @property
     def can_generate_from_scratch(self) -> bool:
         return False
 
-    @property
-    def generation_strategy(self) -> GenerationStrategy:
-        return self.metadata().generation_strategy
-
     @staticmethod
     @abstractmethod
-    def metadata() -> GeneratorMetadata: ...
+    def get_generation_strategy() -> GenerationStrategy: ...
 
     @overload
     @abstractmethod
@@ -108,3 +100,21 @@ def log_pre_generation(self) -> None:
         logger.info(f"  |-- model alias: {self.config.model_alias!r}")
         logger.info(f"  |-- model provider: {self.get_model_provider_name(model_alias=self.config.model_alias)!r}")
         logger.info(f"  |-- inference parameters: {self.inference_parameters.format_for_display()}")
+
+
+class ColumnGeneratorCellByCell(ColumnGenerator[TaskConfigT], ABC):
+    @staticmethod
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.CELL_BY_CELL
+
+    @abstractmethod
+    def generate(self, data: dict) -> dict: ...
+
+
+class ColumnGeneratorFullColumn(ColumnGenerator[TaskConfigT], ABC):
+    @staticmethod
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.FULL_COLUMN
+
+    @abstractmethod
+    def generate(self, data: pd.DataFrame) -> pd.DataFrame: ...

src/data_designer/engine/column_generators/generators/embedding.py

@@ -5,11 +5,7 @@
 from pydantic import BaseModel, computed_field
 
 from data_designer.config.column_configs import EmbeddingColumnConfig
-from data_designer.engine.column_generators.generators.base import (
-    ColumnGeneratorWithModel,
-    GenerationStrategy,
-    GeneratorMetadata,
-)
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModel, GenerationStrategy
 from data_designer.engine.processing.utils import deserialize_json_values, parse_list_string
 
 
@@ -27,12 +23,8 @@ def dimension(self) -> int:
 
 class EmbeddingCellGenerator(ColumnGeneratorWithModel[EmbeddingColumnConfig]):
     @staticmethod
-    def metadata() -> GeneratorMetadata:
-        return GeneratorMetadata(
-            name="embedding_cell_generator",
-            description="Generate embeddings for a text column.",
-            generation_strategy=GenerationStrategy.CELL_BY_CELL,
-        )
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.CELL_BY_CELL
 
     def generate(self, data: dict) -> dict:
         deserialized_record = deserialize_json_values(data)

src/data_designer/engine/column_generators/generators/expression.py

@@ -8,27 +8,15 @@
 import pandas as pd
 
 from data_designer.config.column_configs import ExpressionColumnConfig
-from data_designer.engine.column_generators.generators.base import (
-    ColumnGenerator,
-    GenerationStrategy,
-    GeneratorMetadata,
-)
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorFullColumn
 from data_designer.engine.column_generators.utils.errors import ExpressionTemplateRenderError
 from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
 from data_designer.engine.processing.utils import deserialize_json_values
 
 logger = logging.getLogger(__name__)
 
 
-class ExpressionColumnGenerator(WithJinja2UserTemplateRendering, ColumnGenerator[ExpressionColumnConfig]):
-    @staticmethod
-    def metadata() -> GeneratorMetadata:
-        return GeneratorMetadata(
-            name="expression_generator",
-            description="Generate a column from a jinja2 expression.",
-            generation_strategy=GenerationStrategy.FULL_COLUMN,
-        )
-
+class ExpressionColumnGenerator(WithJinja2UserTemplateRendering, ColumnGeneratorFullColumn[ExpressionColumnConfig]):
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:
         logger.info(f"🧩 Generating column `{self.config.name}` from expression")
 

src/data_designer/engine/column_generators/generators/llm_completion.py

@@ -11,11 +11,7 @@
     LLMTextColumnConfig,
 )
 from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
-from data_designer.engine.column_generators.generators.base import (
-    ColumnGeneratorWithModel,
-    GenerationStrategy,
-    GeneratorMetadata,
-)
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModel, GenerationStrategy
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
     RecordBasedPromptRenderer,
@@ -29,6 +25,10 @@
 
 
 class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfigT]):
+    @staticmethod
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.CELL_BY_CELL
+
     @functools.cached_property
     def response_recipe(self) -> ResponseRecipe:
         return create_response_recipe(self.config, self.model_config)
@@ -87,41 +87,13 @@ def generate(self, data: dict) -> dict:
         return data
 
 
-class LLMTextCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMTextColumnConfig]):
-    @staticmethod
-    def metadata() -> GeneratorMetadata:
-        return GeneratorMetadata(
-            name="llm_text_generator",
-            description="Generate a new dataset cell from a prompt template",
-            generation_strategy=GenerationStrategy.CELL_BY_CELL,
-        )
+class LLMTextCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMTextColumnConfig]): ...
 
 
-class LLMCodeCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMCodeColumnConfig]):
-    @staticmethod
-    def metadata() -> GeneratorMetadata:
-        return GeneratorMetadata(
-            name="llm_code_generator",
-            description="Generate a new dataset cell from a prompt template",
-            generation_strategy=GenerationStrategy.CELL_BY_CELL,
-        )
+class LLMCodeCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMCodeColumnConfig]): ...
 
 
-class LLMStructuredCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMStructuredColumnConfig]):
-    @staticmethod
-    def metadata() -> GeneratorMetadata:
-        return GeneratorMetadata(
-            name="llm_structured_generator",
-            description="Generate a new dataset cell from a prompt template",
-            generation_strategy=GenerationStrategy.CELL_BY_CELL,
-        )
+class LLMStructuredCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMStructuredColumnConfig]): ...
 
 
-class LLMJudgeCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMJudgeColumnConfig]):
-    @staticmethod
-    def metadata() -> GeneratorMetadata:
-        return GeneratorMetadata(
-            name="llm_judge_generator",
-            description="Judge a new dataset cell based on a set of rubrics",
-            generation_strategy=GenerationStrategy.CELL_BY_CELL,
-        )
+class LLMJudgeCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMJudgeColumnConfig]): ...

src/data_designer/engine/column_generators/generators/samplers.py

@@ -11,11 +11,7 @@
 import pandas as pd
 
 from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS
-from data_designer.engine.column_generators.generators.base import (
-    FromScratchColumnGenerator,
-    GenerationStrategy,
-    GeneratorMetadata,
-)
+from data_designer.engine.column_generators.generators.base import FromScratchColumnGenerator, GenerationStrategy
 from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig
 from data_designer.engine.processing.utils import concat_datasets
 from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator
@@ -28,12 +24,8 @@
 
 class SamplerColumnGenerator(FromScratchColumnGenerator[SamplerMultiColumnConfig]):
     @staticmethod
-    def metadata() -> GeneratorMetadata:
-        return GeneratorMetadata(
-            name="sampler_column_generator",
-            description="Generate columns using sampling-based method.",
-            generation_strategy=GenerationStrategy.FULL_COLUMN,
-        )
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.FULL_COLUMN
 
     def generate(self, data: pd.DataFrame) -> pd.DataFrame:
         df_samplers = self.generate_from_scratch(len(data))