Skip to content

fix: keep whole numbers in streamed function call arguments as integers #1796

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: keep whole numbers in streamed function call arguments as integers #1796

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 55 additions & 25 deletions contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.ToolChoice;
import dev.langchain4j.model.chat.request.json.JsonAnyOfSchema;
import dev.langchain4j.model.chat.request.json.JsonArraySchema;
import dev.langchain4j.model.chat.request.json.JsonBooleanSchema;
import dev.langchain4j.model.chat.request.json.JsonIntegerSchema;
Expand Down Expand Up @@ -564,7 +565,11 @@ private List<ToolSpecification> toToolSpecifications(LlmRequest llmRequest) {
}

private JsonObjectSchema toParameters(Schema schema) {
if (schema.type().isPresent() && Type.Known.OBJECT.equals(schema.type().get().knownEnum())) {
// An untyped schema with properties describes an object.
if (schema
.type()
.map(t -> t.knownEnum() == Type.Known.OBJECT)
.orElse(schema.properties().isPresent())) {
return JsonObjectSchema.builder()
.addProperties(toProperties(schema))
.required(schema.required().orElse(List.of()))
Expand All @@ -582,31 +587,56 @@ private Map<String, JsonSchemaElement> toProperties(Schema schema) {
return result;
}

private JsonSchemaElement toJsonSchemaElement(Schema schema) {
if (schema != null && schema.type().isPresent()) {
Type type = schema.type().get();
return switch (type.knownEnum()) {
case STRING ->
JsonStringSchema.builder().description(schema.description().orElse(null)).build();
case NUMBER ->
JsonNumberSchema.builder().description(schema.description().orElse(null)).build();
case INTEGER ->
JsonIntegerSchema.builder().description(schema.description().orElse(null)).build();
case BOOLEAN ->
JsonBooleanSchema.builder().description(schema.description().orElse(null)).build();
case ARRAY ->
JsonArraySchema.builder()
.description(schema.description().orElse(null))
.items(toJsonSchemaElement(schema.items().orElseThrow()))
.build();
case OBJECT -> toParameters(schema);
default ->
throw new UnsupportedFeatureException(
"LangChain4jLlm does not support schema of type: " + type);
};
} else {
throw new IllegalArgumentException("Schema type cannot be null or absent");
private JsonSchemaElement toJsonSchemaElement(@Nullable Schema schema) {
if (schema == null) {
return anyValue(null);
}
// An untyped schema with properties is an object; without them, it accepts any value.
if (schema.type().isEmpty()) {
return schema.properties().isPresent()
? toParameters(schema)
: anyValue(schema.description().orElse(null));
}
Type type = schema.type().get();
return switch (type.knownEnum()) {
case STRING ->
JsonStringSchema.builder().description(schema.description().orElse(null)).build();
case NUMBER ->
JsonNumberSchema.builder().description(schema.description().orElse(null)).build();
case INTEGER ->
JsonIntegerSchema.builder().description(schema.description().orElse(null)).build();
case BOOLEAN ->
JsonBooleanSchema.builder().description(schema.description().orElse(null)).build();
case ARRAY ->
JsonArraySchema.builder()
.description(schema.description().orElse(null))
.items(schema.items().map(this::toJsonSchemaElement).orElseGet(() -> anyValue(null)))
.build();
case OBJECT -> toParameters(schema);
default ->
throw new UnsupportedFeatureException(
"LangChain4jLlm does not support schema of type: " + type);
};
}

/**
* Returns an element that accepts a string, number, boolean, object, or an array of those, for a
* schema without a type. Not a {@code JsonRawSchema}, which the Gemini provider rejects. The
* array branch has items, which Gemini requires.
*/
private static JsonSchemaElement anyValue(@Nullable String description) {
List<JsonSchemaElement> nonArrays =
List.of(
new JsonStringSchema(),
new JsonNumberSchema(),
new JsonBooleanSchema(),
JsonObjectSchema.builder().build());
List<JsonSchemaElement> anyOf = new ArrayList<>(nonArrays);
anyOf.add(
JsonArraySchema.builder()
.items(JsonAnyOfSchema.builder().anyOf(nonArrays).build())
.build());
return JsonAnyOfSchema.builder().description(description).anyOf(anyOf).build();
}

private LlmResponse toLlmResponse(ChatResponse chatResponse, ChatRequest chatRequest) {
Expand Down
125 changes: 125 additions & 0 deletions contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,13 @@
*/
package com.google.adk.models.langchain4j;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.assertj.core.api.Assertions.*;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.adk.agents.LlmAgent;
import com.google.adk.agents.RunConfig;
Expand All @@ -28,7 +30,9 @@
import com.google.adk.runner.InMemoryRunner;
import com.google.adk.sessions.Session;
import com.google.adk.tools.FunctionTool;
import com.google.adk.tools.SetModelResponseTool;
import com.google.genai.types.*;
import com.sun.net.httpserver.HttpServer;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
Expand All @@ -37,17 +41,25 @@
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.json.JsonAnyOfSchema;
import dev.langchain4j.model.chat.request.json.JsonArraySchema;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchemaElement;
import dev.langchain4j.model.chat.request.json.JsonStringSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.googleai.GoogleAiGeminiChatModel;
import dev.langchain4j.model.output.TokenUsage;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.subscribers.TestSubscriber;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
Expand All @@ -57,6 +69,25 @@ class LangChain4jTest {

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

/**
* An output schema with an untyped property, an array without items, and an untyped object with
* properties.
*/
private static final Schema FREE_FORM_SCHEMA =
Schema.builder()
.type("OBJECT")
.properties(
Map.of(
"extra",
Schema.builder().description("Anything").build(),
"tags",
Schema.builder().type("ARRAY").build(),
"nested",
Schema.builder()
.properties(Map.of("a", Schema.builder().type("STRING").build()))
.build()))
.build();

private ChatModel chatModel;
private StreamingChatModel streamingChatModel;
private LangChain4j langChain4j;
Expand Down Expand Up @@ -977,6 +1008,100 @@ void testGenerateContentWithMcpToolParametersJsonSchemaAsSchema() {
assertThat(capturedRequest.toolSpecifications().get(0).description()).isEqualTo("An MCP tool");
}

@Test
@DisplayName("Should convert free-form schemas in tool parameters")
void testGenerateContentWithFreeFormToolParameters() {
// Given
final ChatResponse chatResponse = mock(ChatResponse.class);
when(chatResponse.aiMessage()).thenReturn(AiMessage.from("Done"));
when(chatModel.chat(any(ChatRequest.class))).thenReturn(chatResponse);

// When
final LlmResponse response =
langChain4j.generateContent(freeFormSetModelResponseRequest(), false).blockingFirst();

// Then
assertThat(response.content().get().text()).isEqualTo("Done");
final ArgumentCaptor<ChatRequest> requestCaptor = ArgumentCaptor.forClass(ChatRequest.class);
verify(chatModel).chat(requestCaptor.capture());
final Map<String, JsonSchemaElement> properties =
requestCaptor.getValue().toolSpecifications().get(0).parameters().properties();
final JsonAnyOfSchema extra = (JsonAnyOfSchema) properties.get("extra");
assertThat(extra.description()).isEqualTo("Anything");
assertThat(extra.anyOf()).hasSize(5);
assertThat(((JsonArraySchema) properties.get("tags")).items())
.isInstanceOf(JsonAnyOfSchema.class);
assertThat(((JsonObjectSchema) properties.get("nested")).properties().get("a"))
.isInstanceOf(JsonStringSchema.class);
}

@Test
@DisplayName("Should send free-form tool parameters through the Gemini provider")
void testGenerateContentWithFreeFormToolParameters_geminiProvider() throws Exception {
// Given
final AtomicReference<String> requestBody = new AtomicReference<>();
final HttpServer server =
HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0);
server.createContext(
"/",
exchange -> {
requestBody.set(new String(exchange.getRequestBody().readAllBytes(), UTF_8));
final byte[] reply =
("{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Done\"}]},"
+ "\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":1,"
+ "\"candidatesTokenCount\":1,\"totalTokenCount\":2}}")
.getBytes(UTF_8);
exchange.getResponseHeaders().add("Content-Type", "application/json");
exchange.sendResponseHeaders(200, reply.length);
try (OutputStream body = exchange.getResponseBody()) {
body.write(reply);
}
});

// When
final LlmResponse response;
server.start();
try {
final LangChain4j gemini =
LangChain4j.builder()
.chatModel(
GoogleAiGeminiChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.5-flash")
.baseUrl("http://127.0.0.1:" + server.getAddress().getPort())
.maxRetries(0)
.build())
.build();
response = gemini.generateContent(freeFormSetModelResponseRequest(), false).blockingFirst();
} finally {
server.stop(0);
}

// Then
assertThat(response.content().get().text()).isEqualTo("Done");
// LangChain4j 1.12 sends "tools" as one object, later versions as a list of them.
final JsonNode tools = new ObjectMapper().readTree(requestBody.get()).path("tools");
final JsonNode properties =
(tools.isArray() ? tools.path(0) : tools)
.at("/functionDeclarations/0/parameters/properties");
assertThat(properties.isObject()).isTrue();
assertThat(properties.at("/extra/anyOf").size()).isEqualTo(5);
assertThat(properties.at("/nested/properties/a/type").asText()).isEqualToIgnoringCase("string");
// Gemini rejects an array schema without items.
assertThat(
properties.findParents("type").stream()
.filter(node -> node.get("type").asText().equalsIgnoreCase("array")))
.isNotEmpty()
.allMatch(node -> node.has("items"));
}

private static LlmRequest freeFormSetModelResponseRequest() {
return LlmRequest.builder()
.contents(List.of(Content.fromParts(Part.fromText("Answer"))))
.tools(Map.of(SetModelResponseTool.NAME, new SetModelResponseTool(FREE_FORM_SCHEMA)))
.build();
}

@Test
@DisplayName(
"Should use TokenCountEstimator to estimate token usage when TokenUsage is not available")
Expand Down
56 changes: 32 additions & 24 deletions core/src/main/java/com/google/adk/SchemaUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,11 +41,17 @@ private SchemaUtils() {} // Private constructor for utility class
*/
@SuppressWarnings("unchecked") // For tool parameter type casting.
private static Boolean matchType(Object value, Schema schema, Boolean isInput) {
// A schema with no type and no properties accepts any value, and null unless nullable=false.
boolean freeForm = schema.type().isEmpty() && schema.properties().isEmpty();
if (value == null) {
return schema.nullable().orElse(false);
return schema.nullable().orElse(freeForm);
}
if (freeForm) {
return true;
}
// Based on types from https://cloud.google.com/vertex-ai/docs/reference/rest/v1/Schema
Type.Known type = schema.type().get().knownEnum();
// An untyped schema with properties describes an object.
Type.Known type = schema.type().map(Type::knownEnum).orElse(Type.Known.OBJECT);
switch (type) {
case STRING:
return value instanceof String;
Expand All @@ -57,6 +63,9 @@ private static Boolean matchType(Object value, Schema schema, Boolean isInput) {
return value instanceof Number;
case ARRAY:
if (value instanceof List) {
if (schema.items().isEmpty()) {
return true;
}
for (Object element : (List<?>) value) {
if (!matchType(element, schema.items().get(), isInput)) {
return false;
Expand Down Expand Up @@ -91,29 +100,14 @@ private static Boolean matchType(Object value, Schema schema, Boolean isInput) {
*/
public static void validateMapOnSchema(Map<String, Object> args, Schema schema, Boolean isInput) {
Preconditions.checkNotNull(isInput, "IsInput cannot be null");
Map<String, Schema> properties = schema.properties().get();
for (Entry<String, Object> arg : args.entrySet()) {
// Check if the argument is in the schema.
if (!properties.containsKey(arg.getKey())) {
if (isInput) {
throw new IllegalArgumentException(
"Input arg: " + arg.getKey() + " does not match agent input schema: " + schema);
} else {
throw new IllegalArgumentException(
"Output arg: " + arg.getKey() + " does not match agent output schema: " + schema);
}
}
// Check if the argument type matches the schema type.
if (!matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) {
if (isInput) {
throw new IllegalArgumentException(
"Input arg: " + arg.getKey() + " does not match agent input schema: " + schema);
} else {
throw new IllegalArgumentException(
"Output arg: " + arg.getKey() + " does not match agent output schema: " + schema);
}
}
if (schema.type().map(Type::knownEnum).orElse(Type.Known.OBJECT) != Type.Known.OBJECT) {
throw new IllegalArgumentException(
(isInput ? "Input" : "Output") + " schema does not describe an object: " + schema);
}
// A schema without properties leaves the keys open, as for a free-form map.
schema
.properties()
.ifPresent(properties -> validateProperties(args, properties, schema, isInput));
// Check if all required arguments are present.
if (schema.required().isPresent()) {
for (String required : schema.required().get()) {
Expand All @@ -128,6 +122,20 @@ public static void validateMapOnSchema(Map<String, Object> args, Schema schema,
}
}

/** Checks that every key in {@code args} is one of {@code properties} and matches its schema. */
private static void validateProperties(
Map<String, Object> args, Map<String, Schema> properties, Schema schema, boolean isInput) {
for (Entry<String, Object> arg : args.entrySet()) {
if (!properties.containsKey(arg.getKey())
|| !matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) {
throw new IllegalArgumentException(
isInput
? "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema
: "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema);
}
}
}

/**
* Validates an output string against a schema.
*
Expand Down
Loading
Loading