Skip to content

schema

agentic_ml.data.schema

Description and aggregate profiling of a tabular dataset.

Only aggregate information (dtypes, missing counts, cardinality, target summary) is ever exposed to the LLM. The raw data never leaves the machine through the prompt; it is only read by the code executed inside the sandbox.

DatasetSchema dataclass

DatasetSchema(variables: dict[str, str], target: str)

User-provided description of a dataset, driving both validation and prompting.

A schema is the only thing a user must supply about their data beyond the dataframe itself. It is used to check the data has the expected columns (:meth:validate), to compute the aggregate, LLM-safe profile shown to the agent (:meth:profile and :meth:describe_for_prompt), and to split the dataframe into features/target inside :func:~agentic_ml.core.engine.setup_research.

Attributes:

Name Type Description
variables dict[str, str]

Maps each column name to a short, natural-language description (e.g. {"age": "customer age in years"}). Columns present in the data but missing here are still profiled, just without a human description; keys that are not present in the data make :meth:validate raise a :class:ValueError.

target str

The name of the column to predict. Every other column is treated as a feature by :meth:feature_names.

feature_names

feature_names(data: DataFrame) -> list[str]

Return every column of data except the target.

Source code in src/agentic_ml/data/schema.py
37
38
39
def feature_names(self, data: pd.DataFrame) -> list[str]:
    """Return every column of ``data`` except the target."""
    return [column for column in data.columns if column != self.target]

validate

validate(data: DataFrame) -> None

Check that the target and described variables exist in data.

Source code in src/agentic_ml/data/schema.py
41
42
43
44
45
46
47
def validate(self, data: pd.DataFrame) -> None:
    """Check that the target and described variables exist in ``data``."""
    if self.target not in data.columns:
        raise ValueError(f"target column '{self.target}' is not present in the data")
    unknown = set(self.variables) - set(data.columns)
    if unknown:
        raise ValueError(f"described variables not found in the data: {sorted(unknown)}")

profile

profile(data: DataFrame) -> dict

Compute an aggregate, LLM-safe profile of data.

Source code in src/agentic_ml/data/schema.py
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def profile(self, data: pd.DataFrame) -> dict:
    """Compute an aggregate, LLM-safe profile of ``data``."""
    columns = {}
    for column in data.columns:
        series = data[column]
        info: dict[str, object] = {
            "dtype": str(series.dtype),
            "missing": int(series.isna().sum()),
            "n_unique": int(series.nunique(dropna=True)),
        }
        if pd.api.types.is_numeric_dtype(series):
            described = series.describe()
            info["min"] = float(described.get("min", float("nan")))
            info["mean"] = float(described.get("mean", float("nan")))
            info["max"] = float(described.get("max", float("nan")))
        columns[column] = info
    return {"n_rows": int(len(data)), "target": self.target, "columns": columns}

describe_for_prompt

describe_for_prompt(data: DataFrame) -> str

Render a compact textual description of the schema and profile for the prompt.

Source code in src/agentic_ml/data/schema.py
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def describe_for_prompt(self, data: pd.DataFrame) -> str:
    """Render a compact textual description of the schema and profile for the prompt."""
    profile = self.profile(data)
    lines = [
        f"Rows: {profile['n_rows']}",
        f"Target column: {self.target}",
        "Variables:",
    ]
    for column, info in profile["columns"].items():
        description = self.variables.get(column, "")
        role = "TARGET" if column == self.target else "feature"
        stats = f"dtype={info['dtype']}, missing={info['missing']}, unique={info['n_unique']}"
        if "mean" in info:
            stats += f", min={info['min']:.4g}, mean={info['mean']:.4g}, max={info['max']:.4g}"
        suffix = f" — {description}" if description else ""
        lines.append(f"  - {column} ({role}; {stats}){suffix}")
    return "\n".join(lines)