Quick example¶

This notebook walks through the full agentic-ml workflow on the scikit-learn diabetes toy dataset (a small, bundled tabular regression problem: predict disease progression from baseline patient measurements).

Two clearly separated objects are used:

• TabularRegressionResearcher — runs the agentic experimentation loop: it proposes trials, evaluates each one with a fixed pipeline, and keeps an auditable leaderboard.csv of every trial. It needs an LLM and is only used at research time.
• AgenticModel — the exported, production-ready model with fit / predict / save / load. It has no LLM dependency: it only re-executes the plain Python training code produced by the winning trial.

0. Installation¶

agentic-ml ships with no LLM dependency by default. To run TabularRegressionResearcher.research(...) you need two extra packages:

1. The research extra of jcaste05-agentic-ml, which pulls in the Strands Agents framework that drives the agent loop:

   pip install "jcaste05-agentic-ml[research]"
   

2. A Strands model provider extra matching the LLM backend you want to call — Strands does not bundle any specific provider. This notebook uses Gemini, so it also needs:

   pip install "strands-agents[gemini]"
   

   Swap gemini for another provider extra (e.g. openai, anthropic, ...) if you use a different LLM — see the Strands docs for the full list.

AgenticModel.load(...) / .predict(...) (production use) never needs either of the two packages above: plain pip install jcaste05-agentic-ml is enough.

API key¶

You need a valid API key for whichever provider you choose:

• Gemini (used below): create a key in Google AI Studio and set it as GEMINI_API_KEY (e.g. in a local .env file, loaded via python-dotenv).

Everything except the research cell itself works without a live LLM or API key.

1. Load the dataset¶

We build a pandas DataFrame with the raw scikit-learn features plus the target column, and describe each variable in natural language — this description (never the raw data) is what the agent sees.

In [ ]:

import logging
import os
from dotenv import load_dotenv

load_dotenv()

from sklearn.datasets import load_diabetes

from agentic_ml.core.model import AgenticModel
from agentic_ml.data import DatasetSchema
from agentic_ml.estimators.regression import TabularRegressionResearcher
from agentic_ml.research.agent import AgentConfig

logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(name)s %(levelname)s %(message)s")
logging.getLogger("agentic_ml").setLevel(logging.DEBUG)

raw = load_diabetes(as_frame=True)
data = raw.frame.rename(columns={"target": "disease_progression"})
data.head()

import logging import os from dotenv import load_dotenv load_dotenv() from sklearn.datasets import load_diabetes from agentic_ml.core.model import AgenticModel from agentic_ml.data import DatasetSchema from agentic_ml.estimators.regression import TabularRegressionResearcher from agentic_ml.research.agent import AgentConfig logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(name)s %(levelname)s %(message)s") logging.getLogger("agentic_ml").setLevel(logging.DEBUG) raw = load_diabetes(as_frame=True) data = raw.frame.rename(columns={"target": "disease_progression"}) data.head()

In [ ]:

schema = DatasetSchema(
    variables={
        "age": "Patient age, mean-centered and scaled",
        "sex": "Patient sex, mean-centered and scaled",
        "bmi": "Body mass index, mean-centered and scaled",
        "bp": "Average blood pressure, mean-centered and scaled",
        "s1": "Total serum cholesterol (tc)",
        "s2": "Low-density lipoproteins (ldl)",
        "s3": "High-density lipoproteins (hdl)",
        "s4": "Total cholesterol / HDL ratio (tch)",
        "s5": "Log of serum triglycerides level (ltg)",
        "s6": "Blood sugar level (glu)",
        "disease_progression": (
            "Quantitative measure of disease progression one year after baseline (target)"
        ),
    },
    target="disease_progression",
)
schema.validate(data)
print(schema.describe_for_prompt(data))

schema = DatasetSchema( variables={ "age": "Patient age, mean-centered and scaled", "sex": "Patient sex, mean-centered and scaled", "bmi": "Body mass index, mean-centered and scaled", "bp": "Average blood pressure, mean-centered and scaled", "s1": "Total serum cholesterol (tc)", "s2": "Low-density lipoproteins (ldl)", "s3": "High-density lipoproteins (hdl)", "s4": "Total cholesterol / HDL ratio (tch)", "s5": "Log of serum triglycerides level (ltg)", "s6": "Blood sugar level (glu)", "disease_progression": ( "Quantitative measure of disease progression one year after baseline (target)" ), }, target="disease_progression", ) schema.validate(data) print(schema.describe_for_prompt(data))

