Smarter Mannequin Tuning: An AI Agent with LangGraph + Streamlit That Boosts ML Efficiency

Generalists Can Additionally Dig Deep

3 Methods to Velocity Up and Enhance Your XGBoost Fashions

every day just a little extra whereas working with LangGraph.

Let’s face it: since LangChain is among the first frameworks to deal with the mixing with LLMs, it took off earlier and have become form of a go-to choice in terms of constructing production-ready brokers, whether or not you prefer it or not.

LangChain’s youthful brother is LangGraph. This framework makes use of a graph notation with nodes and edges to construct the functions, making them extremely customizable and really sturdy. That’s what I’m having fun with a lot.

At first, some notations felt unusual to me (possibly it’s simply me!). However I stored digging and studying extra. And I strongly imagine we be taught higher whereas we’re implementing stuff, as a result of that’s when the actual issues pop up. So, after just a few traces of code and a few hours of code debugging, that graph structure began to make rather more sense to me, and I began to take pleasure in creating issues with LangGraph.

Anyway, should you don’t have any introduction to the framework, I like to recommend you take a look at this publish [1].

Now, let’s be taught extra concerning the mission of this text.

The Undertaking

On this mission, we’re going to construct a multi-step agent:

It takes in a machine studying mannequin sort: classification or regression.
And we may even enter the metrics of our mannequin, similar to accuracy, RMSE, confusion matrix, ROC, and so on. The extra we offer to the agent, the higher the response.

The agent, geared up with Google Gemini 2.0 Flash:

Reads the enter
Consider the mannequin’s metric inputted by the consumer
Return an actionable listing of recommendations to tune the mannequin and enhance its efficiency.

That is the mission folder construction:

ml-model-tuning/
├── langgraph_agent/
│ ├── graph.py #LangGraph logic
│ ├── nodes.py #LLMs and instruments
├── predominant.py # Streamlit interface to run the agent
├── necessities.txt

The Agent is reside and deployed on this net app.

Dataset

The dataset for use is a quite simple toy dataset named Suggestions, from the Seaborn bundle, and open-sourced underneath the license BSD 3. I made a decision to make use of a easy dataset like this as a result of it has each categorical and numerical options, being fitted to each varieties of mannequin creation. As well as, the beginning of the article is the agent, so that’s the place we need to spend extra consideration.

To load the information, use the next code.

import seaborn as sns

# Knowledge
df = sns.load_dataset('suggestions')

Subsequent, we are going to construct the nodes.

Nodes

The nodes of a LangGraph object are Python capabilities. They are often instruments that the agent will use or an occasion of an LLM. We construct every node as a separate perform.

However first, we now have to load the modules.

import os
from textwrap import dedent
from dotenv import load_dotenv
load_dotenv()

import streamlit as st
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

Our first node is the one to get the mannequin sort. It merely will get the enter from the consumer on whether or not the mannequin to be enhanced is a regression or a classification.

def get_model_type(state):
    """Test if the consumer is implementing a classification or regression mannequin. """
    
    # Outline the mannequin sort
    modeltype = st.text_input("Please let me know the kind of mannequin you might be engaged on and hit Enter:", 
                              placeholder="(C)lassification or (R)egression", 
                              assist="C for Classification or R for Regression")
    
    # Test if the mannequin sort is legitimate
    if modeltype.decrease() not in ["c", "r", "classification", "regression"]:
        st.information("Please enter a legitimate mannequin sort: C for (C)lassification or R for (R)egression.")
        st.cease()
        
    if modeltype.decrease() in ["c", "classification"]:
        modeltype = "classification"
    elif modeltype.decrease() in ["r", "regression"]:
        modeltype = "regression"
    
    return {"model_type": modeltype.decrease()} # "classification" or "regression"

The opposite two nodes from this graph are virtually the identical, however they differ within the system immediate. One is optimized for regression fashions analysis, whereas the opposite is specialised in classification. I’ll paste solely one in every of them right here. The whole code is offered on GitHub, although. See all of the nodes’ code right here.

