Skip to content

Commit 9d6cc80

Browse files
glaforgecopybara-github
authored andcommitted
feat: add support for Gemma models in LlmRegistry
- Registered 'gemma-.*' models to use the Gemini builder. - Added GemmaTest.java to verify valid and invalid model names. - Restored main BUILD file and added GemmaTest target with env setup. Fixes b/499032158 PiperOrigin-RevId: 893918083
1 parent e0083b3 commit 9d6cc80

2 files changed

Lines changed: 52 additions & 0 deletions

File tree

core/src/main/java/com/google/adk/models/LlmRegistry.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.adk.models;
1818

19+
import com.google.common.annotations.VisibleForTesting;
1920
import java.util.Map;
2021
import java.util.concurrent.ConcurrentHashMap;
2122

@@ -38,6 +39,7 @@ public interface LlmFactory {
3839
static {
3940
registerLlm("gemini-.*", modelName -> Gemini.builder().modelName(modelName).build());
4041
registerLlm("apigee/.*", modelName -> ApigeeLlm.builder().modelName(modelName).build());
42+
registerLlm("gemma-.*", modelName -> Gemini.builder().modelName(modelName).build());
4143
}
4244

4345
/**
@@ -50,6 +52,17 @@ public static void registerLlm(String modelNamePattern, LlmFactory factory) {
5052
llmFactories.put(modelNamePattern, factory);
5153
}
5254

55+
/**
56+
* Checks if the given model name matches any of the registered LLM factory patterns.
57+
*
58+
* @param modelName The model name to check.
59+
* @return {@code true} if the model name matches at least one pattern, {@code false} otherwise.
60+
*/
61+
@VisibleForTesting
62+
static boolean matchesAnyPattern(String modelName) {
63+
return llmFactories.keySet().stream().anyMatch(modelName::matches);
64+
}
65+
5366
/**
5467
* Returns an LLM instance for the given model name, using a cached or new factory-created
5568
* instance.
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.models;
18+
19+
import static com.google.common.truth.Truth.assertThat;
20+
21+
import org.junit.Test;
22+
import org.junit.runner.RunWith;
23+
import org.junit.runners.JUnit4;
24+
25+
@RunWith(JUnit4.class)
26+
public class GemmaTest {
27+
28+
@Test
29+
public void getLlm_withValidGemmaModels_succeeds() {
30+
assertThat(LlmRegistry.matchesAnyPattern("gemma-4-26b-a4b-it")).isTrue();
31+
assertThat(LlmRegistry.matchesAnyPattern("gemma-4-31b-it")).isTrue();
32+
}
33+
34+
@Test
35+
public void getLlm_withInvalidGemmaModels_throwsException() {
36+
assertThat(LlmRegistry.matchesAnyPattern("not-a-gemma")).isFalse();
37+
assertThat(LlmRegistry.matchesAnyPattern("gemma")).isFalse();
38+
}
39+
}

0 commit comments

Comments
 (0)