How We Used Python and OpenAI Agents to Automate Data Cleaning for Our ML Pipelines

If you've ever spent hours wrangling messy CSVs or chasing down weird null values in a dataset, you know data cleaning is the least glamorous part of machine learning. It’s repetitive, it’s error-prone, and it always takes way longer than you think. That’s exactly what we were up against—until we started using Python and OpenAI agents to automate the dirty work. It didn’t just free up our time; it changed how we think about workflow efficiency.

Why Data Cleaning Feels Like Groundhog Day

Honestly, most ML projects stall not because the model is hard, but because the data is messy. You’d expect an email column to have only valid emails, but it’s full of random strings, missing values, and even phone numbers. Dates are all over the place. Sometimes you get weird encodings, or someone used "N/A" instead of null. Multiply that by dozens of columns and hundreds of thousands of rows, and it’s a recipe for frustration.

For our team, cleaning datasets was eating up weeks across sprints. And frankly, we got bored. That boredom led us to ask: what if we could offload the grunt work to an AI agent—one that understands context and can suggest fixes, not just run static scripts?

What Are OpenAI Agents and Why Use Them?

OpenAI agents are basically LLM-powered tools that can interpret instructions, generate code, and even analyze data when you point them at a task. The cool thing is, they’re not just parroting canned responses—they can reason through tasks, ask for clarifications, and adapt as they go.

In our setup, we use Python as the glue. The agent gets a sample of the dataset, we prompt it with cleaning goals, and it returns Python code or suggestions. We review, tweak, and run. It’s not magic, but it’s surprisingly effective.

Setting Up: The Minimum You Need

You don’t need fancy infrastructure. All you need is:

• Python 3.x
• openai Python package (for API access)
• Pandas (for data manipulation)
• An OpenAI API key

Here’s a basic skeleton for querying an OpenAI agent in Python:

import openai

openai.api_key = 'your-api-key'  # Replace with your actual OpenAI key

def ask_agent(prompt):
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",  # Pick a model available to you
        messages=[{"role": "user", "content": prompt}],
        max_tokens=800
    )
    return response.choices[0].message.content

# Example: Ask the agent how to clean missing values in a pandas DataFrame
prompt = (
    "I have a pandas DataFrame with missing values in the 'email' column. "
    "What's the best way to clean it? Please provide Python code."
)
code_suggestion = ask_agent(prompt)
print(code_suggestion)  # Output will be Python code as a string

This function sends your prompt to the agent and prints the code it suggests. I’ve tested this with simple cleaning tasks, and while the code isn’t always plug-and-play, it’s a great starting point—especially for repetitive stuff.

Example 1: Automating Null Handling and Typo Correction

We had a dataset with user emails, and about 15% were missing or invalid. Manually fixing these would have taken days. So we used the agent to suggest a cleaning function.

Here’s how we integrated its suggestion:

import pandas as pd

# Sample data
df = pd.DataFrame({
    'email': ['[e-mail]', 'jane_doe@exampl', None, 'bob@.com', '[e-mail]']
})

def clean_emails(df):
    # Fill missing values with empty string
    df['email'] = df['email'].fillna('')
    # Use a regex to filter valid emails (simple pattern)
    df['email'] = df['email'].where(
        df['email'].str.match(r'^[\w\.-]+@[\w\.-]+\.[a-z]{2,}$'),
        ''
    )
    return df

df_clean = clean_emails(df)
print(df_clean)

Key lines:

• fillna(''): replaces None with empty strings, which makes downstream processing easier.
• str.match(...): checks for valid email format. Invalid emails are replaced with empty strings.

The agent suggested a similar regex, which we then tweaked for our use case. Crucially, we could ask it for improvements: “How can I catch more edge cases?” This iterative approach sped things up.

Example 2: Standardizing Date Formats

Dates are a pain. Our data had dates like "2024-01-05", "05/01/2024", "January 5th, 2024", and even "5 Jan 24". We wanted everything in ISO format.

We prompted the agent: "I have a pandas DataFrame with a 'date' column in mixed formats. Write code to standardize to YYYY-MM-DD."

Here’s the solution (with some tweaks):

import pandas as pd
from dateutil import parser

# Sample data
df = pd.DataFrame({
    'date': ['2024-01-05', '05/01/2024', 'January 5th, 2024', '5 Jan 24', None]
})

def standardize_dates(df):
    # Apply dateutil.parser to each value, handle missing values
    def parse_date(x):
        if pd.isnull(x):
            return None
        try:
            return parser.parse(x).strftime('%Y-%m-%d')
        except Exception:
            return None  # Return None if parsing fails

    df['date'] = df['date'].apply(parse_date)
    return df

df_clean = standardize_dates(df)
print(df_clean)

Key lines:

• parser.parse(x): handles a ridiculous number of date formats.
• strftime('%Y-%m-%d'): outputs in ISO format.

Turns out, the agent’s code was almost right—except it didn't handle exceptions or nulls. We fixed that, and now date cleaning is totally automated.

Example 3: Using Agents for Column Type Detection

Another headache: columns mislabeled as strings when they should be numeric, or vice versa. We asked the agent to suggest a function to auto-detect column types.

Here’s a practical version:

import pandas as pd

# Sample data
df = pd.DataFrame({
    'age': ['25', 'thirty', '40', None, '18'],
    'income': ['50000', '70000', 'not available', '60000', None]
})

def detect_and_convert_types(df):
    for col in df.columns:
        # Try converting to numeric, errors go to NaN
        df[col] = pd.to_numeric(df[col], errors='coerce')
    return df

df_converted = detect_and_convert_types(df)
print(df_converted)

Key lines:

• pd.to_numeric(..., errors='coerce'): turns invalid entries into NaN, so downstream models won’t choke.

The agent's suggestion was similar, but missed handling some edge cases (like custom "not available" strings). We added pre-processing to replace those with NaN before conversion.

Common Mistakes When Using AI for Data Cleaning

I’ve seen (and made) these mistakes more than once:

• Blindly trusting agent output: The agent can make mistakes or oversimplify. Always review the code and test it on a sample before running it on production data. I spent a weekend debugging a column where the agent’s regex was too permissive.
• Not providing enough context: If your prompt is vague (“Clean my data”), you’ll get generic answers. Be specific: mention column names, sample values, and desired formats.
• Skipping error handling: Agents often suggest code without try/except blocks. If you don’t add robust error handling, you’ll end up with silent failures or cryptic stack traces.

Key Takeaways

• Automating data cleaning with Python and OpenAI agents can save weeks of repetitive work, but it’s not totally hands-off—you still need to review and tweak.
• Providing detailed, context-rich prompts leads to much better agent suggestions. The more specific you are, the more useful the code.
• AI-generated code is a starting point, not an endpoint. Test, validate, and customize for your actual data quirks.
• Handling exceptions and edge cases is essential. Don’t assume the agent’s code covers everything.
• Cleaning data is where workflow efficiency matters most. The time you save here pays off down the line when iterating on models.

Closing Thoughts

Honestly, automating data cleaning with AI agents gave our team superpowers and made the whole ML pipeline smoother. If you’re tired of grunt work, it’s worth experimenting with. Just remember: AI helps, but you’re still the one in charge of quality.

---

If you found this helpful, check out more programming tutorials on our blog. We cover Python, JavaScript, Java, Data Science, and more.