Skip to content

Commit 55fffb7

Browse files
Jacksunweicopybara-github
Jacksunwei
authored and
copybara-github
committed
feat(config): Adds ComponentRegistry for loading objects in yaml config
PiperOrigin-RevId: 795193080
1 parent 7695e21 commit 55fffb7

File tree

2 files changed

+361
-0
lines changed

2 files changed

+361
-0
lines changed

‎core/src/main/java/com/google/adk/utils/ComponentRegistry.java

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
/*
2+
* Copyright 2025 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.utils;
18+
19+
import com.google.adk.tools.BuiltInCodeExecutionTool;
20+
import com.google.adk.tools.GoogleSearchTool;
21+
import java.util.Map;
22+
import java.util.Optional;
23+
import java.util.concurrent.ConcurrentHashMap;
24+
import org.slf4j.Logger;
25+
import org.slf4j.LoggerFactory;
26+
27+
/**
28+
* A registry for storing and retrieving ADK instances by name.
29+
*
30+
* <p>This class provides a base registry with common ADK components and is designed to be extended
31+
* by users who want to add their own pre-wired entries. The registry is fully thread-safe and
32+
* supports storing any type of object.
33+
*
34+
* <p><strong>Thread Safety:</strong>
35+
*
36+
* <ul>
37+
* <li>All instance methods are thread-safe due to the underlying ConcurrentHashMap
38+
* <li>The singleton instance access is thread-safe using volatile semantics
39+
* <li>The setInstance() method is synchronized to ensure atomic singleton replacement
40+
* </ul>
41+
*
42+
* <p>Base pre-wired entries include:
43+
*
44+
* <ul>
45+
* <li>"google_search" - GoogleSearchTool instance
46+
* <li>"code_execution" - BuiltInCodeExecutionTool instance
47+
* <li>"exit_loop" - ExitLoopTool instance
48+
* </ul>
49+
*
50+
* <p>Example usage:
51+
*
52+
* <pre>{@code
53+
* // Use the singleton instance
54+
* ComponentRegistry registry = ComponentRegistry.getInstance();
55+
* Optional<GoogleSearchTool> searchTool = registry.get("google_search", GoogleSearchTool.class);
56+
*
57+
* // Extend ComponentRegistry to add custom pre-wired entries
58+
* public class MyComponentRegistry extends ComponentRegistry {
59+
* public MyComponentRegistry() {
60+
* super(); // Initialize base pre-wired entries
61+
* register("my_custom_tool", new MyCustomTool());
62+
* register("my_agent", new MyCustomAgent());
63+
* }
64+
* }
65+
*
66+
* // Replace the singleton with custom registry when server starts
67+
* ComponentRegistry.setInstance(new MyComponentRegistry());
68+
* }</pre>
69+
*/
70+
public class ComponentRegistry {
71+
72+
private static final Logger logger = LoggerFactory.getLogger(ComponentRegistry.class);
73+
private static volatile ComponentRegistry instance = new ComponentRegistry();
74+
75+
private final Map<String, Object> registry = new ConcurrentHashMap<>();
76+
77+
public ComponentRegistry() {
78+
initializePreWiredEntries();
79+
}
80+
81+
/** Initializes the registry with base pre-wired ADK instances. */
82+
private void initializePreWiredEntries() {
83+
registry.put("google_search", new GoogleSearchTool());
84+
registry.put("code_execution", new BuiltInCodeExecutionTool());
85+
86+
logger.debug("Initialized base pre-wired entries in ComponentRegistry");
87+
}
88+
89+
/**
90+
* Registers an object with the given name. This can override pre-wired entries.
91+
*
92+
* <p>This method is thread-safe due to the underlying ConcurrentHashMap.
93+
*
94+
* @param name the name to associate with the object
95+
* @param value the object to register (can be an instance, class, function, etc.)
96+
* @throws IllegalArgumentException if name is null or empty, or if value is null
97+
*/
98+
public void register(String name, Object value) {
99+
if (name == null || name.trim().isEmpty()) {
100+
throw new IllegalArgumentException("Name cannot be null or empty");
101+
}
102+
if (value == null) {
103+
throw new IllegalArgumentException("Value cannot be null");
104+
}
105+
106+
Object previous = registry.put(name, value);
107+
if (previous != null) {
108+
logger.info(
109+
"Overriding existing registration for name: {} (was: {}, now: {})",
110+
name,
111+
previous.getClass().getSimpleName(),
112+
value.getClass().getSimpleName());
113+
} else {
114+
logger.debug(
115+
"Registered new object of type {} with name: {}", value.getClass().getSimpleName(), name);
116+
}
117+
}
118+
119+
/**
120+
* Retrieves an object by name and attempts to cast it to the specified type.
121+
*
122+
* @param name the name of the object to retrieve
123+
* @param type the expected type of the object
124+
* @param <T> the type parameter
125+
* @return an Optional containing the object if found and castable to the specified type, or an
126+
* empty Optional otherwise
127+
*/
128+
public <T> Optional<T> get(String name, Class<T> type) {
129+
return get(name)
130+
.filter(
131+
value -> {
132+
if (type.isInstance(value)) {
133+
return true;
134+
} else {
135+
logger.warn(
136+
"Object with name '{}' is of type {} but expected type {}",
137+
name,
138+
value.getClass().getSimpleName(),
139+
type.getSimpleName());
140+
return false;
141+
}
142+
})
143+
.map(type::cast);
144+
}
145+
146+
/**
147+
* Retrieves an object by name without type checking.
148+
*
149+
* @param name the name of the object to retrieve
150+
* @return an Optional containing the object if found, or an empty Optional otherwise
151+
*/
152+
public Optional<Object> get(String name) {
153+
if (name == null || name.trim().isEmpty()) {
154+
return Optional.empty();
155+
}
156+
157+
return Optional.ofNullable(registry.get(name));
158+
}
159+
160+
/**
161+
* Returns the global singleton instance of ComponentRegistry.
162+
*
163+
* @return the singleton ComponentRegistry instance
164+
*/
165+
public static ComponentRegistry getInstance() {
166+
return instance;
167+
}
168+
169+
/**
170+
* Updates the global singleton instance with a new ComponentRegistry. This is useful for
171+
* replacing the default registry with a custom one when the server starts.
172+
*
173+
* <p>This method is thread-safe and ensures that all threads see the updated instance atomically.
174+
*
175+
* @param newInstance the new ComponentRegistry instance to use as the singleton
176+
* @throws IllegalArgumentException if newInstance is null
177+
*/
178+
public static synchronized void setInstance(ComponentRegistry newInstance) {
179+
if (newInstance == null) {
180+
throw new IllegalArgumentException("ComponentRegistry instance cannot be null");
181+
}
182+
instance = newInstance;
183+
logger.info("ComponentRegistry singleton instance updated");
184+
}
185+
}

