DataDesigner/docs/plugins/models.md

Using Models in Plugins

Model access belongs in column generator implementations, not config objects. Keep the config declarative by asking users for model aliases, then resolve those aliases at runtime through the model registry.

Do not construct model clients in plugin configs, read API keys in configs, or bypass Data Designer's model providers. The engine builds a ResourceProvider and exposes its model registry to every generator at:

self.resource_provider.model_registry

Access the registry

Use a model-aware column generator base whenever your plugin needs the registry:

Need	Base class	Registry access
Primary model alias	ColumnGeneratorWithModel	Use self.model, self.model_config, and self.inference_parameters.
Multiple aliases or provider inspection	ColumnGeneratorWithModelRegistry	Use self.get_model(alias), self.get_model_config(alias), and self.get_model_provider_name(alias).

ColumnGeneratorWithModel is a convenience subclass of ColumnGeneratorWithModelRegistry. It expects the config to have a model_alias field and resolves that one alias for you. For independent model calls, return GenerationStrategy.CELL_BY_CELL so the runtime can fan out rows like the built-in LLM, embedding, and image generators. Use full-column generation only when your plugin intentionally calls a batched API for the whole DataFrame.

from __future__ import annotations

from data_designer.config.column_configs import GenerationStrategy
from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModel
from data_designer.engine.models.parsers.errors import ParserException

from data_designer_sentiment_label.config import SentimentLabelColumnConfig


def parse_sentiment_label(response: str) -> str:
    label = response.strip().lower()
    if label not in {"positive", "neutral", "negative"}:
        raise ParserException("Expected exactly one of: positive, neutral, negative.", source=response)
    return label


class SentimentLabelColumnGenerator(ColumnGeneratorWithModel[SentimentLabelColumnConfig]):
    @staticmethod
    def get_generation_strategy() -> GenerationStrategy:
        return GenerationStrategy.CELL_BY_CELL

    async def agenerate(self, data: dict) -> dict:
        label, _ = await self.model.agenerate(
            prompt=f"Classify the sentiment of this text: {data[self.config.source_column]}",
            system_prompt="Return exactly one label: positive, neutral, or negative.",
            parser=parse_sentiment_label,
            max_correction_steps=self.resource_provider.run_config.max_conversation_correction_steps,
            max_conversation_restarts=self.resource_provider.run_config.max_conversation_restarts,
            purpose=f"running generation for column '{self.config.name}'",
        )
        data[self.config.name] = label
        return data

The matching config must include model_alias: str as a normal user-facing field:

from __future__ import annotations

from typing import Literal

from data_designer.config.base import SingleColumnConfig


class SentimentLabelColumnConfig(SingleColumnConfig):
    column_type: Literal["sentiment-label"] = "sentiment-label"
    source_column: str
    model_alias: str

    @property
    def required_columns(self) -> list[str]:
        return [self.source_column]

    @property
    def side_effect_columns(self) -> list[str]:
        return []

Users set that alias from default model settings or from DataDesignerConfigBuilder(model_configs=...).

Use multiple models

If your plugin uses multiple model aliases, inherit from ColumnGeneratorWithModelRegistry and resolve each alias explicitly with self.get_model(...).

The config must include a primary model_alias: str field. Startup health checks read it directly from any column config whose generator inherits from ColumnGeneratorWithModelRegistry, including generators that inherit through ColumnGeneratorWithModel. A config for this pattern might also define judge_model_alias, critic_model_alias, or another task-specific alias.

Validate additional alias fields in _validate() or _initialize() with get_model_config(...) so missing aliases fail before generation starts. get_model_config(alias) only verifies that the alias is registered; it does not call the endpoint. Endpoint reachability is only exercised for the primary model_alias collected by the standard startup health check.

The matching config shows which alias gets the standard startup health check and which alias the plugin validates itself:

from __future__ import annotations

from typing import Literal

from data_designer.config.base import SingleColumnConfig


class PairwiseJudgeColumnConfig(SingleColumnConfig):
    column_type: Literal["pairwise-judge"] = "pairwise-judge"
    question_column: str
    model_alias: str
    judge_model_alias: str

    @property
    def required_columns(self) -> list[str]:
        return [self.question_column]

    @property
    def side_effect_columns(self) -> list[str]:
        return []

from __future__ import annotations

from data_designer.config.column_configs import GenerationStrategy
from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModelRegistry
from data_designer.engine.models.parsers.errors import ParserException

