ML Guardian: AI Agent for ML Risk Detection

• 

  ML Guardian dashboard showing ML asset information and DataHub connection.

Inspiration

Machine learning models can fail silently when the data, schemas, features, or upstream datasets they depend on change. These risks can be difficult to trace across complex ML pipelines. We wanted to build an AI agent that uses DataHub metadata and lineage context to help developers understand these risks before they affect a model.

What it does

ML Guardian is an AI-powered assistant for analyzing ML pipeline risks. A user can ask questions such as: "Is my customer churn model at risk?" The agent uses DataHub metadata and lineage information to identify relevant datasets, models, dependencies, and changes. It then produces a risk assessment, explains detected issues, and recommends practical actions.

Key Features

• Discover ML assets and their dependencies
• Analyze lineage between datasets, features, and models
• Identify potential schema and upstream-data risks
• Generate an understandable risk assessment
• Explain why a model may be at risk
• Recommend actions such as validation or retraining

How it works

ML Guardian connects AI reasoning with DataHub's metadata and lineage context. The agent retrieves relevant information about ML assets and uses that context to analyze possible risks.

The application provides a simple interface where users can ask questions about their ML assets and receive a structured risk analysis.

Why DataHub

DataHub provides the metadata and lineage context needed to understand relationships between data assets and ML systems. Instead of acting as a generic chatbot, ML Guardian uses structured metadata and lineage information to ground its analysis.

Challenges

The main challenge is turning complex metadata and lineage relationships into a concise and useful risk assessment. Another challenge is making the results understandable to users who may not be familiar with metadata systems.

What we learned

This project explores how metadata, lineage, and AI agents can work together to solve practical problems in machine learning operations.

Future Improvements

• Automated monitoring of schema and lineage changes
• More detailed ML and data-quality risk scoring
• Alerts for high-risk model dependencies
• Integration with additional ML monitoring tools
• Automated recommendations for model validation and retraining

What's next for ML Guardian: AI Agent for ML Risk Detection

We plan to make ML Guardian more proactive and reliable by adding automated monitoring for schema, data-quality, and lineage changes.

Future versions will provide real-time alerts for high-risk ML dependencies, improved risk scoring, deeper integration with ML monitoring tools, and automated recommendations for validation and retraining.

Our goal is to evolve ML Guardian from a risk-analysis assistant into a proactive AI agent that helps teams maintain safer and more reliable ML pipelines.

Built With

• ai
• datahub
• flask
• github
• machine-learning
• pandas
• python
• scikit-learn