def llm_node_regression(state):
    """
    Processes the consumer question and search outcomes utilizing the LLM and returns a solution.
    """
    llm = ChatGoogleGenerativeAI(
        mannequin="gemini-2.5-flash",
        api_key=os.environ.get("GEMINI_API_KEY"),
        temperature=0.5,
        max_tokens=None,
        timeout=None,
        max_retries=2
    )

    # Create a immediate
    messages = ChatPromptTemplate.from_messages([
        ("system", dedent("""
                          You are a seasoned data scientist, specialized in regression models. 
                          You have a deep understanding of regression models and their applications.
                          You will get the user's result for a regression model and your task is to build a summary of how to improve the model.
                          Use the context to answer the question.
                          Give me actionable suggestions in the form of bullet points.
                          Be concise and avoid unnecessary details. 
                          If the question is not about regression, say 'Please input regression model metrics.'.
                          
                          """)),
        MessagesPlaceholder(variable_name="messages"),
        ("user", state["metrics_to_tune"])
    ])
    
    # Create a sequence
    chain = messages | llm
    response = chain.invoke(state)
    return {"final_answer": [response]}

Nice. Now it’s time to stick these nodes collectively by constructing the sides to attach them. In different phrases, constructing the movement of the data from the consumer enter to the ultimate output.

Graph

The file graph.py shall be used to generate the LangGraph object. First, we have to import the modules.

from langgraph.graph import StateGraph, END
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
from langgraph_agent.nodes import llm_node_classification, llm_node_regression, get_model_type

The following step is to create the state of the graph. StateGraph manages the agent’s state all through the workflow. It retains monitor of the data the agent has gathered and processed. It’s nothing however a category with the names of the variables and their sort written in dictionary fashion.