‎core/src/test/java/com/google/adk/utils/ComponentRegistryTest.java

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
/*
2+
* Copyright 2025 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.utils;
18+
19+
import static com.google.common.truth.Truth.assertThat;
20+
import static org.junit.Assert.assertThrows;
21+
22+
import com.google.adk.tools.BuiltInCodeExecutionTool;
23+
import com.google.adk.tools.GoogleSearchTool;
24+
import java.util.Optional;
25+
import org.junit.Test;
26+
import org.junit.runner.RunWith;
27+
import org.junit.runners.JUnit4;
28+
29+
@RunWith(JUnit4.class)
30+
public final class ComponentRegistryTest {
31+
32+
@Test
33+
public void testPreWiredEntries() {
34+
ComponentRegistry registry = new ComponentRegistry();
35+
36+
Optional<GoogleSearchTool> searchTool = registry.get("google_search", GoogleSearchTool.class);
37+
assertThat(searchTool).isPresent();
38+
39+
Optional<BuiltInCodeExecutionTool> codeTool =
40+
registry.get("code_execution", BuiltInCodeExecutionTool.class);
41+
assertThat(codeTool).isPresent();
42+
}
43+
44+
@Test
45+
public void testRegisterAndGet() {
46+
ComponentRegistry registry = new ComponentRegistry();
47+
String testValue = "test value";
48+
49+
registry.register("test_key", testValue);
50+
51+
Optional<String> result = registry.get("test_key", String.class);
52+
assertThat(result).hasValue(testValue);
53+
}
54+
55+
@Test
56+
public void testGetWithoutType() {
57+
ComponentRegistry registry = new ComponentRegistry();
58+
String testValue = "test value";
59+
60+
registry.register("test_key", testValue);
61+
62+
Optional<Object> result = registry.get("test_key");
63+
assertThat(result).hasValue(testValue);
64+
}
65+
66+
@Test
67+
public void testGetNonExistentKey() {
68+
ComponentRegistry registry = new ComponentRegistry();
69+
70+
Optional<String> result = registry.get("non_existent", String.class);
71+
assertThat(result).isEmpty();
72+
73+
Optional<Object> resultNoType = registry.get("non_existent");
74+
assertThat(resultNoType).isEmpty();
75+
}
76+
77+
@Test
78+
public void testGetWithWrongType() {
79+
ComponentRegistry registry = new ComponentRegistry();
80+
registry.register("test_key", "string value");
81+
82+
Optional<Integer> result = registry.get("test_key", Integer.class);
83+
assertThat(result).isEmpty();
84+
}
85+
86+
@Test
87+
public void testOverridePreWiredEntry() {
88+
ComponentRegistry registry = new ComponentRegistry();
89+
String customSearchTool = "custom search tool";
90+
91+
registry.register("google_search", customSearchTool);
92+
93+
Optional<String> result = registry.get("google_search", String.class);
94+
assertThat(result).hasValue(customSearchTool);
95+
96+
Optional<GoogleSearchTool> originalTool = registry.get("google_search", GoogleSearchTool.class);
97+
assertThat(originalTool).isEmpty();
98+
}
99+
100+
@Test
101+
public void testRegisterWithNullName() {
102+
ComponentRegistry registry = new ComponentRegistry();
103+
104+
IllegalArgumentException thrown =
105+
assertThrows(IllegalArgumentException.class, () -> registry.register(null, "value"));
106+
107+
assertThat(thrown.getMessage()).contains("Name cannot be null or empty");
108+
}
109+
110+
@Test
111+
public void testRegisterWithEmptyName() {
112+
ComponentRegistry registry = new ComponentRegistry();
113+
114+
IllegalArgumentException thrown1 =
115+
assertThrows(IllegalArgumentException.class, () -> registry.register("", "value"));
116+
IllegalArgumentException thrown2 =
117+
assertThrows(IllegalArgumentException.class, () -> registry.register(" ", "value"));
118+
119+
assertThat(thrown1.getMessage()).contains("Name cannot be null or empty");
120+
assertThat(thrown2.getMessage()).contains("Name cannot be null or empty");
121+
}
122+
123+
@Test
124+
public void testGetWithNullName() {
125+
ComponentRegistry registry = new ComponentRegistry();
126+
127+
Optional<String> result = registry.get(null, String.class);
128+
assertThat(result).isEmpty();
129+
130+
Optional<Object> resultNoType = registry.get(null);
131+
assertThat(resultNoType).isEmpty();
132+
}
133+
134+
@Test
135+
public void testGetWithEmptyName() {
136+
ComponentRegistry registry = new ComponentRegistry();
137+
138+
Optional<String> result = registry.get("", String.class);
139+
assertThat(result).isEmpty();
140+
141+
Optional<String> resultWhitespace = registry.get(" ", String.class);
142+
assertThat(resultWhitespace).isEmpty();
143+
}
144+
145+
@Test
146+
public void testRegisterNullValue() {
147+
ComponentRegistry registry = new ComponentRegistry();
148+
149+
IllegalArgumentException thrown =
150+
assertThrows(IllegalArgumentException.class, () -> registry.register("null_test", null));
151+
152+
assertThat(thrown.getMessage()).contains("Value cannot be null");
153+
}
154+
155+
@Test
156+
public void testSubclassExtension() {
157+
class CustomComponentRegistry extends ComponentRegistry {
158+
public CustomComponentRegistry() {
159+
super();
160+
register("custom_tool", "my custom tool");
161+
register("custom_agent", new Object());
162+
}
163+
}
164+
165+
CustomComponentRegistry registry = new CustomComponentRegistry();
166+
167+
Optional<GoogleSearchTool> prewiredTool = registry.get("google_search", GoogleSearchTool.class);
168+
assertThat(prewiredTool).isPresent();
169+
170+
Optional<String> customTool = registry.get("custom_tool", String.class);
171+
assertThat(customTool).hasValue("my custom tool");
172+
173+
Optional<Object> customAgent = registry.get("custom_agent");
174+
assertThat(customAgent).isPresent();
175+
}
176+
}

0 commit comments

Comments
 (0)