train_linear_regression_model

Train a linear regression model on uploaded CSV data to predict values and evaluate performance using RMSE metrics.

Instructions

This function trains linear regression model.

Args: Takes input for output column name.

Returns: String which contains the RMSE value.

Input Schema

TableJSON Schema

Name	Required	Description	Default
`output_column`	Yes

Implementation Reference

server.py:123-165 (handler)

Implements the training of a linear regression model using scikit-learn. It retrieves data from context, prepares features and target, splits into train/test sets (90/10), fits the model, predicts on test set, computes and returns RMSE.

def train_linear_regression_model(output_column: str) -> str:
    """
    This function trains linear regression model.

    Args:
        Takes input for output column name.

    Returns:
        String which contains the RMSE value.
    """

    try:
        data = context.get_data()

        # Check if the output column exists in the dataset
        if output_column not in data.columns:
            return f"Error: '{output_column}' column not found in the dataset."

        # Prepare the features (X) and target variable (y)
        X = data.drop(columns=[output_column])  # Drop the target column for features
        y = data[output_column]  # The target variable is the output column

        # Split the data into training and test sets (80% train, 20% test)
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=42)

        # Initialize the Linear Regression model
        model = LinearRegression()

        # Train the model
        model.fit(X_train, y_train)

        # Predict on the test set
        y_pred = model.predict(X_test)

        # Calculate RMSE (Root Mean Squared Error)
        rmse = np.sqrt(mean_squared_error(y_test, y_pred))

        # Return the RMSE value
        return f"Model trained successfully. RMSE: {rmse:.4f}"

    except Exception as e:
        return f"An error occurred while training the model: {str(e)}"

server.py:122-122 (registration)
Registers the train_linear_regression_model function as an MCP tool using the FastMCP decorator.
```
@mcp.tool()
```

server.py:14-35 (helper)

DataContext dataclass provides shared storage for the pandas DataFrame used by the tool (via global context instance). Used in the handler to get_data().

@dataclass
class DataContext():
    """
    A class that stores the DataFrame in the context.
    """
    _data: pd.DataFrame = None

    def set_data(self, new_data: pd.DataFrame):
        """
        Method to set or update the data.
        """
        self._data = new_data

    def get_data(self) -> pd.DataFrame:
        """
        Method to get the data from the context.
        """
        return self._data

# Initialize the DataContext instance globally
context = DataContext()

Linear Regression MCP