Build ML agents that use local, pre-trained models

AI agents are no longer limited to acting like simple assistants. Modern agents can execute complex machine learning (ML) workflows—including forecasting, anomaly detection, classification, and automated decision support. By combining reasoning with tools, agents can independently complete tasks that traditionally required skilled ML engineers.

However, most agents rely on general-purpose large language models (LLMs) for reasoning and orchestration. This creates a challenge for enterprises: sensitive or proprietary data may be exposed to external foundation models during processing. Organizations with confidential financial data, customer information, or intellectual property often cannot risk sending raw data to public LLM services due to security, compliance, and governance requirements.

One effective solution is to pre-train or fine-tune smaller domain-specific models on internal enterprise data and securely store those models locally in protected serialized formats such as password-protected pickle files. Instead of exposing raw enterprise data to a general LLM, the agent can call a dedicated tool that interacts with the trained local ML model to perform the required task. In this architecture, the agent only receives the prediction, insight, or output generated by the model, while the underlying training data remains completely hidden and inaccessible to the LLM itself. This enables organizations to build intelligent agents while maintaining strong data privacy, security, and governance controls.

Sample application: used car price predictor

In this article, an ML model is trained to predict the used car price based on 2 variables: km_run and engine_size, using synthetic data. A tool and agent in watsonx Orchestrate are then created to use this model.

The following diagram shows (1) watsonx.ai, used to build and train the model and generate the pickle file, and (2) watsonx Orchestrate with (a) the tool that calls the trained model and (b) the agent that users interact with to predict used car prices.

Build and train the model

While this example uses watsonx.ai to build and train the model, you can use any environment you like (such as PyTorch).

Use the following Python code to generate some synthetic data, train a random forest regressor model, and then store the trained model as a pickle (.pkl) file named used_car_price_model.pkl.

import pandas as pd
import json
import pickle
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, r2_score

# -----------------------------
# 1. Create synthetic JSON dataset
# -----------------------------

np.random.seed(42)

n = 100

km_run = np.random.randint(5_000, 200_000, n)
engine_size = np.random.randint(800, 3000, n)

# Synthetic pricing formula
price = (
    900000
    - (km_run * 2.7)
    + (engine_size * 130)
    + np.random.normal(0, 45000, n)
)

# Avoid unrealistic negative prices
price = np.maximum(price, 80_000)

dataset = []

for i in range(n):
    dataset.append({
        "km_run": int(km_run[i]),
        "engine_size": int(engine_size[i]),
        "price_rupees": int(price[i])
    })

# Save dataset as JSON
json_dataset_path = "<path>/used_car_dataset.json"

with open(json_dataset_path, "w") as f:
    json.dump(dataset, f, indent=4)

# -----------------------------
# 2. Load JSON dataset for training
# -----------------------------

with open(json_dataset_path, "r") as f:
    training_data = json.load(f)

df = pd.DataFrame(training_data)

X = df[["km_run", "engine_size"]]
y = df["price_rupees"]

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# -----------------------------
# 3. Train regression model
# -----------------------------

model = RandomForestRegressor(
    n_estimators=100,
    random_state=42
)

model.fit(X_train, y_train)

# Evaluate model
y_pred = model.predict(X_test)

mae = mean_absolute_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

# -----------------------------
# 4. Save model as pickle
# -----------------------------

model_path = "<path>/used_car_price_model.pkl"

with open(model_path, "wb") as f:
    pickle.dump(model, f)

# -----------------------------
# 5. Load model and run inference using JSON input
# -----------------------------

with open(model_path, "rb") as f:
    loaded_model = pickle.load(f)

# JSON input from agent/user
inference_json = {
    "km_run": 40000,
    "engine_size": 1500
}

# Convert JSON input to DataFrame
inference_df = pd.DataFrame([{
    "km_run": inference_json["km_run"],
    "engine_size": inference_json["engine_size"]
}])

predicted_price = loaded_model.predict(inference_df)[0]

# JSON output
inference_output_json = {
    "input": inference_json,
    "predicted_price_rupees": int(round(predicted_price))
}

print("Model Training Complete")
print("MAE:", round(mae, 2))
print("R2 Score:", round(r2, 4))
print("\nInference JSON Output:")
print(json.dumps(inference_output_json, indent=4))

json_dataset_path, model_path

Create a tool and an agent

Now that you have your pre-trained ML model, you can create an agent that uses that model. Follow these steps.

