Skip to content

Commit c650ca3

Browse files
Mateusz Krawieccopybara-github
Mateusz Krawiec
authored and
copybara-github
committed
fix: declare set_model_response parameters as an object for untyped schemas
PiperOrigin-RevId: 1018280172
1 parent f688c80 commit c650ca3

10 files changed

Lines changed: 780 additions & 74 deletions

File tree

‎contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java‎

Lines changed: 55 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959
import dev.langchain4j.model.chat.request.ChatRequest;
6060
import dev.langchain4j.model.chat.request.ChatRequestParameters;
6161
import dev.langchain4j.model.chat.request.ToolChoice;
62+
import dev.langchain4j.model.chat.request.json.JsonAnyOfSchema;
6263
import dev.langchain4j.model.chat.request.json.JsonArraySchema;
6364
import dev.langchain4j.model.chat.request.json.JsonBooleanSchema;
6465
import dev.langchain4j.model.chat.request.json.JsonIntegerSchema;
@@ -564,7 +565,11 @@ private List<ToolSpecification> toToolSpecifications(LlmRequest llmRequest) {
564565
}
565566

566567
private JsonObjectSchema toParameters(Schema schema) {
567-
if (schema.type().isPresent() && Type.Known.OBJECT.equals(schema.type().get().knownEnum())) {
568+
// An untyped schema with properties describes an object.
569+
if (schema
570+
.type()
571+
.map(t -> t.knownEnum() == Type.Known.OBJECT)
572+
.orElse(schema.properties().isPresent())) {
568573
return JsonObjectSchema.builder()
569574
.addProperties(toProperties(schema))
570575
.required(schema.required().orElse(List.of()))
@@ -582,31 +587,56 @@ private Map<String, JsonSchemaElement> toProperties(Schema schema) {
582587
return result;
583588
}
584589

585-
private JsonSchemaElement toJsonSchemaElement(Schema schema) {
586-
if (schema != null && schema.type().isPresent()) {
587-
Type type = schema.type().get();
588-
return switch (type.knownEnum()) {
589-
case STRING ->
590-
JsonStringSchema.builder().description(schema.description().orElse(null)).build();
591-
case NUMBER ->
592-
JsonNumberSchema.builder().description(schema.description().orElse(null)).build();
593-
case INTEGER ->
594-
JsonIntegerSchema.builder().description(schema.description().orElse(null)).build();
595-
case BOOLEAN ->
596-
JsonBooleanSchema.builder().description(schema.description().orElse(null)).build();
597-
case ARRAY ->
598-
JsonArraySchema.builder()
599-
.description(schema.description().orElse(null))
600-
.items(toJsonSchemaElement(schema.items().orElseThrow()))
601-
.build();
602-
case OBJECT -> toParameters(schema);
603-
default ->
604-
throw new UnsupportedFeatureException(
605-
"LangChain4jLlm does not support schema of type: " + type);
606-
};
607-
} else {
608-
throw new IllegalArgumentException("Schema type cannot be null or absent");
590+
private JsonSchemaElement toJsonSchemaElement(@Nullable Schema schema) {
591+
if (schema == null) {
592+
return anyValue(null);
593+
}
594+
// An untyped schema with properties is an object; without them, it accepts any value.
595+
if (schema.type().isEmpty()) {
596+
return schema.properties().isPresent()
597+
? toParameters(schema)
598+
: anyValue(schema.description().orElse(null));
609599
}
600+
Type type = schema.type().get();
601+
return switch (type.knownEnum()) {
602+
case STRING ->
603+
JsonStringSchema.builder().description(schema.description().orElse(null)).build();
604+
case NUMBER ->
605+
JsonNumberSchema.builder().description(schema.description().orElse(null)).build();
606+
case INTEGER ->
607+
JsonIntegerSchema.builder().description(schema.description().orElse(null)).build();
608+
case BOOLEAN ->
609+
JsonBooleanSchema.builder().description(schema.description().orElse(null)).build();
610+
case ARRAY ->
611+
JsonArraySchema.builder()
612+
.description(schema.description().orElse(null))
613+
.items(schema.items().map(this::toJsonSchemaElement).orElseGet(() -> anyValue(null)))
614+
.build();
615+
case OBJECT -> toParameters(schema);
616+
default ->
617+
throw new UnsupportedFeatureException(
618+
"LangChain4jLlm does not support schema of type: " + type);
619+
};
620+
}
621+
622+
/**
623+
* Returns an element that accepts a string, number, boolean, object, or an array of those, for a
624+
* schema without a type. Not a {@code JsonRawSchema}, which the Gemini provider rejects. The
625+
* array branch has items, which Gemini requires.
626+
*/
627+
private static JsonSchemaElement anyValue(@Nullable String description) {
628+
List<JsonSchemaElement> nonArrays =
629+
List.of(
630+
new JsonStringSchema(),
631+
new JsonNumberSchema(),
632+
new JsonBooleanSchema(),
633+
JsonObjectSchema.builder().build());
634+
List<JsonSchemaElement> anyOf = new ArrayList<>(nonArrays);
635+
anyOf.add(
636+
JsonArraySchema.builder()
637+
.items(JsonAnyOfSchema.builder().anyOf(nonArrays).build())
638+
.build());
639+
return JsonAnyOfSchema.builder().description(description).anyOf(anyOf).build();
610640
}
611641

612642
private LlmResponse toLlmResponse(ChatResponse chatResponse, ChatRequest chatRequest) {

‎contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,13 @@
1515
*/
1616
package com.google.adk.models.langchain4j;
1717

18+
import static java.nio.charset.StandardCharsets.UTF_8;
1819
import static org.assertj.core.api.Assertions.*;
1920
import static org.junit.jupiter.api.Assertions.assertThrows;
2021
import static org.mockito.ArgumentMatchers.any;
2122
import static org.mockito.Mockito.*;
2223

24+
import com.fasterxml.jackson.databind.JsonNode;
2325
import com.fasterxml.jackson.databind.ObjectMapper;
2426
import com.google.adk.agents.LlmAgent;
2527
import com.google.adk.agents.RunConfig;
@@ -28,7 +30,9 @@
2830
import com.google.adk.runner.InMemoryRunner;
2931
import com.google.adk.sessions.Session;
3032
import com.google.adk.tools.FunctionTool;
33+
import com.google.adk.tools.SetModelResponseTool;
3134
import com.google.genai.types.*;
35+
import com.sun.net.httpserver.HttpServer;
3236
import dev.langchain4j.agent.tool.ToolExecutionRequest;
3337
import dev.langchain4j.data.message.AiMessage;
3438
import dev.langchain4j.data.message.UserMessage;
@@ -37,17 +41,25 @@
3741
import dev.langchain4j.model.chat.StreamingChatModel;
3842
import dev.langchain4j.model.chat.request.ChatRequest;
3943
import dev.langchain4j.model.chat.request.ChatRequestParameters;
44+
import dev.langchain4j.model.chat.request.json.JsonAnyOfSchema;
45+
import dev.langchain4j.model.chat.request.json.JsonArraySchema;
4046
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
47+
import dev.langchain4j.model.chat.request.json.JsonSchemaElement;
4148
import dev.langchain4j.model.chat.request.json.JsonStringSchema;
4249
import dev.langchain4j.model.chat.response.ChatResponse;
4350
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
51+
import dev.langchain4j.model.googleai.GoogleAiGeminiChatModel;
4452
import dev.langchain4j.model.output.TokenUsage;
4553
import io.reactivex.rxjava3.core.Flowable;
4654
import io.reactivex.rxjava3.subscribers.TestSubscriber;
55+
import java.io.OutputStream;
56+
import java.net.InetAddress;
57+
import java.net.InetSocketAddress;
4758
import java.util.ArrayList;
4859
import java.util.List;
4960
import java.util.Map;
5061
import java.util.Optional;
62+
import java.util.concurrent.atomic.AtomicReference;
5163
import org.junit.jupiter.api.BeforeEach;
5264
import org.junit.jupiter.api.DisplayName;
5365
import org.junit.jupiter.api.Test;
@@ -57,6 +69,25 @@ class LangChain4jTest {
5769

5870
private static final String MODEL_NAME = "test-model";
5971

72+
/**
73+
* An output schema with an untyped property, an array without items, and an untyped object with
74+
* properties.
75+
*/
76+
private static final Schema FREE_FORM_SCHEMA =
77+
Schema.builder()
78+
.type("OBJECT")
79+
.properties(
80+
Map.of(
81+
"extra",
82+
Schema.builder().description("Anything").build(),
83+
"tags",
84+
Schema.builder().type("ARRAY").build(),
85+
"nested",
86+
Schema.builder()
87+
.properties(Map.of("a", Schema.builder().type("STRING").build()))
88+
.build()))
89+
.build();
90+
6091
private ChatModel chatModel;
6192
private StreamingChatModel streamingChatModel;
6293
private LangChain4j langChain4j;
@@ -977,6 +1008,100 @@ void testGenerateContentWithMcpToolParametersJsonSchemaAsSchema() {
9771008
assertThat(capturedRequest.toolSpecifications().get(0).description()).isEqualTo("An MCP tool");
9781009
}
9791010

1011+
@Test
1012+
@DisplayName("Should convert free-form schemas in tool parameters")
1013+
void testGenerateContentWithFreeFormToolParameters() {
1014+
// Given
1015+
final ChatResponse chatResponse = mock(ChatResponse.class);
1016+
when(chatResponse.aiMessage()).thenReturn(AiMessage.from("Done"));
1017+
when(chatModel.chat(any(ChatRequest.class))).thenReturn(chatResponse);
1018+
1019+
// When
1020+
final LlmResponse response =
1021+
langChain4j.generateContent(freeFormSetModelResponseRequest(), false).blockingFirst();
1022+
1023+
// Then
1024+
assertThat(response.content().get().text()).isEqualTo("Done");
1025+
final ArgumentCaptor<ChatRequest> requestCaptor = ArgumentCaptor.forClass(ChatRequest.class);
1026+
verify(chatModel).chat(requestCaptor.capture());
1027+
final Map<String, JsonSchemaElement> properties =
1028+
requestCaptor.getValue().toolSpecifications().get(0).parameters().properties();
1029+
final JsonAnyOfSchema extra = (JsonAnyOfSchema) properties.get("extra");
1030+
assertThat(extra.description()).isEqualTo("Anything");
1031+
assertThat(extra.anyOf()).hasSize(5);
1032+
assertThat(((JsonArraySchema) properties.get("tags")).items())
1033+
.isInstanceOf(JsonAnyOfSchema.class);
1034+
assertThat(((JsonObjectSchema) properties.get("nested")).properties().get("a"))
1035+
.isInstanceOf(JsonStringSchema.class);
1036+
}
1037+
1038+
@Test
1039+
@DisplayName("Should send free-form tool parameters through the Gemini provider")
1040+
void testGenerateContentWithFreeFormToolParameters_geminiProvider() throws Exception {
1041+
// Given
1042+
final AtomicReference<String> requestBody = new AtomicReference<>();
1043+
final HttpServer server =
1044+
HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0);
1045+
server.createContext(
1046+
"/",
1047+
exchange -> {
1048+
requestBody.set(new String(exchange.getRequestBody().readAllBytes(), UTF_8));
1049+
final byte[] reply =
1050+
("{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Done\"}]},"
1051+
+ "\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":1,"
1052+
+ "\"candidatesTokenCount\":1,\"totalTokenCount\":2}}")
1053+
.getBytes(UTF_8);
1054+
exchange.getResponseHeaders().add("Content-Type", "application/json");
1055+
exchange.sendResponseHeaders(200, reply.length);
1056+
try (OutputStream body = exchange.getResponseBody()) {
1057+
body.write(reply);
1058+
}
1059+
});
1060+
1061+
// When
1062+
final LlmResponse response;
1063+
server.start();
1064+
try {
1065+
final LangChain4j gemini =
1066+
LangChain4j.builder()
1067+
.chatModel(
1068+
GoogleAiGeminiChatModel.builder()
1069+
.apiKey("test-key")
1070+
.modelName("gemini-2.5-flash")
1071+
.baseUrl("http://127.0.0.1:" + server.getAddress().getPort())
1072+
.maxRetries(0)
1073+
.build())
1074+
.build();
1075+
response = gemini.generateContent(freeFormSetModelResponseRequest(), false).blockingFirst();
1076+
} finally {
1077+
server.stop(0);
1078+
}
1079+
1080+
// Then
1081+
assertThat(response.content().get().text()).isEqualTo("Done");
1082+
// LangChain4j 1.12 sends "tools" as one object, later versions as a list of them.
1083+
final JsonNode tools = new ObjectMapper().readTree(requestBody.get()).path("tools");
1084+
final JsonNode properties =
1085+
(tools.isArray() ? tools.path(0) : tools)
1086+
.at("/functionDeclarations/0/parameters/properties");
1087+
assertThat(properties.isObject()).isTrue();
1088+
assertThat(properties.at("/extra/anyOf").size()).isEqualTo(5);
1089+
assertThat(properties.at("/nested/properties/a/type").asText()).isEqualToIgnoringCase("string");
1090+
// Gemini rejects an array schema without items.
1091+
assertThat(
1092+
properties.findParents("type").stream()
1093+
.filter(node -> node.get("type").asText().equalsIgnoreCase("array")))
1094+
.isNotEmpty()
1095+
.allMatch(node -> node.has("items"));
1096+
}
1097+
1098+
private static LlmRequest freeFormSetModelResponseRequest() {
1099+
return LlmRequest.builder()
1100+
.contents(List.of(Content.fromParts(Part.fromText("Answer"))))
1101+
.tools(Map.of(SetModelResponseTool.NAME, new SetModelResponseTool(FREE_FORM_SCHEMA)))
1102+
.build();
1103+
}
1104+
9801105
@Test
9811106
@DisplayName(
9821107
"Should use TokenCountEstimator to estimate token usage when TokenUsage is not available")

‎core/src/main/java/com/google/adk/SchemaUtils.java‎

Lines changed: 32 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -41,11 +41,17 @@ private SchemaUtils() {} // Private constructor for utility class
4141
*/
4242
@SuppressWarnings("unchecked") // For tool parameter type casting.
4343
private static Boolean matchType(Object value, Schema schema, Boolean isInput) {
44+
// A schema with no type and no properties accepts any value, and null unless nullable=false.
45+
boolean freeForm = schema.type().isEmpty() && schema.properties().isEmpty();
4446
if (value == null) {
45-
return schema.nullable().orElse(false);
47+
return schema.nullable().orElse(freeForm);
48+
}
49+
if (freeForm) {
50+
return true;
4651
}
4752
// Based on types from https://cloud.google.com/vertex-ai/docs/reference/rest/v1/Schema
48-
Type.Known type = schema.type().get().knownEnum();
53+
// An untyped schema with properties describes an object.
54+
Type.Known type = schema.type().map(Type::knownEnum).orElse(Type.Known.OBJECT);
4955
switch (type) {
5056
case STRING:
5157
return value instanceof String;
@@ -57,6 +63,9 @@ private static Boolean matchType(Object value, Schema schema, Boolean isInput) {
5763
return value instanceof Number;
5864
case ARRAY:
5965
if (value instanceof List) {
66+
if (schema.items().isEmpty()) {
67+
return true;
68+
}
6069
for (Object element : (List<?>) value) {
6170
if (!matchType(element, schema.items().get(), isInput)) {
6271
return false;
@@ -91,29 +100,14 @@ private static Boolean matchType(Object value, Schema schema, Boolean isInput) {
91100
*/
92101
public static void validateMapOnSchema(Map<String, Object> args, Schema schema, Boolean isInput) {
93102
Preconditions.checkNotNull(isInput, "IsInput cannot be null");
94-
Map<String, Schema> properties = schema.properties().get();
95-
for (Entry<String, Object> arg : args.entrySet()) {
96-
// Check if the argument is in the schema.
97-
if (!properties.containsKey(arg.getKey())) {
98-
if (isInput) {
99-
throw new IllegalArgumentException(
100-
"Input arg: " + arg.getKey() + " does not match agent input schema: " + schema);
101-
} else {
102-
throw new IllegalArgumentException(
103-
"Output arg: " + arg.getKey() + " does not match agent output schema: " + schema);
104-
}
105-
}
106-
// Check if the argument type matches the schema type.
107-
if (!matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) {
108-
if (isInput) {
109-
throw new IllegalArgumentException(
110-
"Input arg: " + arg.getKey() + " does not match agent input schema: " + schema);
111-
} else {
112-
throw new IllegalArgumentException(
113-
"Output arg: " + arg.getKey() + " does not match agent output schema: " + schema);
114-
}
115-
}
103+
if (schema.type().map(Type::knownEnum).orElse(Type.Known.OBJECT) != Type.Known.OBJECT) {
104+
throw new IllegalArgumentException(
105+
(isInput ? "Input" : "Output") + " schema does not describe an object: " + schema);
116106
}
107+
// A schema without properties leaves the keys open, as for a free-form map.
108+
schema
109+
.properties()
110+
.ifPresent(properties -> validateProperties(args, properties, schema, isInput));
117111
// Check if all required arguments are present.
118112
if (schema.required().isPresent()) {
119113
for (String required : schema.required().get()) {
@@ -128,6 +122,20 @@ public static void validateMapOnSchema(Map<String, Object> args, Schema schema,
128122
}
129123
}
130124

125+
/** Checks that every key in {@code args} is one of {@code properties} and matches its schema. */
126+
private static void validateProperties(
127+
Map<String, Object> args, Map<String, Schema> properties, Schema schema, boolean isInput) {
128+
for (Entry<String, Object> arg : args.entrySet()) {
129+
if (!properties.containsKey(arg.getKey())
130+
|| !matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) {
131+
throw new IllegalArgumentException(
132+
isInput
133+
? "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema
134+
: "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema);
135+
}
136+
}
137+
}
138+
131139
/**
132140
* Validates an output string against a schema.
133141
*

0 commit comments

Comments
 (0)