2. Configure the LLM and run research¶

TabularRegressionResearcher doesn't hardcode a model provider: you instantiate the Strands model object yourself and pass it in, so you can point it at any provider Strands supports. This requires the jcaste05-agentic-ml[research] extra plus the matching strands-agents[<provider>] extra (e.g. strands-agents[gemini] for Gemini, strands-agents[openai] for OpenAI) and a valid API key for that provider — see the installation note above (for Gemini, get one from Google AI Studio).

research(...) drives the agent loop: the agent proposes trials via create_trial, the fixed evaluation pipeline scores each one (K-Fold + save/load roundtrip check), and the result is appended to leaderboard.csv. The agent decides on its own when to stop calling tools, but iterations is a hard cap enforced in Python: create_trial itself refuses to create more than iterations trials, regardless of what the agent asks for.

Following along live: with the INFO logging enabled above you'll see, for each trial, the trial id, status, metrics and runtime (via agentic_ml.core.engine). Switch to DEBUG to also see the exact prompt sent to the agent. A single trial can take a while with small local models: generating a full model_py as a tool call is much harder than a plain chat reply, so give it a few minutes before assuming it's stuck.

In [ ]:

###
# Create your LLM instance:

# from strands.models.openai import OpenAIModel
# llm = OpenAIModel(
#     model_id="gpt-4o",
#     client_args={"api_key": os.environ["OPENAI_API_KEY"]},
# )
# from strands.models.gemini import GeminiModel
# llm = GeminiModel(
#     client_args={
#         "api_key": os.environ["GEMINI_API_KEY"],
#     },
#     model_id="gemini-flash-latest",  # "gemini-flash-latest" - "gemma-4-31b-it"
# )
# from strands.models.ollama import OllamaModel
# llm = OllamaModel(
#     host="http://192.168.1.136:31434",  # Replace with your Ollama server host if different
#     model_id="qwen3:1.7b",
#     keep_alive="5m",
#     temperature=0.1,  # lower = more deterministic, fewer malformed tool calls
#     max_tokens=1536,  # hard cap so a stuck/rambling generation can't run forever
#     options={
#         "num_ctx": 8192,
#     },
#     additional_args={
#         "think": False,
#     },
# )
from strands.models.openai import OpenAIModel
llm = OpenAIModel(
    client_args={
        "api_key": os.environ["GROQ_API_KEY"],
        "base_url": "https://api.groq.com/openai/v1",
    },
    model_id="openai/gpt-oss-120b",  # "qwen/qwen3.8-27b" - "openai/gpt-oss-120b"
    context_window_limit=6_000,  # match the chosen model's real context window
    params={
        "max_tokens": 6_000,  # raise the output cap so a full model_py isn't truncated
        # "temperature": 0.1,  # lower = more deterministic, fewer malformed tool calls
    },
)

###
# SANITY CHECK: Make sure the LLM is working as expected

from strands import Agent
sanity_check = (
    "You only need to remember this number: 902718. What is the number you need to "
    "remember? Be concise with your answer."
)
sanity_agent = Agent(model=llm)
response = sanity_agent(sanity_check)

### # Create your LLM instance: # from strands.models.openai import OpenAIModel # llm = OpenAIModel( # model_id="gpt-4o", # client_args={"api_key": os.environ["OPENAI_API_KEY"]}, # ) # from strands.models.gemini import GeminiModel # llm = GeminiModel( # client_args={ # "api_key": os.environ["GEMINI_API_KEY"], # }, # model_id="gemini-flash-latest", # "gemini-flash-latest" - "gemma-4-31b-it" # ) # from strands.models.ollama import OllamaModel # llm = OllamaModel( # host="http://192.168.1.136:31434", # Replace with your Ollama server host if different # model_id="qwen3:1.7b", # keep_alive="5m", # temperature=0.1, # lower = more deterministic, fewer malformed tool calls # max_tokens=1536, # hard cap so a stuck/rambling generation can't run forever # options={ # "num_ctx": 8192, # }, # additional_args={ # "think": False, # }, # ) from strands.models.openai import OpenAIModel llm = OpenAIModel( client_args={ "api_key": os.environ["GROQ_API_KEY"], "base_url": "https://api.groq.com/openai/v1", }, model_id="openai/gpt-oss-120b", # "qwen/qwen3.8-27b" - "openai/gpt-oss-120b" context_window_limit=6_000, # match the chosen model's real context window params={ "max_tokens": 6_000, # raise the output cap so a full model_py isn't truncated # "temperature": 0.1, # lower = more deterministic, fewer malformed tool calls }, ) ### # SANITY CHECK: Make sure the LLM is working as expected from strands import Agent sanity_check = ( "You only need to remember this number: 902718. What is the number you need to " "remember? Be concise with your answer." ) sanity_agent = Agent(model=llm) response = sanity_agent(sanity_check)

