[Python] How to invoke a Gemini Fine tuned (via Vertex AI) model in Google ADK Agent

Problem Statement: I fine-tuned a google's gemini flash 2.0 model via GCP Vertex AI. But I am unable to invoke the endpoint of that fine-tuned model trained on my own data in the below agent.py. I have provided all the code I have so far below:-

This is my agent.py

import os
from google.adk.agents import Agent
from google.adk.tools import google_search
from google.adk.tools.agent_tool import AgentTool
# from .fine_tune_llama3_8b import FineTuneLlama3 # Assuming this is in the same directory

# --- Set API Key ---
os.environ["GOOGLE_API_KEY"] = "<my_api_key>"

# --- Root Agent ---
root_agent = Agent(
name="RootAgent",
model="gemini-2.0-flash-exp",
# model="gemini-2.0-flash-001"
description="AI Agent",
instruction=f"""
You are a helpful AI assistant.
""",
tools=[google_search]
)

# --- Root Agent for the Runner ---
from google.genai import types
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner

# Instantiate constants
APP_NAME = "snack_creations_app"
USER_ID = "12345"
SESSION_ID = "123344"

# Session and Runner
session_service = InMemorySessionService()
session = session_service.create_session(
app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
)
runner = Runner(
agent=root_agent, app_name=APP_NAME, session_service=session_service
)

# Agent Interaction with Session for Context
def call_agent(query):
content = types.Content(role="user", parts=[types.Part(text=query)])
events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content)

final_response = None
for event in events:
if event.is_final_response():
final_response = event.content.parts[0].text
print("Agent Response: ", final_response)
return final_response
return final_response

# Running Conversation
print("Starting the interactive conversation (type 'quit' to exit).")
while True:
user_query = input("You: ")
if user_query.lower() == 'quit':
break
call_agent(user_query)
print("\n")

print("Conversation ended.")


This is my fined tuned model code given by vertex AI post fine-tuning my model

from google import genai
from google.genai import types
import base64

def generate():
client = genai.Client(
vertexai=True,
project="9456734556734",
location="us-central1",
)

model = "projects/9456734556734/locations/us-central1/endpoints/5797762931831894471"
contents = [
types.Content(
role="user",
parts=[
types.Part.from_text(text="What is the maximum temperature used in the modulated DSC protocol ?")
]
)
]
generate_content_config = types.GenerateContentConfig(
temperature = 0.7,
top_p = 0.95,
max_output_tokens = 8192,
response_modalities = ["TEXT"],
speech_config = types.SpeechConfig(
voice_config = types.VoiceConfig(
prebuilt_voice_config = types.PrebuiltVoiceConfig(
voice_name = "zephyr"
)
),
),
safety_settings = [types.SafetySetting(
category="HARM_CATEGORY_HATE_SPEECH",
threshold="OFF"
),types.SafetySetting(
category="HARM_CATEGORY_DANGEROUS_CONTENT",
threshold="OFF"
),types.SafetySetting(
category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
threshold="OFF"
),types.SafetySetting(
category="HARM_CATEGORY_HARASSMENT",
threshold="OFF"
)],
)

for chunk in client.models.generate_content_stream(
model = model,
contents = contents,
config = generate_content_config,
):
print(chunk.text, end="")

generate()


I tried invoking my fine tuned model end point into Agent but it doesn't work. Is there any way I can use my model in agent.py ? I am not sure if that has been implemented by google or not. Any help would be appreciated.

Continue reading...