feat(models): add support for gemma model via gemini api

Merge https://github.com/google/adk-python/pull/2857

Adds support for invoking Gemma models via the Gemini API endpoint. To support agentic function, callbacks are added which can extract and transform function calls and responses into user and model messages in the history.

This change is intended to allow developers to explore the use of Gemma models for agentic purposes without requiring local deployment of the models. This should ease the burden of experimentation and testing for developers.

A basic "hello world" style agent example is provided to demonstrate proper functioning of Gemma 3 models inside an Agent container, using the dice roll + prime check framework of similar examples for other models.

## Testing

### Testing Plan
- add and run integration and unit tests
- manual run of example `multi_tool_agent` from quickstart using new `Gemma` model
- manual run of `hello_world_gemma` agent

### Automated Test Results:
| Test Command | Results |
|----------------|---------|
| pytest ./tests/unittests | 4386 passed, 2849 warnings in 58.43s |
| pytest ./tests/unittests/models/test_google_llm.py | 100 passed, 4 warnings in 5.83s |
| pytest ./tests/integration/models/test_google_llm.py | 5 passed, 2 warnings in 3.73s |

### Manual Testing

Here is a log of `multi_tool_agent` run with locally-built wheel and using Gemma model.
```
❯ adk run multi_tool_agent
Log setup complete: /var/folders/bg/_133c0ds2kb7cn699cpmmh_h0061bp/T/agents_log/agent.20250904_152617.log
To access latest log: tail -F /var/folders/bg/_133c0ds2kb7cn699cpmmh_h0061bp/T/agents_log/agent.latest.log
/Users/<redacted>/venvs/adk-quickstart/lib/python3.11/site-packages/google/adk/cli/cli.py:143: UserWarning: [EXPERIMENTAL] InMemoryCredentialService: This feature is experimental and may change or be removed in future versions without notice. It may introduce breaking changes at any time.
  credential_service = InMemoryCredentialService()
/Users/<redacted>/venvs/adk-quickstart/lib/python3.11/site-packages/google/adk/auth/credential_service/in_memory_credential_service.py:33: UserWarning: [EXPERIMENTAL] BaseCredentialService: This feature is experimental and may change or be removed in future versions without notice. It may introduce breaking changes at any time.
  super().__init__()
Running agent weather_time_agent, type exit to exit.
[user]: what's the weather like today?
[weather_time_agent]: Which city are you asking about?

[user]: new york
[weather_time_agent]: OK. The weather in New York is sunny with a temperature of 25 degrees Celsius (77 degrees Fahrenheit).
```

And here is a snippet of a log generated with DEBUG level logging of the `hello_world_gemma` sample. It demonstrates how function calls are extracted and inserted based on Gemma model interactions:

```
...
2025-09-04 15:32:41,708 - DEBUG - google_llm.py:138 -
LLM Request:
-----------------------------------------------------------
System Instruction:
None
-----------------------------------------------------------
Contents:
{"parts":[{"text":"\n      You roll dice and answer questions about the outcome of the dice rolls.\n      You can roll dice of different sizes...\n"}],"role":"user"}
{"parts":[{"text":"Hi, introduce yourself."}],"role":"user"}
{"parts":[{"text":"Hello! I am data_processing_agent, a hello world agent that can roll many-sided dice and check if numbers are prime. I'm ready to assist you with those tasks. Let's begin!\n\n\n\n"}],"role":"model"}
{"parts":[{"text":"Roll a die with 100 sides and check if it is prime"}],"role":"user"}
{"parts":[{"text":"{\"args\":{\"sides\":100},\"name\":\"roll_die\"}"}],"role":"model"}
{"parts":[{"text":"Invoking tool `roll_die` produced: `{\"result\": 82}`."}],"role":"user"}
{"parts":[{"text":"{\"args\":{\"nums\":[82]},\"name\":\"check_prime\"}"}],"role":"model"}
{"parts":[{"text":"Invoking tool `check_prime` produced: `{\"result\": \"No prime numbers found.\"}`."}],"role":"user"}
{"parts":[{"text":"The die roll was 82, and it is not a prime number.\n\n\n\n"}],"role":"model"}
{"parts":[{"text":"Roll it again."}],"role":"user"}
-----------------------------------------------------------
Functions:

-----------------------------------------------------------

2025-09-04 15:32:41,708 - INFO - models.py:8165 - AFC is enabled with max remote calls: 10.
2025-09-04 15:32:42,693 - INFO - google_llm.py:180 - Response received from the model.
2025-09-04 15:32:42,693 - DEBUG - google_llm.py:181 -
LLM Response:
-----------------------------------------------------------
Text:
{"args":{"sides":100},"name":"roll_die"}
-----------------------------------------------------------
...
```
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2857 from douglas-reid:add-gemma-via-api e6d015f6a9ccbcf20ef7a7af8e4bbe1e9a5936b6
PiperOrigin-RevId: 816451001
This commit is contained in:
Douglas Reid
2025-10-07 17:38:35 -07:00
committed by Copybara-Service
parent 84f2f417f7
commit 2b5acb98f5
8 changed files with 1087 additions and 3 deletions
@@ -0,0 +1,16 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from . import agent
@@ -0,0 +1,95 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import random
from google.adk.agents.llm_agent import Agent
from google.adk.models.gemma_llm import Gemma
from google.genai.types import GenerateContentConfig
def roll_die(sides: int) -> int:
"""Roll a die and return the rolled result.
Args:
sides: The integer number of sides the die has.
Returns:
An integer of the result of rolling the die.
"""
return random.randint(1, sides)
async def check_prime(nums: list[int]) -> str:
"""Check if a given list of numbers are prime.
Args:
nums: The list of numbers to check.
Returns:
A str indicating which number is prime.
"""
primes = set()
for number in nums:
number = number
if number <= 1:
continue
is_prime = True
for i in range(2, int(number**0.5) + 1):
if number % i == 0:
is_prime = False
break
if is_prime:
primes.add(number)
return (
"No prime numbers found."
if not primes
else f"{', '.join(str(num) for num in primes)} are prime numbers."
)
root_agent = Agent(
model=Gemma(model="gemma-3-27b-it"),
name="data_processing_agent",
description=(
"hello world agent that can roll many-sided dice and check if numbers"
" are prime."
),
instruction="""
You roll dice and answer questions about the outcome of the dice rolls.
You can roll dice of different sizes.
You can use multiple tools in parallel by calling functions in parallel(in one request and in one round).
It is ok to discuss previous dice roles, and comment on the dice rolls.
When you are asked to roll a die, you must call the roll_die tool with the number of sides. Be sure to pass in an integer. Do not pass in a string.
You should never roll a die on your own.
When checking prime numbers, call the check_prime tool with a list of integers. Be sure to pass in a list of integers. You should never pass in a string.
You should not check prime numbers before calling the tool.
When you are asked to roll a die and check prime numbers, you should always make the following two function calls:
1. You should first call the roll_die tool to get a roll. Wait for the function response before calling the check_prime tool.
2. After the user reports a response from roll_die tool, you should call the check_prime tool with the roll_die result.
2.1 If user asks you to check primes based on previous rolls, make sure you include the previous rolls in the list.
3. When you respond, you must include the roll_die result from step 1.
You should always perform the previous 3 steps when asking for a roll and checking prime numbers.
You should not rely on the previous history on prime results.
""",
tools=[
roll_die,
check_prime,
],
generate_content_config=GenerateContentConfig(
temperature=1.0,
top_p=0.95,
),
)
@@ -0,0 +1,77 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
import logging
import time
import agent
from dotenv import load_dotenv
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
from google.adk.cli.utils import logs
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.sessions.session import Session
from google.genai import types
load_dotenv(override=True)
logs.log_to_tmp_folder(level=logging.INFO)
async def main():
app_name = 'my_gemma_app'
user_id_1 = 'user1'
session_service = InMemorySessionService()
artifact_service = InMemoryArtifactService()
runner = Runner(
app_name=app_name,
agent=agent.root_agent,
artifact_service=artifact_service,
session_service=session_service,
)
session_11 = await session_service.create_session(
app_name=app_name, user_id=user_id_1
)
async def run_prompt(session: Session, new_message: str):
content = types.Content(
role='user', parts=[types.Part.from_text(text=new_message)]
)
print('** User says:', content.model_dump(exclude_none=True))
async for event in runner.run_async(
user_id=user_id_1,
session_id=session.id,
new_message=content,
):
if event.content.parts and event.content.parts[0].text:
print(f'** {event.author}: {event.content.parts[0].text}')
start_time = time.time()
print('Start time:', start_time)
print('------------------------------------')
await run_prompt(session_11, 'Hi, introduce yourself.')
await run_prompt(
session_11, 'Roll a die with 100 sides and check if it is prime'
)
await run_prompt(session_11, 'Roll it again.')
await run_prompt(session_11, 'What numbers did I get?')
end_time = time.time()
print('------------------------------------')
print('End time:', end_time)
print('Total time:', end_time - start_time)
if __name__ == '__main__':
asyncio.run(main())
+4 -2
View File
@@ -15,6 +15,7 @@
"""Defines the interface to support a model."""
from .base_llm import BaseLlm
from .gemma_llm import Gemma
from .google_llm import Gemini
from .llm_request import LlmRequest
from .llm_response import LlmResponse
@@ -23,9 +24,10 @@ from .registry import LLMRegistry
__all__ = [
'BaseLlm',
'Gemini',
'Gemma',
'LLMRegistry',
]
for regex in Gemini.supported_models():
LLMRegistry.register(Gemini)
LLMRegistry.register(Gemini)
LLMRegistry.register(Gemma)
+331
View File
@@ -0,0 +1,331 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from functools import cached_property
import json
import logging
import re
from typing import Any
from typing import AsyncGenerator
from google.adk.models.google_llm import Gemini
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.utils.variant_utils import GoogleLLMVariant
from google.genai import types
from google.genai.types import Content
from google.genai.types import FunctionDeclaration
from google.genai.types import Part
from pydantic import AliasChoices
from pydantic import BaseModel
from pydantic import Field
from pydantic import ValidationError
from typing_extensions import override
logger = logging.getLogger('google_adk.' + __name__)
class GemmaFunctionCallModel(BaseModel):
"""Flexible Pydantic model for parsing inline Gemma function call responses."""
name: str = Field(validation_alias=AliasChoices('name', 'function'))
parameters: dict[str, Any] = Field(
validation_alias=AliasChoices('parameters', 'args')
)
class Gemma(Gemini):
"""Integration for Gemma models exposed via the Gemini API.
Only Gemma 3 models are supported at this time. For agentic use cases,
use of gemma-3-27b-it and gemma-3-12b-it are strongly recommended.
For full documentation, see: https://ai.google.dev/gemma/docs/core/
NOTE: Gemma does **NOT** support system instructions. Any system instructions
will be replaced with an initial *user* prompt in the LLM request. If system
instructions change over the course of agent execution, the initial content
**SHOULD** be replaced. Special care is warranted here.
See: https://ai.google.dev/gemma/docs/core/prompt-structure#system-instructions
NOTE: Gemma's function calling support is limited. It does not have full access to the
same built-in tools as Gemini. It also does not have special API support for tools and
functions. Rather, tools must be passed in via a `user` prompt, and extracted from model
responses based on approximate shape.
NOTE: Vertex AI API support for Gemma is not currently included. This **ONLY** supports
usage via the Gemini API.
"""
model: str = (
'gemma-3-27b-it' # Others: [gemma-3-1b-it, gemma-3-4b-it, gemma-3-12b-it]
)
@classmethod
@override
def supported_models(cls) -> list[str]:
"""Provides the list of supported models.
Returns:
A list of supported models.
"""
return [
r'gemma-3.*',
]
@cached_property
def _api_backend(self) -> GoogleLLMVariant:
return GoogleLLMVariant.GEMINI_API
def _move_function_calls_into_system_instruction(
self, llm_request: LlmRequest
):
if llm_request.model is None or not llm_request.model.startswith('gemma-3'):
return
# Iterate through the existing contents to find and convert function calls and responses
# from text parts, as Gemma models don't directly support function calling.
new_contents: list[Content] = []
for content_item in llm_request.contents:
(
new_parts_for_content,
has_function_response_part,
has_function_call_part,
) = _convert_content_parts_for_gemma(content_item)
if has_function_response_part:
if new_parts_for_content:
new_contents.append(Content(role='user', parts=new_parts_for_content))
elif has_function_call_part:
if new_parts_for_content:
new_contents.append(
Content(role='model', parts=new_parts_for_content)
)
else:
new_contents.append(content_item)
llm_request.contents = new_contents
if not llm_request.config.tools:
return
all_function_declarations: list[FunctionDeclaration] = []
for tool_item in llm_request.config.tools:
if isinstance(tool_item, types.Tool) and tool_item.function_declarations:
all_function_declarations.extend(tool_item.function_declarations)
if all_function_declarations:
system_instruction = _build_gemma_function_system_instruction(
all_function_declarations
)
llm_request.append_instructions([system_instruction])
llm_request.config.tools = []
def _extract_function_calls_from_response(self, llm_response: LlmResponse):
if llm_response.partial or (llm_response.turn_complete is True):
return
if not llm_response.content:
return
if not llm_response.content.parts:
return
if len(llm_response.content.parts) > 1:
return
response_text = llm_response.content.parts[0].text
if not response_text:
return
try:
json_candidate = None
markdown_code_block_pattern = re.compile(
r'```(?:(json|tool_code))?\s*(.*?)\s*```', re.DOTALL
)
block_match = markdown_code_block_pattern.search(response_text)
if block_match:
json_candidate = block_match.group(2).strip()
else:
found, json_text = _get_last_valid_json_substring(response_text)
if found:
json_candidate = json_text
if not json_candidate:
return
function_call_parsed = GemmaFunctionCallModel.model_validate_json(
json_candidate
)
function_call = types.FunctionCall(
name=function_call_parsed.name,
args=function_call_parsed.parameters,
)
function_call_part = Part(function_call=function_call)
llm_response.content.parts = [function_call_part]
except (json.JSONDecodeError, ValidationError) as e:
logger.debug(
f'Error attempting to parse JSON into function call. Leaving as text'
f' response. %s',
e,
)
except Exception as e:
logger.warning('Error processing Gemma function call response: %s', e)
@override
async def _preprocess_request(self, llm_request: LlmRequest) -> None:
self._move_function_calls_into_system_instruction(llm_request=llm_request)
if system_instruction := llm_request.config.system_instruction:
contents = llm_request.contents
instruction_content = Content(
role='user', parts=[Part.from_text(text=system_instruction)]
)
# NOTE: if history is preserved, we must include the system instructions ONLY once at the beginning
# of any chain of contents.
if contents:
if contents[0] != instruction_content:
# only prepend if it hasn't already been done
llm_request.contents = [instruction_content] + contents
llm_request.config.system_instruction = None
return await super()._preprocess_request(llm_request)
@override
async def generate_content_async(
self, llm_request: LlmRequest, stream: bool = False
) -> AsyncGenerator[LlmResponse, None]:
"""Sends a request to the Gemma model.
Args:
llm_request: LlmRequest, the request to send to the Gemini model.
stream: bool = False, whether to do streaming call.
Yields:
LlmResponse: The model response.
"""
# print(f'{llm_request=}')
assert llm_request.model.startswith('gemma-'), (
f'Requesting a non-Gemma model ({llm_request.model}) with the Gemma LLM'
' is not supported.'
)
async for response in super().generate_content_async(llm_request, stream):
self._extract_function_calls_from_response(response)
yield response
def _convert_content_parts_for_gemma(
content_item: Content,
) -> tuple[list[Part], bool, bool]:
"""Converts function call/response parts within a content item to text parts.
Args:
content_item: The original Content item.
Returns:
A tuple containing:
- A list of new Part objects with function calls/responses converted to text.
- A boolean indicating if any function response parts were found.
- A boolean indicating if any function call parts were found.
"""
new_parts: list[Part] = []
has_function_response_part = False
has_function_call_part = False
for part in content_item.parts:
if func_response := part.function_response:
has_function_response_part = True
response_text = (
f'Invoking tool `{func_response.name}` produced:'
f' `{json.dumps(func_response.response)}`.'
)
new_parts.append(Part.from_text(text=response_text))
elif func_call := part.function_call:
has_function_call_part = True
new_parts.append(
Part.from_text(text=func_call.model_dump_json(exclude_none=True))
)
else:
new_parts.append(part)
return new_parts, has_function_response_part, has_function_call_part
def _build_gemma_function_system_instruction(
function_declarations: list[FunctionDeclaration],
) -> str:
"""Constructs the system instruction string for Gemma function calling."""
if not function_declarations:
return ''
system_instruction_prefix = 'You have access to the following functions:\n['
instruction_parts = []
for func in function_declarations:
instruction_parts.append(func.model_dump_json(exclude_none=True))
separator = ',\n'
system_instruction = (
f'{system_instruction_prefix}{separator.join(instruction_parts)}\n]\n'
)
system_instruction += (
'When you call a function, you MUST respond in the format of: '
"""{"name": function name, "parameters": dictionary of argument name and its value}\n"""
'When you call a function, you MUST NOT include any other text in the'
' response.\n'
)
return system_instruction
def _get_last_valid_json_substring(text: str) -> tuple[bool, str | None]:
"""Attempts to find and return the last valid JSON object in a string.
This function is designed to extract JSON that might be embedded in a larger
text, potentially with introductory or concluding remarks. It will always chose
the last block of valid json found within the supplied text (if it exists).
Args:
text: The input string to search for JSON objects.
Returns:
A tuple:
- bool: True if a valid JSON substring was found, False otherwise.
- str | None: The last valid JSON substring found, or None if none was
found.
"""
decoder = json.JSONDecoder()
last_json_str = None
start_pos = 0
while start_pos < len(text):
try:
first_brace_index = text.index('{', start_pos)
_, end_index = decoder.raw_decode(text[first_brace_index:])
last_json_str = text[first_brace_index : first_brace_index + end_index]
start_pos = first_brace_index + end_index
except json.JSONDecodeError:
start_pos = first_brace_index + 1
except ValueError:
break
if last_json_str:
return True, last_json_str
return False, None
+1 -1
View File
@@ -114,6 +114,6 @@ def pytest_generate_tests(metafunc: Metafunc):
def _is_explicitly_marked(mark_name: str, metafunc: Metafunc) -> bool:
if hasattr(metafunc.function, 'pytestmark'):
for mark in metafunc.function.pytestmark:
if mark.name == 'parametrize' and mark.args[0] == mark_name:
if mark.name == 'parametrize' and mark_name in mark.args[0]:
return True
return False
@@ -0,0 +1,57 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from google.adk.models.gemma_llm import Gemma
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.genai import types
from google.genai.types import Content
from google.genai.types import Part
import pytest
DEFAULT_GEMMA_MODEL = "gemma-3-1b-it"
@pytest.fixture
def gemma_llm():
return Gemma(model=DEFAULT_GEMMA_MODEL)
@pytest.fixture
def gemma_request():
return LlmRequest(
model=DEFAULT_GEMMA_MODEL,
contents=[
Content(
role="user",
parts=[
Part.from_text(text="You are a helpful assistant."),
Part.from_text(text="Hello!"),
],
)
],
config=types.GenerateContentConfig(
temperature=0.1,
response_modalities=[types.Modality.TEXT],
system_instruction="Talk like a pirate.",
),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"])
async def test_generate_content_async(gemma_llm, gemma_request):
async for response in gemma_llm.generate_content_async(gemma_request):
assert isinstance(response, LlmResponse)
assert response.content.parts[0].text
File diff suppressed because it is too large Load Diff