# Create a state graph
class AgentState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        messages: A listing of messages within the dialog, together with consumer enter and agent outputs.
        model_type: The kind of mannequin getting used, both "classification" or "regression".
        query: The preliminary query from the consumer.
        final_answer: The ultimate reply supplied by the agent.
    """
    messages: Annotated[AnyMessage, add_messages] # accumulate messages
    model_type: str
    metrics_to_tune: str
    final_answer: str

To construct the graph, we are going to use a perform that:

Provides every node with a tuple ("identify", node_function_name)
Defines the start line on the get_model_type node. .set_entry_point("get_model_type")
Then, there’s a conditional edge, that decides to go to the suitable node relying on the response from the get_model_type node.
Lastly, join the LLM nodes to the END state.
Compile the graph to make it prepared to be used.

def build_graph():
    # Construct the LangGraph movement
    builder = StateGraph(AgentState)

    # Add nodes
    builder.add_node("get_model_type", get_model_type)
    builder.add_node("classification", llm_node_classification)
    builder.add_node("regression", llm_node_regression)

    # Outline edges and movement
    builder.set_entry_point("get_model_type")

    builder.add_conditional_edges(
        "get_model_type",
        lambda state: state["model_type"],
        {
            "classification": "classification",
            "regression": "regression"
        }
    )

    builder.add_edge("classification", END)
    builder.add_edge("regression", END)

    # Compile the graph
    return builder.compile()

If you wish to see the graph, you should utilize this little snippet.

# Create the graph picture and save png
from IPython.show import show, Picture
graph = build_graph()
show(Picture(graph.get_graph().draw_mermaid_png(output_file_path="graph.png")))

Picture of the graph created. Picture by the writer.

It’s a easy agent, but it surely works very nicely. We’ll get to that quickly. However we have to construct the front-end piece first.

Constructing the Consumer Interface

The consumer interface is a Streamlit app. I’ve chosen this selection resulting from straightforward prototyping and deployment options.

Let’s load the libraries wanted as soon as once more.

import os
from langgraph_agent.graph import AgentState, build_graph
from textwrap import dedent
import streamlit as st

Configuring the web page format (title, icon, sidebar and so on).

## Config web page
st.set_page_config(page_title="ML Mannequin Tuning Assistant",
                   page_icon='🤖',
                   format="huge",
                   initial_sidebar_state="expanded")

Creating the sidebar that holds the sphere so as to add a Google Gemini API Key and the restart session button.

## SIDEBAR | Add a spot to enter the API key
with st.sidebar:
    api_key = st.text_input("GOOGLE_API_KEY", sort="password")

    # Save the API key to the atmosphere variable
    if api_key:
        os.environ["GEMINI_API_KEY"] = api_key

    # Clear
    if st.button('Clear'):
        st.rerun()

Now, we add the web page’s title and directions to make use of the agent. That is all easy code utilizing principally the perform st.write().

## Title and Directions
if not api_key:
    st.warning("Please enter your OpenAI API key within the sidebar.")
    
st.title('ML Mannequin Tuning Assistant | 🤖')
st.caption('This AI Agent is will allow you to tuning your machine studying mannequin.')
st.write(':pink[**1**] | 👨‍💻 Add the metrics of your ML mannequin to be tuned within the textual content field. The extra metrics you add, the higher.')
st.write(':pink[**2**] | ℹ️ Inform the AI Agent what sort of mannequin you might be engaged on.')
st.write(':pink[**3**] | 🤖 The AI Agent will reply with recommendations on  enhance your mannequin.')
st.divider()

# Get the consumer enter
textual content = st.text_area('**👨‍💻 Add right here the metrics of your ML mannequin to be tuned:**')

st.divider()

And, lastly, the code to:

Run the build_graph() perform and create the agent.
Create the preliminary state of the agent, with an empty messages.
Invoke the agent.
Print the outcomes on display screen.

## Run the graph

# Spinner
with st.spinner("Gathering Tuning Solutions...", show_time=True):
    from langgraph_agent.graph import build_graph
    agent = build_graph()

    # Create the preliminary state for the agent, with clean messages and the consumer enter
    immediate = {
        "messages": [],
        "metrics_to_tune": textual content
    }

    # Invoke the agent
    outcome = agent.invoke(immediate)

    # Print the agent's response
    st.write('**🤖 Agent Response:**')
    st.write(outcome['final_answer'][0].content material)

All created. It’s time to put this AI Agent to work!

So, we are going to construct some fashions and ask the agent for tuning recommendations.

Operating the Agent

Effectively, as that is an agent that helps us with mannequin tuning recommendations, we should have a mannequin to tune.

Regression Mannequin

We’ll strive the regression mannequin first. We will shortly construct a easy mannequin.

# Imports
import pandas as pd
import numpy as np
import seaborn as sns
from sklearn.pipeline import Pipeline
from feature_engine.encoding import OneHotEncoder
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import root_mean_squared_error

## Baseline Mannequin
# Knowledge
df = sns.load_dataset('suggestions')

# Practice Take a look at Break up
X = df.drop('tip', axis=1)
y = df['tip']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Categorical
cat_vars = df.select_dtypes(embody=['object']).columns

# Pipeline
pipe = Pipeline([
    ('encoder', OneHotEncoder(variables=['sex', 'smoker', 'day', 'time'],
                              drop_last=True)),
    ('mannequin', LinearRegression())
])

# Match
pipe.match(X_train, y_train)

Now, we now have to collect metrics information to current to our AI Agent as a way to get tuning recommendations. The extra information, the higher. As I’m working with a regression mannequin, I selected to current the next data:

Function names
Statistical description of the dataset
R²
Root Imply Squared Error (RMSE)
Regression intercept and coefficients
VIF

        total_bill         tip   intercourse smoker  day    time        dimension
depend   244.000000  244.000000   244    244  244     244  244.000000
distinctive         NaN         NaN     2      2    4       2         NaN
prime            NaN         NaN  Male     No  Sat  Dinner         NaN
freq           NaN         NaN   157    151   87     176         NaN
imply     19.785943    2.998279   NaN    NaN  NaN     NaN    2.569672
std       8.902412    1.383638   NaN    NaN  NaN     NaN    0.951100
min       3.070000    1.000000   NaN    NaN  NaN     NaN    1.000000
25%      13.347500    2.000000   NaN    NaN  NaN     NaN    2.000000
50%      17.795000    2.900000   NaN    NaN  NaN     NaN    2.000000
75%      24.127500    3.562500   NaN    NaN  NaN     NaN    3.000000
max      50.810000   10.000000   NaN    NaN  NaN     NaN    6.000000

---
Mannequin: Linear Regression
Rating: 0.44
RMSE: 0.84
Intercept: 0.45

Coefficients:
       function  coefficient
0   total_bill     0.094700
1         dimension     0.233484
2     sex_Male     0.028819
3    smoker_No     0.192353
4      day_Sat    -0.006064
5      day_Fri     0.179721
6      day_Sun     0.128928
7  time_Dinner    -0.094957

VIF:
total_bill    2.226294
tip           1.879238
dimension          1.590524

Now I’ll run the agent.

Right here is the agent’s response:

🤖 Agent Response:

Listed here are actionable recommendations to enhance your regression mannequin:

Discover non-linear relationships: Think about including polynomial options for total_bill and dimension, or interplay phrases between options (e.g., total_bill * dimension), as the present linear mannequin is perhaps too simplistic for the underlying information patterns.

Consider different regression fashions: Given the R-squared of 0.44, check different fashions like Random Forest Regressor, Gradient Boosting Regressor, or Assist Vector Regressor, which might seize extra complicated, non-linear relationships.

Deal with information distribution and outliers: Examine and deal with outliers in total_bill and the goal variable tip. Think about making use of transformations (e.g., log rework) to skewed options to raised meet linearity assumptions and enhance mannequin efficiency.

Analyze function statistical significance: Receive p-values for every coefficient to determine options that is probably not statistically vital. Eradicating or re-evaluating such options can simplify the mannequin and probably enhance generalization.

There are a few recommendations right here. We will now select what we are going to settle for or not. Here’s what I attempted (code in GitHub):

I educated a Random Forest Regressor, however the outcome was not good with the out of the field mannequin, dropping the R² to 0.25 and the RMSE to 0.97. So I discarded that choice.
So, if I’m conserving the Linear Regression, one other suggestion is to make use of log transformations and deal with outliers. I attempted that, and the result’s higher. The mannequin goes to an R² of 0.55 and RMSE of 0.23. A big enchancment.

Classification Mannequin

I adopted the identical drill right here, however now engaged on a classification mannequin, utilizing the identical dataset and attempting to foretell if the restaurant’s buyer is a smoker or not.

Skilled a classification mannequin
Obtained the preliminary metrics: Rating = 0.69; RMSE = 0.55
Ran the AI Agent for recommendations
Utilized some tuning recommendations: class_weight='balanced' and BayesSearchCV.
Obtained the tuned metrics: Rating = 0.71; RMSE = 0.52

AI Agent’s recommendations. Picture by the writer.

Discover how the Precision vs. Recall is extra balanced as nicely.

Rating Earlier than vs. After tuning. Picture by the writer.

Our job is full. The agent is working as designed.

Earlier than You Go

We’ve got reached the tip of this mission. Total, I’m happy with the outcome. This mission is kind of easy and fast to construct, and but it delivers loads of worth!

Tuning fashions shouldn’t be a one-size-fits-all motion. There are numerous choices to strive. Thus, having the assistance of an AI Agent to offer us just a few concepts to strive could be very precious and makes our job simpler with out changing us.

Strive the app for your self and let me know if it helped you get an improved efficiency metric!

https://ml-tuning-assistant.streamlit.app