1. Set up the IBM watsonx Orchestrate Agent Development Kit (ADK) to quickly build your agent. (Follow the watsonx Orchestrate ADK installation guide to get started.)

2. Create a folder for your agent, named wxo_ML_agent. Within that folder, create a sub-folder named car_price_predictor. In this sub-folder, create your agent YAML file, price_predictor_agent.yaml that contains the following YAML. (To learn more about how to build agents with the watsonx Orchestrate ADK, read the Authoring agents documentation).

    . spec_version: v1
    kind: native
    name: used_car_price_predictor
    description: Predicts used car prices based on kilometers driven and engine size.
    instructions: |
      You are an agent that predicts used car prices.
   
          When a user asks about car price:
   
      1. Collect two values: kilometers driven and engine size in cc.
       - If engine size is given in liters or another unit, convert it to cc.If the user just provides a random number you must ask him about the units. Your task is to convert the engine size into cc if any other units are provided by the user.
       - If distance is in miles, convert it to kilometers.similarly the total distance covered by the car must be in kilometers . you must ask and confirm 
       - Ask only if these values are missing.
      2. Once you have both values, call price_prediction_tool with:
       {"km_run": <int>, "engine_size": <int>}
      3. Do NOT show the JSON to the user. Do NOT estimate the price yourself.
       Always call the tool and use its output as the predicted price.
      4. Present the tool's predicted price to the user.
   
    llm: groq/openai/gpt-oss-120b
    style: default
    collaborators: []
    tools: 
    - price_prediction_tool
   

3. In the car_price_predictor sub-folder, create the tool, price_prediction_tool.py, with this Python code. (To learn more about how to build tools with the watsonx Orchestrate ADK, read the Authoring Python-based tools documentation).

    # price_prediction_tool.py
    import pickle
    import pandas as pd
    import traceback
    from pathlib import Path
    from ibm_watsonx_orchestrate.agent_builder.tools import tool
   
    MODEL_PATH = Path(__file__).parent / "used_car_price_model.pkl"
    with open(MODEL_PATH, "rb") as f:
        loaded_model = pickle.load(f)
   
    EXPECTED_COLUMNS = list(loaded_model.feature_names_in_)
   
    @tool(
        name="price_prediction_tool",
        description="Predicts the price of a used car. Input: JSON with 'km_run' (int) and 'engine_size' (int, in cc). Returns the predicted price as a string.",
    )
    def price_prediction_tool(json_input: dict) -> str:
        """Predicts the price of a car given its kilometers driven and engine size."""
        try:
            cleaned = {
   
                "km_run": int(json_input["km_run"]),
                "engine_size": int(json_input["engine_size"]),
            }
            input_df = pd.DataFrame([cleaned])[EXPECTED_COLUMNS]
            predicted_price = loaded_model.predict(input_df)[0]
            return f"Predicted price: {int(round(float(predicted_price)))}"
        except Exception as e:
            return f"ERROR: {type(e).__name__}: {str(e)} | input received: {json_input}"
   

4. Move the trained model file used_car_price_model.pkl into the car_price_predictor sub-folder.

5. Because the trained ML model requires specific libraries to work, create a requirements.txt file in the car_price_predictor sub-folder and paste the following code in it:

    scikit-learn==1.8.0
    numpy==2.4.4
    pandas==3.0.2
   

6. Import the tool and the agent into the watsonx Orchestrate ADK. (Learn more about how to import and deploy agents in the watsonx Orchestrate ADK documentation.)

Test the ML agent in watsonx Orchestrate

Now it's time to test the agent working with the trained model.

When you log in to your watsonx Orchestrate instance, the used_car_price_predictor agent is available. Open the agent and interact with it in the chat window. Ask it to estimate the car price by providing an engine size and kilometer run. The agent invokes the price_prediction_tool, which uses the trained model to estimate the car price.

Summary

In this article, you learned how to pre-train a domain-specific machine learning model on proprietary data and securely package it as a password-protected serialized model. This process enables AI agents to use a lightweight, specialized model tailored to the use case, rather than relying solely on a general-purpose foundation model within the agent framework.

Then, you learned how to develop a custom tool to expose the trained model's inference capabilities to the agent. The agent can invoke the tool to perform predictions and ML-driven tasks without direct access to the underlying training data, ensuring data privacy, security, and intellectual property protection while maintaining accurate, domain-specific outcomes.