ai-data-science-team

Implements a SupervisorDSTeam agent that routes natural language data science tasks to 10+ specialized agents using a state machine pattern built on LangGraph. The supervisor decomposes user requests, selects appropriate agents (DataLoaderAgent, DataCleaningAgent, FeatureEngineeringAgent, etc.), and chains their outputs together, maintaining dataset lineage across multi-step workflows. Uses CompiledStateGraph with conditional routing logic to dynamically dispatch to domain-specific agents based on task type.

Implements a coding agent pattern where specialized agents generate Python code via LLM, execute it in isolated subprocess sandboxes using run_code_sandboxed_subprocess(), capture errors, and automatically attempt fixes by re-prompting the LLM with error context. The BaseAgent class wraps a CompiledStateGraph with nodes for execution, error fixing, and explanation, enabling autonomous error recovery without user intervention. Supports multiple LLM providers (OpenAI, Anthropic, Ollama) through LangChain abstraction.

Implements a DataCleaningAgent that detects data quality issues (missing values, duplicates, outliers, type inconsistencies) and generates code to fix them. The agent analyzes data distributions, identifies anomalies, and applies appropriate cleaning techniques (imputation, deduplication, outlier removal, type conversion). Supports both statistical and domain-specific cleaning rules, with generated code that is transparent and modifiable.

Implements a DataWranglingAgent that generates code for complex data transformations (pivoting, melting, grouping, joining, filtering, sorting). The agent understands pandas operations and generates appropriate transformations from natural language descriptions. Supports multi-table operations (merges, concatenation) and complex aggregations, with generated code that is transparent and reusable.

Implements a DataLoaderAgent that loads data from multiple sources (CSV, Excel, JSON, Parquet, SQL databases, APIs) and returns pandas DataFrames. The agent handles format detection, encoding issues, and connection management. Supports both local files and remote data sources, with automatic schema inference and optional data preview.

Implements the AI Pipeline Studio application, a Streamlit-based visual interface for composing multi-agent workflows without code. Users drag-and-drop agent nodes (DataLoader, DataCleaner, FeatureEngineer, etc.), connect them with data flow edges, configure parameters through UI forms, and execute the pipeline. The studio generates the underlying agent orchestration code and provides real-time execution monitoring with error visualization.

Implements a PandasDataAnalyst workflow that orchestrates multiple agents (DataLoader, DataCleaner, DataWrangler, EDATools, FeatureEngineer, MLAgent) to perform end-to-end pandas-based data analysis. The workflow accepts a natural language task description, automatically decomposes it into sub-tasks, routes to appropriate agents, and chains results together. Generates a complete, reproducible pandas analysis script as output.

Implements a SQLDataAnalyst workflow that orchestrates SQL-based analysis using the SQLDatabaseAgent, with optional pandas integration for visualization and advanced analysis. The workflow accepts natural language queries, generates SQL code, executes against connected databases, and returns results as DataFrames. Supports exploratory queries, aggregations, and complex joins without requiring manual SQL writing.

Maintains a dataset registry that tracks parent-child relationships between datasets as they flow through the agent pipeline, recording which agent performed which transformation and when. Each dataset is assigned metadata including source, transformations applied, and downstream dependencies. The registry enables reproducible pipelines by allowing users to trace any output dataset back to its original source and understand the exact sequence of operations that produced it.

Provides 10+ pre-built specialized agents (DataLoaderAgent, DataCleaningAgent, DataWranglingAgent, FeatureEngineeringAgent, DataVisualizationAgent, EDAToolsAgent, SQLDatabaseAgent, MLAgent, ExperimentTrackingAgent) that inherit from BaseAgent and implement domain-specific prompts and tool bindings. Each agent is instantiated via create_coding_agent_graph() factory function, which configures the agent's system prompt, available tools, and execution environment. Agents can work independently or be composed by the SupervisorDSTeam for complex workflows.

Abstracts LLM provider selection through LangChain's language model interface, enabling seamless switching between OpenAI, Anthropic, Ollama, and other providers without code changes. Configuration is handled via environment variables or explicit provider specification at agent instantiation. Supports both cloud-based APIs (OpenAI GPT-4, Claude) and local models (Ollama) for air-gapped or privacy-sensitive deployments.

Generates complete, executable Python scripts that encapsulate the entire data science workflow performed by agents. Each script includes all data loading, transformation, visualization, and ML steps in a single reproducible file that can be version-controlled, shared, and re-executed independently of the agent system. Scripts include error handling, logging, and comments explaining each step, making them suitable for production deployment or team collaboration.

Implements a specialized SQLDatabaseAgent that generates SQL queries from natural language descriptions, executes them against connected databases, and returns results as pandas DataFrames. The agent understands database schema, handles connection management, and can perform exploratory queries, data extraction, and aggregations. Supports multiple database backends (PostgreSQL, MySQL, SQLite, etc.) through SQLAlchemy abstraction.

Implements an EDAToolsAgent that automatically generates exploratory visualizations, statistical summaries, and data quality reports from datasets. The agent analyzes column types, distributions, correlations, and missing values, then generates appropriate visualizations (histograms, scatter plots, heatmaps, box plots) using Plotly. Results are returned as interactive HTML visualizations and JSON summaries suitable for stakeholder communication.

Implements a FeatureEngineeringAgent that generates feature transformations (scaling, encoding, polynomial features, interactions, domain-specific features) from natural language descriptions. The agent analyzes the target variable and existing features, then generates code to create new features that improve model predictability. Supports both numeric and categorical feature engineering, with automatic selection of appropriate techniques (StandardScaler, OneHotEncoder, PolynomialFeatures, etc.).

Implements MLAgent and ExperimentTrackingAgent that generate model training code, execute training pipelines, and automatically log experiments to MLflow. The agent supports multiple model types (linear regression, decision trees, random forests, gradient boosting, neural networks), hyperparameter tuning, and cross-validation. Experiment metadata (parameters, metrics, artifacts) is logged to MLflow for tracking model performance across iterations.