Building a Reasoning Agent with Gemini and Vertex AI

Sebastian Gomez

Jun 24, 2026

Updated: Jun 25, 2026

Building a Reasoning Agent with Gemini and Vertex AI

Many of the most powerful reasoning language models today are proprietary and closed source. This limits access for developers and slows research into these crucial capabilities. But don't worry! In this article, we'll build our own reasoning agent using Google's Gemini 2.5 Flash, fine tuning it with a publicly available dataset. Our goal is to replicate the chain of thought reasoning often found in closed models. With this method, the model explains its thought process step by step before giving the final answer, making this powerful technique far more accessible to everyone.

Why does explicit reasoning matter?

Having a model "think out loud" isn't just a neat trick. It has concrete benefits:

More accurate results: By following step by step logic, the model is less likely to make mistakes or jump to wrong conclusions.
Transparency: It lets us understand how the model reached an answer, which is vital for trust and debugging. We get to see its "thought process."
Better generalization: A model that learns to reason can apply that skill to new, unfamiliar problems more effectively.

We'll achieve this using Vertex AI on Google Cloud Platform (GCP), walking through the full workflow from data preparation to deploying the fine tuned model.

Creating a Storage Bucket

Before we can train our model, we need a safe and accessible place to store our training data. We'll use Google Cloud Storage (GCS) for this. Think of GCS as a giant, secure hard drive in the cloud where you can store files of any kind.

A "bucket" in GCS is essentially a top level container for your data. It's similar to a folder on your computer, but at a much larger scale and optimized for the cloud. Why do we need it? Because Vertex AI (the platform where we'll train the model) needs to access the training data from a centralized, efficient location. A GCS bucket provides a reliable, scalable way to make our dataset available to the Vertex AI training process, and it integrates seamlessly with other GCP services.

Steps to create a bucket:

Open the Google Cloud Console: Make sure you're signed in with the Google account you want to use for this project and that you've selected the correct project in the top dropdown.
Go to Cloud Storage: In the left navigation menu (sometimes called the "hamburger" menu), find Cloud Storage and click Buckets. You may need to scroll down or use the search bar.
Click + CREATE: A form will open to configure your new bucket.
Pick a globally unique name: This name must be unique across all of Google Cloud, not just your project! Something like bucket-dataset-reasoning-my-project could work (replace my-project with something unique to you).
Leave the other settings at their defaults: For this tutorial, the defaults are usually enough (location, storage class, etc.).
Click "CREATE". And that's it! You now have your bucket.

IAM Setup (Identity and Access Management)

To let our code (running, say, in a Colab notebook or a Vertex AI instance) interact securely with Vertex AI and Cloud Storage, we need the right permissions.

Security recommendation (updated): Instead of downloading a long lived service account JSON key, a pattern Google now discourages, we'll use Application Default Credentials (ADC). It's safer, leaves no secrets on your disk or in your notebook, and works just as well in Colab, Vertex AI Workbench, or your local machine. If you deploy on GCP infrastructure, prefer Workload Identity.

Concept: the service account. A service account is like a special Google identity used by applications or services rather than individual users. This is crucial for security: instead of putting your personal credentials in the code (never do that!), we give our code its own identity with specific, limited permissions.

Steps to grant the necessary permissions:

▶ Go to IAM: Open the Google Cloud Console, select your project, and go to IAM & Admin > IAM.
Grant the roles to the identity you'll use (your user for ADC, or a dedicated service account like finetune-gemini-agent if you're working on infrastructure):

- Vertex AI User (`roles/aiplatform.user`): To create and manage training and inference jobs in Vertex AI. - Storage Object Admin (`roles/storage.objectAdmin`): To read the dataset from the bucket and write results if needed.

Save the changes.

With ADC, you don't need to download or handle any key file: you authenticate once with gcloud auth application-default login (we'll cover this in the next step) and the SDK picks up your credentials automatically.

Installing & Importing Libraries and Configuration

Before we start working with the dataset and Vertex AI, we need to install and import the required Python libraries. Let's get to it!

Migration note: This tutorial now uses the Google Gen AI SDK (google-genai), the official, current library for Gemini models. It replaces the old vertexai.generative_models / vertexai.tuning, which is deprecated.

# Install the required libraries
!pip install datasets --quiet
!pip install google-cloud-storage --quiet
!pip install --upgrade google-genai --quiet

# Authenticate with Application Default Credentials (ADC).
# This may open a login window in Colab/notebooks.
!gcloud auth application-default login

# Import the libraries
import datasets
import json
import time  # For pauses and monitoring
import os
from copy import deepcopy  # To copy complex data structures