In [ ]:

###
# Create a TabularRegressionResearcher instance

researcher = TabularRegressionResearcher(model=llm)

### # Create a TabularRegressionResearcher instance researcher = TabularRegressionResearcher(model=llm)

In [ ]:

###
# Execute the research agent loop

# Configuration for small models or free plan rate-limits
agent_config = AgentConfig(
    pin_first=1,
    preserve_recent_messages=1,  # always keep the last 4 messages uncompressed
    summary_ratio=0.15,  # fraction of older messages summarized when compression triggers
    compression_threshold=0.80,  # % of the context window that triggers proactive summarization
    offload_max_result_tokens=500,  # tool results larger than this get truncated
    offload_preview_tokens=100,  # size of the truncated preview kept in place
)

prompt_idea = (
    "Study linear models"
)
researcher.research(
    data,
    schema,
    metrics=["rmse", "mae", "r2"],
    iterations=1,
    ideas=prompt_idea,
    research_dir="./agentic_ml_runs/diabetes_quick_example_2",
    sleep=0.0,
    new_session=True,
    agent_config=agent_config,
)

### # Execute the research agent loop # Configuration for small models or free plan rate-limits agent_config = AgentConfig( pin_first=1, preserve_recent_messages=1, # always keep the last 4 messages uncompressed summary_ratio=0.15, # fraction of older messages summarized when compression triggers compression_threshold=0.80, # % of the context window that triggers proactive summarization offload_max_result_tokens=500, # tool results larger than this get truncated offload_preview_tokens=100, # size of the truncated preview kept in place ) prompt_idea = ( "Study linear models" ) researcher.research( data, schema, metrics=["rmse", "mae", "r2"], iterations=1, ideas=prompt_idea, research_dir="./agentic_ml_runs/diabetes_quick_example_2", sleep=0.0, new_session=True, agent_config=agent_config, )

3. Inspect the leaderboard¶

The research directory is the audit trail: every trial_N/ folder holds the agent-written model.py (+ optional helpers.py), and leaderboard.csv records the metrics for all of them, in the order they were tried.

In [ ]:

leaderboard = researcher.context.leaderboard.read()
leaderboard.sort_values("rmse")

leaderboard = researcher.context.leaderboard.read() leaderboard.sort_values("rmse")

4. Export and fit the production model¶

export_model() binds a light AgenticModel to the best trial's code (by the primary metric, rmse here). fit re-trains it — possibly on different data than was used during research — and save persists both the code and the fitted state.

In [ ]:

X = data.drop(columns=[schema.target])
y = data[schema.target]

model = researcher.export_model()  # best trial by default
model.fit(X, y)
model.save("./artifacts/diabetes_model")

X = data.drop(columns=[schema.target]) y = data[schema.target] model = researcher.export_model() # best trial by default model.fit(X, y) model.save("./artifacts/diabetes_model")

5. Load in "production" (no LLM dependency)¶

AgenticModel.load only needs the base install (jcaste05-agentic-ml, without the research extra): it executes the saved trial's code and restores the fitted state.

In [ ]:

production_model = AgenticModel.load("./artifacts/diabetes_model")
predictions = production_model.predict(X.head())
predictions

production_model = AgenticModel.load("./artifacts/diabetes_model") predictions = production_model.predict(X.head()) predictions

Security note¶

Both research(...) and AgenticModel.load(...) execute Python code written by the agent. Trials run in an isolated subprocess with a timeout, but treat saved artifacts with the same trust level as a pickle: only load artifacts you produced or trust.