Custom Columns

1
import data_designer.config as dd
2
3
@dd.custom_column_generator(required_columns=["name"])
4
def create_greeting(row: dict) -> dict:
5
    row["greeting"] = f"Hello, {row['name']}!"
6
    return row
7
8
config_builder.add_column(
9
    dd.CustomColumnConfig(
10
        name="greeting",
11
        generator_function=create_greeting,
12
    )
13
)

1
@dd.custom_column_generator(required_columns=["name"], model_aliases=["my-model"])
2
def generate_message(row: dict, generator_params: None, models: dict) -> dict:
3
    response, _ = models["my-model"].generate(prompt=f"Greet {row['name']}")
4
    row["greeting"] = response
5
    return row

1
@dd.custom_column_generator(
2
    required_columns=["col1"],        # DAG ordering
3
    side_effect_columns=["extra"],    # Additional columns created
4
    model_aliases=["model1"],         # Required for LLM access
5
)

1
@dd.custom_column_generator(model_aliases=["my-model"])
2
def my_generator(row: dict, generator_params: None, models: dict) -> dict:
3
    model = models["my-model"]
4
    response, trace = model.generate(
5
        prompt="...",
6
        parser=my_custom_parser,  # optional, defaults to identity
7
        system_prompt="...",
8
        max_correction_steps=3,
9
    )
10
    row["result"] = response
11
    return row

1
@dd.custom_column_generator(
2
    required_columns=["topic"],
3
    side_effect_columns=["variation_id"],
4
)
5
def expand_topics(df: pd.DataFrame, params: None, models: dict) -> pd.DataFrame:
6
    rows = []
7
    for _, row in df.iterrows():
8
        for i in range(3):  # Generate 3 variations per input
9
            rows.append({
10
                "topic": row["topic"],
11
                "question": f"Question {i+1} about {row['topic']}",
12
                "variation_id": i,
13
            })
14
    return pd.DataFrame(rows)
15
16
dd.CustomColumnConfig(
17
    name="question",
18
    generator_function=expand_topics,
19
    generation_strategy=dd.GenerationStrategy.FULL_COLUMN,
20
    allow_resize=True,
21
)

1
@dd.custom_column_generator(required_columns=["id"])
2
def expand_row(row: dict) -> list[dict]:
3
    return [
4
        {**row, "variant": "a"},
5
        {**row, "variant": "b"},
6
    ]
7
8
dd.CustomColumnConfig(
9
    name="variant",
10
    generator_function=expand_row,
11
    generation_strategy=dd.GenerationStrategy.CELL_BY_CELL,
12
    allow_resize=True,
13
)

1
@dd.custom_column_generator(
2
    required_columns=["topic"],
3
    side_effect_columns=["draft", "critique"],
4
    model_aliases=["writer", "editor"],
5
)
6
def writer_editor(row: dict, generator_params: None, models: dict) -> dict:
7
    draft, _ = models["writer"].generate(prompt=f"Write about '{row['topic']}'")
8
    critique, _ = models["editor"].generate(prompt=f"Critique: {draft}")
9
    revised, _ = models["writer"].generate(prompt=f"Revise based on: {critique}\n\nOriginal: {draft}")
10
11
    row["final_text"] = revised
12
    row["draft"] = draft
13
    row["critique"] = critique
14
    return row

1
data_designer = DataDesigner()
2
models = data_designer.get_models(["my-model"])
3
result = my_generator({"name": "Alice"}, None, models)

1
from unittest.mock import MagicMock
2
from data_designer.engine.models.facade import ModelFacade
3
4
mock_model = MagicMock(spec=ModelFacade)