From 3992b48baec764776acba7e5c53b882d0519f1c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Wed, 4 Nov 2026 03:13:51 -0800 Subject: [PATCH] fix(langchain4j): end a streamed turn with one response that holds the whole answer PiperOrigin-RevId: 1013707586 --- .../adk/models/langchain4j/LangChain4j.java | 36 ++++---- .../LangChain4jIntegrationTest.java | 22 ++++- .../models/langchain4j/LangChain4jTest.java | 92 +++++++++++++++++++ .../adk/tools/SetModelResponseTool.java | 28 ++++-- .../adk/tools/SetModelResponseToolTest.java | 91 ++++++++++++++++++ 5 files changed, 238 insertions(+), 31 deletions(-) diff --git a/contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java b/contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java index ae2544e6e..34e730a15 100644 --- a/contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java +++ b/contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java @@ -217,30 +217,28 @@ public Flowable generateContent(LlmRequest llmRequest, boolean stre public void onPartialResponse(String s) { emitter.onNext( LlmResponse.builder() - .content(Content.fromParts(Part.fromText(s))) + .content( + Content.builder().role("model").parts(Part.fromText(s)).build()) + .partial(true) .build()); } @Override public void onCompleteResponse(ChatResponse chatResponse) { - if (chatResponse.aiMessage().hasToolExecutionRequests()) { - AiMessage aiMessage = chatResponse.aiMessage(); - toParts(aiMessage).stream() - .map(Part::functionCall) - .forEach( - functionCall -> { - functionCall.ifPresent( - function -> { - emitter.onNext( - LlmResponse.builder() - .content( - Content.fromParts( - Part.fromFunctionCall( - function.name().orElse(""), - function.args().orElse(Map.of())))) - .build()); - }); - }); + // The chunks were partial, so one response with the whole turn ends it. + AiMessage aiMessage = chatResponse.aiMessage(); + List parts = new ArrayList<>(); + if (aiMessage.text() != null && !aiMessage.text().isEmpty()) { + parts.add(Part.fromText(aiMessage.text())); + } + if (aiMessage.hasToolExecutionRequests()) { + parts.addAll(toParts(aiMessage)); + } + if (!parts.isEmpty()) { + emitter.onNext( + LlmResponse.builder() + .content(Content.builder().role("model").parts(parts).build()) + .build()); } emitter.onComplete(); } diff --git a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java index 4a11cd886..a3dbf87ac 100644 --- a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java +++ b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java @@ -375,6 +375,7 @@ void testSimpleStreamingResponse() { "", responses .blockingStream() + .filter(llmResponse -> !llmResponse.partial().orElse(false)) .map(llmResponse -> llmResponse.content().get().text()) .toList()); @@ -430,15 +431,30 @@ void testStreamingRunConfig() { // when List eventsHi = askAgentStreaming(agent, "Hi"); String responseToHi = - String.join("", eventsHi.stream().map(event -> event.content().get().text()).toList()); + String.join( + "", + eventsHi.stream() + .filter(event -> !event.partial().orElse(false)) + .map(event -> event.content().get().text()) + .toList()); List eventsQubit = askAgentStreaming(agent, "Tell me about qubits"); String responseToQubit = - String.join("", eventsQubit.stream().map(event -> event.content().get().text()).toList()); + String.join( + "", + eventsQubit.stream() + .filter(event -> !event.partial().orElse(false)) + .map(event -> event.content().get().text()) + .toList()); List eventsWeather = askAgentStreaming(agent, "What's the weather in Paris?"); String responseToWeather = - String.join("", eventsWeather.stream().map(Event::stringifyContent).toList()); + String.join( + "", + eventsWeather.stream() + .filter(event -> !event.partial().orElse(false)) + .map(Event::stringifyContent) + .toList()); // then diff --git a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java index 9c56a0743..ff6af9614 100644 --- a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java +++ b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java @@ -350,6 +350,17 @@ void testGenerateContentWithStreamingChatModel() { assertThat(responses).hasSize(5); // 4 partial responses + 1 function call // Verify the partial responses + assertThat(responses.stream().map(response -> response.content().orElseThrow().role())) + .containsExactly( + Optional.of("model"), + Optional.of("model"), + Optional.of("model"), + Optional.of("model"), + Optional.of("model")); + assertThat(responses.subList(0, 4).stream().map(LlmResponse::partial)) + .containsExactly( + Optional.of(true), Optional.of(true), Optional.of(true), Optional.of(true)); + assertThat(responses.get(4).partial()).isEmpty(); assertThat(responses.get(0).content().orElseThrow().text()).isEqualTo("Hello"); assertThat(responses.get(1).content().orElseThrow().text()).isEqualTo(", how"); assertThat(responses.get(2).content().orElseThrow().text()).isEqualTo(" can I help"); @@ -369,12 +380,93 @@ void testGenerateContentWithStreamingChatModel() { .get(0) .functionCall() .orElseThrow(); + assertThat(functionCall.id()).hasValue("123"); assertThat(functionCall.name()).isEqualTo(Optional.of("getWeather")); assertThat(functionCall.args().orElseThrow()).containsEntry("city", "Paris"); disposable.dispose(); } + @Test + @DisplayName("Should end a streamed text turn with its full text") + void testGenerateContentWithStreamingChatModel_endsWithTheFullText() { + final LlmRequest llmRequest = + LlmRequest.builder().contents(List.of(Content.fromParts(Part.fromText("Hello")))).build(); + final List responses = new ArrayList<>(); + final StreamingChatResponseHandler[] handlerRef = new StreamingChatResponseHandler[1]; + doAnswer( + invocation -> { + handlerRef[0] = invocation.getArgument(1); + return null; + }) + .when(streamingChatModel) + .chat(any(ChatRequest.class), any(StreamingChatResponseHandler.class)); + final var disposable = + streamingLangChain4j.generateContent(llmRequest, true).subscribe(responses::add); + + handlerRef[0].onPartialResponse("Par"); + handlerRef[0].onPartialResponse("is"); + final ChatResponse chatResponse = mock(ChatResponse.class); + when(chatResponse.aiMessage()).thenReturn(AiMessage.from("Paris")); + handlerRef[0].onCompleteResponse(chatResponse); + + assertThat(responses).hasSize(3); + final LlmResponse last = responses.get(2); + assertThat(last.partial()).isEmpty(); + assertThat(last.content().orElseThrow().role()).hasValue("model"); + assertThat(last.content().orElseThrow().text()).isEqualTo("Paris"); + disposable.dispose(); + } + + @Test + @DisplayName("Should end a streamed turn with its text and calls in one response") + void testGenerateContentWithStreamingChatModel_endsWithTextAndCallsInOneResponse() { + final LlmRequest llmRequest = + LlmRequest.builder().contents(List.of(Content.fromParts(Part.fromText("Hello")))).build(); + final List responses = new ArrayList<>(); + final StreamingChatResponseHandler[] handlerRef = new StreamingChatResponseHandler[1]; + doAnswer( + invocation -> { + handlerRef[0] = invocation.getArgument(1); + return null; + }) + .when(streamingChatModel) + .chat(any(ChatRequest.class), any(StreamingChatResponseHandler.class)); + final var disposable = + streamingLangChain4j.generateContent(llmRequest, true).subscribe(responses::add); + + handlerRef[0].onPartialResponse("Let me check."); + final AiMessage aiMessage = + AiMessage.builder() + .text("Let me check.") + .toolExecutionRequests( + List.of( + ToolExecutionRequest.builder() + .id("c1") + .name("getWeather") + .arguments("{\"city\":\"Paris\"}") + .build(), + ToolExecutionRequest.builder() + .id("c2") + .name("getWeather") + .arguments("{\"city\":\"Rome\"}") + .build())) + .build(); + final ChatResponse chatResponse = mock(ChatResponse.class); + when(chatResponse.aiMessage()).thenReturn(aiMessage); + handlerRef[0].onCompleteResponse(chatResponse); + + assertThat(responses).hasSize(2); + final LlmResponse last = responses.get(1); + assertThat(last.partial()).isEmpty(); + final List parts = last.content().orElseThrow().parts().orElseThrow(); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).text()).hasValue("Let me check."); + assertThat(parts.subList(1, 3).stream().map(part -> part.functionCall().orElseThrow().id())) + .containsExactly(Optional.of("c1"), Optional.of("c2")); + disposable.dispose(); + } + @Test @DisplayName("Should pass configuration options to LangChain4j") void testGenerateContentWithConfigOptions() { diff --git a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java index 4f7022914..88f52cf02 100644 --- a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java +++ b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java @@ -17,6 +17,7 @@ package com.google.adk.tools; import com.google.adk.SchemaUtils; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.Schema; @@ -77,34 +78,43 @@ public Single> runAsync(Map args, ToolContex + "\nRecall the set_model_response function correctly, fix the errors, and" + " call it again with all required fields using the correct types."); } - // Match Python's model_dump(exclude_none=True) for Java's map-shaped response. - Map validatedResponse = excludeNullFields(args); + Map validatedResponse = excludeNullFields(args, outputSchema); toolContext.actions().setSetModelResponse(validatedResponse); return validatedResponse; }); } - private static Map excludeNullFields(Map values) { + /** + * Drops nulls like Python's {@code model_dump(exclude_none=True)}, but keeps the null of a + * property that {@code schema} marks required and nullable, which the schema demands. + */ + private static Map excludeNullFields(Map values, Schema schema) { + // Validation passed, so each key has a property schema and each non-empty list an item schema. + Map properties = schema.properties().orElseThrow(); + List required = schema.required().orElse(ImmutableList.of()); Map result = new LinkedHashMap<>(); for (Map.Entry entry : values.entrySet()) { + Schema property = properties.get(entry.getKey()); Object value = entry.getValue(); if (value != null) { - result.put(entry.getKey(), excludeNullsFromValue(value)); + result.put(entry.getKey(), excludeNullsFromValue(value, property)); + } else if (property.nullable().orElse(false) && required.contains(entry.getKey())) { + result.put(entry.getKey(), null); } } return result; } - private static Object excludeNullsFromValue(Object value) { + private static Object excludeNullsFromValue(Object value, Schema schema) { if (value instanceof Map) { @SuppressWarnings("unchecked") // Function-call args always have String keys. Map map = (Map) value; - return excludeNullFields(map); + return excludeNullFields(map, schema); } - if (value instanceof List) { + if (value instanceof List list) { List result = new ArrayList<>(); - for (Object item : (List) value) { - result.add(excludeNullsFromValue(item)); + for (Object item : list) { + result.add(excludeNullsFromValue(item, schema.items().orElseThrow())); } return result; } diff --git a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java index 1964daa98..ea35f2f85 100644 --- a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java @@ -21,6 +21,7 @@ import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.common.truth.Truth.assertThat; +import com.google.adk.SchemaUtils; import com.google.adk.models.LlmResponse; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -123,6 +124,96 @@ public void runAsync_nullableNull_excludesNullFromReturnedAndRecordedResponse() assertThat(toolContext.actions().setModelResponse()).hasValue(ImmutableMap.of("a", "x")); } + @Test + public void runAsync_requiredNullableNull_keepsTheNull() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "a", + Schema.builder().type("STRING").build(), + "b", + Schema.builder().type("STRING").nullable(true).build())) + .required(ImmutableList.of("a", "b")) + .build(); + Map args = new LinkedHashMap<>(); + args.put("a", "x"); + args.put("b", null); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema).runAsync(args, toolContext).blockingGet(); + + assertThat(result).containsExactly("a", "x", "b", null).inOrder(); + assertThat(toolContext.actions().setModelResponse()).hasValue(result); + // Throws if the result does not match the schema, as dropping the null did. + SchemaUtils.validateMapOnSchema(result, outputSchema, /* isInput= */ false); + } + + @Test + public void runAsync_nestedRequiredNullableNull_keepsTheNull() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("addr", requiredNullableZipAddressSchema())) + .required(ImmutableList.of("addr")) + .build(); + Map address = new LinkedHashMap<>(); + address.put("city", "NYC"); + address.put("zip", null); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema) + .runAsync(ImmutableMap.of("addr", address), toolContext) + .blockingGet(); + + assertThat(result).isEqualTo(ImmutableMap.of("addr", address)); + assertThat(toolContext.actions().setModelResponse()).hasValue(result); + } + + @Test + public void runAsync_arrayItemRequiredNullableNull_keepsTheNull() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "addresses", + Schema.builder() + .type("ARRAY") + .items(requiredNullableZipAddressSchema()) + .build())) + .required(ImmutableList.of("addresses")) + .build(); + Map address = new LinkedHashMap<>(); + address.put("city", "NYC"); + address.put("zip", null); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema) + .runAsync(ImmutableMap.of("addresses", ImmutableList.of(address)), toolContext) + .blockingGet(); + + assertThat(result).isEqualTo(ImmutableMap.of("addresses", ImmutableList.of(address))); + assertThat(toolContext.actions().setModelResponse()).hasValue(result); + } + + private static Schema requiredNullableZipAddressSchema() { + return Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "city", + Schema.builder().type("STRING").build(), + "zip", + Schema.builder().type("STRING").nullable(true).build())) + .required(ImmutableList.of("city", "zip")) + .build(); + } + @Test public void runAsync_nestedNullableNull_excludesNullFromReturnedAndRecordedResponse() { Schema addressSchema =