from google.cloud import storage
from google import genai
from google.genai import types

Now we'll define a few constants we'll use throughout the notebook. With ADC, you no longer need to upload or reference any key file, just adjust your project ID, region, and bucket name.

# --- CONFIGURATION ---
# !!! CHANGE THESE VALUES TO YOUR OWN !!!
PROJECT_ID = "your-gcp-project-id"        #  Replace with YOUR GCP project ID
LOCATION = "us-central1"                   # A region that supports Vertex AI and Gemini
BUCKET_NAME = "reasoning_dataset_bucket"   #  Replace with YOUR bucket name

# Base Gemini model to fine-tune.
#  Check the Vertex AI docs for the current tunable Flash model.
MODEL_BASE = "gemini-2.5-flash"  # We use the Flash version for efficiency and cost

# Name we'll give our fine-tuned model in Vertex AI
FINETUNED_MODEL_NAME = "gemini_flash_reasoner_tuned_v1"
# --- END CONFIGURATION ---

# Gen AI SDK client pointing at Vertex AI.
# It picks up ADC credentials automatically; no JSON key needed.
client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

print(" Configuration ready and Gen AI client initialized.")

Downloading the Dataset and Converting the Format

Now we'll download the reasoning dataset and transform it into the specific JSON Lines (JSONL) format Vertex AI requires for fine tuning. JSONL is simply a text file where each line is a valid JSON object. This format is efficient for processing large amounts of data line by line.

Here we'll use a Hugging Face dataset called `KingNish/reasoning-base-20k`.

# Download the dataset from Hugging Face
dataset_hf = datasets.load_dataset("KingNish/reasoning-base-20k")
print("Dataset downloaded:")
print(dataset_hf)

This dataset is designed specifically to train models to reason step by step about problems, much like a human would. It contains a variety of problems (science, code, math, general logic) along with a detailed chain of thought (CoT) that leads to the correct answer. Crucially for our goal, the dataset already explicitly separates the "thinking" stage (reasoning) from the final "answer" (assistant). Perfect for what we want!