from data_designer_pairwise_judge.config import PairwiseJudgeColumnConfig


def parse_score(response: str) -> int:
    text = response.strip()
    if text not in {"1", "2", "3", "4", "5"}:
        raise ParserException("Expected an integer score from 1 to 5.", source=response)
    return int(text)


class PairwiseJudgeColumnGenerator(ColumnGeneratorWithModelRegistry[PairwiseJudgeColumnConfig]):
    @staticmethod
    def get_generation_strategy() -> GenerationStrategy:
        return GenerationStrategy.CELL_BY_CELL

    def _validate(self) -> None:
        self.get_model_config(self.config.model_alias)
        self.get_model_config(self.config.judge_model_alias)

    async def agenerate(self, data: dict) -> dict:
        generator_model = self.get_model(self.config.model_alias)
        judge_model = self.get_model(self.config.judge_model_alias)
        retry_kwargs = {
            "max_correction_steps": self.resource_provider.run_config.max_conversation_correction_steps,
            "max_conversation_restarts": self.resource_provider.run_config.max_conversation_restarts,
        }

        draft, _ = await generator_model.agenerate(
            prompt=f"Draft an answer for: {data[self.config.question_column]}",
            purpose=f"drafting an answer for column '{self.config.name}'",
            **retry_kwargs,
        )
        score, _ = await judge_model.agenerate(
            prompt=f"Score this answer from 1 to 5: {draft}",
            system_prompt="Return exactly one integer from 1 to 5.",
            parser=parse_score,
            purpose=f"judging an answer for column '{self.config.name}'",
            **retry_kwargs,
        )
        data[self.config.name] = {"draft": draft, "score": score}
        return data

What the registry returns

get_model(...) returns a ModelFacade. Call the facade based on the modality your plugin needs:

• Chat completion aliases use model.generate(...) or await model.agenerate(...) and return (parsed_output, trace).
• Embedding aliases use model.generate_text_embeddings(...) or await model.agenerate_text_embeddings(...) and return list[list[float]].
• Image aliases use model.generate_image(...) or await model.agenerate_image(...) and return list[str] of base64-encoded image data.

Choose a model alias whose ModelConfig.inference_parameters.generation_type matches the facade method you call. The facade merges the alias's configured inference parameters into each request.

Pass runtime context such as prompt, system_prompt, parser, tool_alias, multi_modal_context, max_correction_steps, max_conversation_restarts, and purpose at the call site. Parser functions should raise ParserException for invalid model responses; that is what allows ModelFacade.generate(...) and ModelFacade.agenerate(...) to run correction turns and conversation restarts.

Prefer implementing agenerate(...) for model-backed plugins. The base generate(...) method can bridge to agenerate(...) for sync runs when the subclass only implements async generation. If your plugin has a sync-specific path, implement both generate(...) and agenerate(...), as the built-in generators do.

Health checks and scheduling

The model-aware bases mark the generator as LLM-bound, so the async scheduler treats the work like other model calls.

Plugin discovery treats column generator implementations that inherit from ColumnGeneratorWithModelRegistry as model-generated column types for startup model health checks. The standard health-check collection reads a primary model_alias field directly from the config. Additional alias fields should be registration-validated by the plugin implementation; their endpoints are not pinged by the standard startup health check.

Built-in patterns

The built-in model-backed generators use these same hooks:

• LLMTextCellGenerator, LLMCodeCellGenerator, LLMStructuredCellGenerator, and LLMJudgeCellGenerator inherit through a chat-completion base that uses ColumnGeneratorWithModel. They render prompts from row data, call self.model.generate(...) or self.model.agenerate(...), pass parsers into the ModelFacade, and store optional trace side-effect columns.
• EmbeddingCellGenerator uses ColumnGeneratorWithModel but calls the facade's embedding methods instead of chat completion.
• ImageCellGenerator uses ColumnGeneratorWithModel, renders a prompt, calls the facade's image methods, and writes generated media through the artifact storage supplied by the same ResourceProvider.
• CustomColumnGenerator is the inline-function counterpart: when users declare model_aliases, it builds a models dict from resource_provider.model_registry. Packaged plugins usually use ColumnGeneratorWithModel or ColumnGeneratorWithModelRegistry directly instead of recreating that dict.

See Column Generators for the full base-class API and Custom Model Settings for configuring model aliases.