To fine tune Gemini on Vertex AI, our data needs to follow a specific JSON structure for each training example. The expected structure looks like this (pseudo JSON; the // comments are explanatory only):

// Example structure for ONE line in the JSONL file
{
  "systemInstruction": {            // Optional: general instruction for the model
    "role": "system",
    "parts": [{ "text": "You are an expert reasoning model..." }]
  },
  "contents": [                     // List of conversation turns
    {
      "role": "user",               // What the user asks
      "parts": [{ "text": "{dataset_question}" }]
    },
    {
      "role": "model",              // What the model should answer
      "parts": [{ "text": "<think>{reasoning}</think>\n<answer>{answer}</answer>" }]
    }
  ]
}

Verify the schema: The exact JSONL format for SFT (field names like contents / messages and systemInstruction) depends on the version of the Vertex AI tuning API. Confirm the current schema in the docs before launching the job.

Let's define a standard system instruction and a template to format our data:

# System instruction that will guide our fine-tuned model
instruction = """
You are an exceptional language model specialized in logical reasoning.
You must always think step by step before providing a final answer.
Enclose your thought process inside <think> </think> tags.
After thinking, summarize your reasoning and present the final, clear and concise
answer inside <answer> </answer> tags.
"""

# Base template for each JSONL example
template_format = {
    "contents": [
        {"role": "user", "parts": [{"text": "{question}"}]},   # Placeholder for the question
        {"role": "model", "parts": [{"text": "{answer}"}]},    # Formatted answer
    ]
}

# Function to preprocess each example from the Hugging Face dataset
def preproc_train_data(example):
    now_input = example["user"]
    now_reasoning = example["reasoning"]
    now_assistant_answer = example["assistant"]

    # Build the formatted output for the model (thinking + answer)
    now_output = (
        f"<think>\n{now_reasoning}\n</think>\n"
        f"<answer>\n{now_assistant_answer}\n</answer>"
    )

    template_now = deepcopy(template_format)
    template_now["contents"][0]["parts"][0]["text"] = now_input
    template_now["contents"][1]["parts"][0]["text"] = now_output

    # Each JSONL line must be a JSON string
    return {"data_jsonl": json.dumps(template_now)}

# Apply preprocessing to the whole dataset ('map' uses multiple processes if available)
dataset_processed = dataset_hf.map(preproc_train_data, num_proc=os.cpu_count())

# We focus on the 'train' split for this example
training_data_jsonl_strings = dataset_processed["train"]["data_jsonl"]

print(f" Preprocessed {len(training_data_jsonl_strings)} training examples.")
print("Example of one formatted JSONL line:")
print(training_data_jsonl_strings[0])

Uploading the Dataset to the Bucket

Now that we've prepared our dataset in the correct JSONL format, we need to upload it to the Google Cloud Storage bucket we created earlier. This makes the data accessible to Vertex AI during fine tuning.

We'll use the google-cloud-storage library. Since we're already authenticated with ADC, the GCS client picks up the credentials automatically, no key files:

def upload_jsonl_to_gcs(bucket_name, destination_blob_name, json_string_list):
    """Uploads a list of JSON strings to GCS as a JSON Lines (.jsonl) file.

    Args:
        bucket_name (str): Name of the GCS bucket.
        destination_blob_name (str): Destination file name (must end in .jsonl).
        json_string_list (list): A list where each item is a valid JSON string.
    """
    try:
        # GCS client using ADC (no JSON key)
        storage_client = storage.Client()
        bucket = storage_client.bucket(bucket_name)
        blob = bucket.blob(destination_blob_name)

        # Join the JSON lines with newlines to form the JSONL content
        jsonl_content = "\n".join(json_string_list)

        # 'application/jsonl' is not a registered MIME type; we use a conventional one.
        blob.upload_from_string(jsonl_content, content_type="application/x-ndjson")

        gcs_uri = f"gs://{bucket_name}/{destination_blob_name}"
        print(f" File '{destination_blob_name}' uploaded to {gcs_uri}")
        return gcs_uri

    except Exception as e:
        print(f" Error uploading the JSON Lines file: {e}")
        return None


TRAINING_DATA_GCS_FILENAME = "training_dataset_reasoning_gemini.jsonl"

training_data_gcs_uri = upload_jsonl_to_gcs(
    BUCKET_NAME,
    TRAINING_DATA_GCS_FILENAME,
    training_data_jsonl_strings,
)

if training_data_gcs_uri:
    print(f"Training dataset GCS URI: {training_data_gcs_uri}")
else:
    print(" There was a problem uploading the dataset. Check the error messages.")

You can verify the upload by going to the Cloud Storage section in the GCP Console and navigating to your bucket. You should see your .jsonl file there!

Fine tuning Gemini

Now for the heart of the process! Fine tuning the base Gemini model (our MODEL_BASE) with our prepared dataset. With the Gen AI SDK, we launch a supervised fine tuning (SFT) job through client.tunings.

Verify the tuning API: The exact method names (client.tunings.tune), parameters (epoch_count, adapter_size, learning_rate_multiplier), and how the dataset is passed evolve over time. Confirm the current signature in the Vertex AI tuning docs before running.

print(f" Starting the fine-tuning job for base model: {MODEL_BASE}")
print(f"Using training dataset at: {training_data_gcs_uri}")

if not training_data_gcs_uri:
    raise ValueError("The GCS dataset URI is not defined. Cannot start training.")

tuning_job = client.tunings.tune(
    base_model=MODEL_BASE,
    training_dataset=types.TuningDataset(gcs_uri=training_data_gcs_uri),
    config=types.CreateTuningJobConfig(
        tuned_model_display_name=FINETUNED_MODEL_NAME,
        epoch_count=2,                 # 2-4 is a good starting point
        adapter_size="ADAPTER_SIZE_FOUR",  # LoRA: larger = more capacity (and cost)
        learning_rate_multiplier=1.0,
        # validation_dataset=types.TuningValidationDataset(gcs_uri="gs://.../validation.jsonl"),
    ),
)

print(" Fine-tuning job submitted successfully to Vertex AI.")
print(f"Job name: {tuning_job.name}")
print("You can monitor progress in the 'Tuning' section of Vertex AI in the GCP console.")

Monitoring the Job:

Fine tuning can take a while ⏳ (from under an hour to several hours, depending on the dataset size, base model, and parameters like epoch_count and adapter_size). The following code checks the job status every 60 seconds until it finishes:

import time

print(" Waiting for the fine-tuning job to finish... (This may take a while)")

# States that mean the job is still in progress
running_states = {
    "JOB_STATE_PENDING",
    "JOB_STATE_RUNNING",
    "JOB_STATE_QUEUED",
}

while tuning_job.state in running_states:
    time.sleep(60)
    tuning_job = client.tunings.get(name=tuning_job.name)  # Refresh the state
    print(f"Current job state ({time.strftime('%H:%M:%S')}): {tuning_job.state}")

print("\n The fine-tuning job has finished!")

if tuning_job.state == "JOB_STATE_SUCCEEDED":
    print(" The job completed successfully.")
    # The fine-tuned model is ready for inference via its endpoint
    tuned_endpoint = tuning_job.tuned_model.endpoint
    print(f"Fine-tuned model endpoint: {tuned_endpoint}")
else:
    print(f" The job failed or was cancelled. Final state: {tuning_job.state}")
    print(f"Error (if available): {getattr(tuning_job, 'error', None)}")

Good news: With the Gen AI SDK, a successful tuning job registers the fine tuned model and exposes an endpoint ready for inference (tuning_job.tuned_model.endpoint), without the manual deployment step the old flow required. Confirm this behavior for your model and region in the docs.

Inference (Using the Fine tuned Model)

Congratulations! If you made it this far, you've fine tuned your own Gemini model. Now we can send it prompts (questions or instructions) and get back responses that show off its new reasoning abilities, following the <think>...</think><answer>...</answer> format we taught it.

We'll point at the fine tuned model's endpoint and reinforce the behavior with the same system instruction. We also set some generation parameters (we use a low temperature to favor deterministic reasoning):

# The fine-tuned model's endpoint. In a new session, you can get it like this:
#   tuning_job = client.tunings.get(name="projects/.../tuningJobs/...")
#   tuned_endpoint = tuning_job.tuned_model.endpoint
TUNED_MODEL = tuned_endpoint  # from the previous step

system_instruction_inference = """
You are an exceptional language model specialized in logical reasoning.
You must always think step by step before providing a final answer.
Enclose your thought process inside <think> </think> tags.
After thinking, summarize your reasoning and present the final, clear and concise
answer inside <answer> </answer> tags.
"""

def generate_with_tuned_model(prompt_text):
    print(f"\n Sending prompt to the fine-tuned model: '{prompt_text}'")
    try:
        response = client.models.generate_content(
            model=TUNED_MODEL,
            contents=prompt_text,
            config=types.GenerateContentConfig(
                system_instruction=system_instruction_inference,
                temperature=0.3,        # Low = more deterministic, ideal for reasoning
                top_p=0.95,
                max_output_tokens=2048,
            ),
        )
        print(" Response received.")
        return response
    except Exception as e:
        print(f" Error during inference: {e}")
        return None


# --- Let's test the model! ---
question = "What happens if I don't have enough money to pay for GCP and I haven't linked my credit card?"
full_response = generate_with_tuned_model(question)

# Process the response to extract the thinking and answer parts
if full_response and full_response.text:
    full_text = full_response.text

    # Extract the <think>...</think> block
    thinking_part = ""
    if "<think>" in full_text and "</think>" in full_text:
        thinking_part = full_text.split("<think>", 1)[1].split("</think>", 1)[0].strip()

    # Extract the <answer>...</answer> block (with fallbacks)
    if "<answer>" in full_text and "</answer>" in full_text:
        answer_part = full_text.split("<answer>", 1)[1].split("</answer>", 1)[0].strip()
    elif "</think>" in full_text:
        answer_part = full_text.split("</think>", 1)[1].strip()
    else:
        answer_part = full_text.strip()

    print("\n--- Model's Thought Process --- ")
    print(thinking_part or "(No <think> tag found)")

    print("\n--- Model's Final Answer --- ")
    print(answer_part or "(No <answer> tag found)")
else:
    print("No valid response received from the model.")

Conclusion

In this article, we've walked the full path to successfully fine tune a Google Gemini Flash model and turn it into a reasoning agent that can show its chain of thought process.

We started with a common problem: the opacity of many powerful models. Our solution was to take an accessible model (Gemini Flash), a public dataset designed for reasoning (KingNish/reasoning-base-20k), and use Google Cloud Platform tools, specifically Vertex AI and Google Cloud Storage, to:

Prepare and format the data (JSONL).
Store it securely and accessibly (GCS bucket).
Configure the necessary permissions securely (IAM + ADC).
Launch and monitor a supervised fine tuning job (Vertex AI SFT).
Run inference to observe explicit reasoning (the famous <think>...</think><answer>...</answer>).

What might look like a daunting task, building a model that not only answers but explains how it reached the answer, is surprisingly accessible thanks to modern platforms like Google Cloud. The process, while it involves several steps, is fairly straightforward once you break it down.

The Gen AI SDK in particular greatly simplifies infrastructure management and large model training, letting you focus on what matters: teaching your model to think and reason better for your specific needs. Now you have the power to build more transparent, trustworthy models!

That's all, I hope this post is useful to you. If you have questions or want to share what you built, drop it in the comments. Good luck!