diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java index ece9b1e2f..abf161edc 100644 --- a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java @@ -18,20 +18,26 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.adk.models.BaseLlm; import com.google.adk.models.BaseLlmConnection; +import com.google.adk.models.FunctionCallIds; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; import com.google.adk.models.springai.error.SpringAIErrorMapper; import com.google.adk.models.springai.observability.SpringAIObservabilityHandler; import com.google.adk.models.springai.properties.SpringAIProperties; -import io.reactivex.rxjava3.core.BackpressureStrategy; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Flowable; +import java.util.ArrayList; +import java.util.List; import java.util.Objects; +import org.jspecify.annotations.Nullable; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.StreamingChatModel; import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; -import reactor.core.publisher.Flux; /** * Spring AI implementation of BaseLlm that wraps Spring AI ChatModel and StreamingChatModel. @@ -195,60 +201,158 @@ private Flowable generateContent(LlmRequest llmRequest) { } private Flowable generateStreamingContent(LlmRequest llmRequest) { - SpringAIObservabilityHandler.RequestContext context = - observabilityHandler.startRequest(model(), "streaming"); - - return Flowable.create( - emitter -> { - try { - Prompt prompt = messageConverter.toLlmPrompt(llmRequest, resolveDefaultOptions()); - observabilityHandler.logRequest(prompt.toString(), model()); - - Flux responseFlux = streamingChatModel.stream(prompt); - - responseFlux - .doOnError( - error -> { - observabilityHandler.recordError(context, error); - SpringAIErrorMapper.MappedError mappedError = - SpringAIErrorMapper.mapError(error); - emitter.onError( - new RuntimeException(mappedError.getNormalizedMessage(), error)); - }) - .subscribe( - chatResponse -> { - try { - // Use enhanced streaming-aware conversion - LlmResponse llmResponse = - messageConverter.toLlmResponse(chatResponse, true); - emitter.onNext(llmResponse); - } catch (Exception e) { - observabilityHandler.recordError(context, e); - SpringAIErrorMapper.MappedError mappedError = - SpringAIErrorMapper.mapError(e); - emitter.onError( - new RuntimeException(mappedError.getNormalizedMessage(), e)); - } - }, - error -> { - observabilityHandler.recordError(context, error); - SpringAIErrorMapper.MappedError mappedError = - SpringAIErrorMapper.mapError(error); - emitter.onError( - new RuntimeException(mappedError.getNormalizedMessage(), error)); - }, - () -> { - // Record success for streaming completion - observabilityHandler.recordSuccess(context, 0, 0, 0); - emitter.onComplete(); - }); - } catch (Exception e) { - observabilityHandler.recordError(context, e); - SpringAIErrorMapper.MappedError mappedError = SpringAIErrorMapper.mapError(e); - emitter.onError(new RuntimeException(mappedError.getNormalizedMessage(), e)); + return Flowable.defer( + () -> { + SpringAIObservabilityHandler.RequestContext context = + observabilityHandler.startRequest(model(), "streaming"); + // Deferred so a failure building the prompt also reaches onErrorResumeNext. + return Flowable.defer(() -> streamResponses(llmRequest)) + .doOnComplete(() -> observabilityHandler.recordSuccess(context, 0, 0, 0)) + .onErrorResumeNext( + error -> { + observabilityHandler.recordError(context, error); + SpringAIErrorMapper.MappedError mappedError = + SpringAIErrorMapper.mapError(error); + return Flowable.error( + new RuntimeException(mappedError.getNormalizedMessage(), error)); + }); + }); + } + + /** + * Emits each chunk that has content as a partial response and, if the model streamed anything, + * then one aggregated response, as {@code Gemini} does. ADK saves only the aggregated response to + * the session and runs its function calls. + */ + private Flowable streamResponses(LlmRequest llmRequest) { + Prompt prompt = messageConverter.toLlmPrompt(llmRequest, resolveDefaultOptions()); + observabilityHandler.logRequest(prompt.toString(), model()); + StreamAggregation aggregation = new StreamAggregation(); + return Flowable.fromPublisher(streamingChatModel.stream(prompt)) + .concatMap(aggregation::processChunk) + .concatWith(Flowable.defer(aggregation::aggregatedResponse)); + } + + /** The state of one stream, used to build its aggregated response. */ + private final class StreamAggregation { + // Parts in stream order. Consecutive text, or consecutive thought text, merges into one part; + // function calls, media and other parts are kept as they arrived. + private final List parts = new ArrayList<>(); + private final StringBuilder text = new StringBuilder(); + private boolean textIsThought; + private @Nullable GenerateContentResponseUsageMetadata usage; + private @Nullable String modelVersion; + private @Nullable LlmResponse finishingChunk; + private boolean receivedChunk; + + private Flowable processChunk(ChatResponse chatResponse) { + receivedChunk = true; + LlmResponse chunk = withFunctionCallIds(messageConverter.toLlmResponse(chatResponse)); + chunk.usageMetadata().ifPresent(chunkUsage -> usage = chunkUsage); + chunk.modelVersion().ifPresent(version -> modelVersion = version); + if (chunk.finishReason().isPresent()) { + finishingChunk = chunk; + } + List chunkParts = chunk.content().flatMap(Content::parts).orElse(List.of()); + if (chunkParts.isEmpty()) { + return Flowable.empty(); + } + for (Part part : chunkParts) { + if (part.text().isPresent()) { + boolean isThought = part.thought().orElse(false); + if (isThought != textIsThought) { + flushText(); } - }, - BackpressureStrategy.BUFFER); + textIsThought = isThought; + text.append(part.text().get()); + } else { + flushText(); + parts.add(part); + } + } + // As in Gemini, only the aggregated response reports an error. + return Flowable.just( + chunk.toBuilder() + .partial(true) + .turnComplete(null) + .errorCode(null) + .errorMessage(null) + .build()); + } + + private void flushText() { + if (text.isEmpty()) { + return; + } + Part.Builder part = Part.builder().text(text.toString()); + if (textIsThought) { + part.thought(true); + } + parts.add(part.build()); + text.setLength(0); + } + + private Flowable aggregatedResponse() { + if (!receivedChunk) { + return Flowable.empty(); + } + flushText(); + LlmResponse.Builder aggregated = + LlmResponse.builder().partial(false).turnComplete(true).modelVersion(modelVersion); + // A stream without content still ends with a response, as Gemini's does. + if (!parts.isEmpty()) { + aggregated.content(Content.builder().role("model").parts(parts).build()); + } + if (usage != null) { + aggregated.usageMetadata(usage); + } + // Report how the stream ended, from the last chunk that said so, as ADK Python does. + if (finishingChunk != null) { + aggregated + .finishReason(finishingChunk.finishReason().orElse(null)) + .errorCode(finishingChunk.errorCode().orElse(null)) + .errorMessage(finishingChunk.errorMessage().orElse(null)); + } + return Flowable.just(aggregated.build()); + } + } + + /** + * Gives each function call without an id a client-side one, so the partial and aggregated + * responses carrying it share the same id. + */ + private static LlmResponse withFunctionCallIds(LlmResponse response) { + if (response.content().isEmpty()) { + return response; + } + Content content = response.content().get(); + List parts = content.parts().orElse(List.of()); + if (parts.stream() + .noneMatch(part -> part.functionCall().filter(SpringAI::lacksId).isPresent())) { + return response; + } + return response.toBuilder() + .content( + content.toBuilder() + .parts(parts.stream().map(SpringAI::withFunctionCallId).toList()) + .build()) + .build(); + } + + private static Part withFunctionCallId(Part part) { + return part.functionCall() + .filter(SpringAI::lacksId) + .map( + call -> + part.toBuilder() + .functionCall( + call.toBuilder().id(FunctionCallIds.generateClientFunctionCallId()).build()) + .build()) + .orElse(part); + } + + private static boolean lacksId(FunctionCall call) { + return call.id().map(String::isEmpty).orElse(true); } @Override diff --git a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAITest.java b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAITest.java index 51fe60abb..538dc90b5 100644 --- a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAITest.java +++ b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAITest.java @@ -23,20 +23,34 @@ import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; +import com.google.genai.types.Blob; import com.google.genai.types.Content; +import com.google.genai.types.FinishReason; +import com.google.genai.types.FunctionCall; import com.google.genai.types.Part; +import io.reactivex.rxjava3.functions.Consumer; +import io.reactivex.rxjava3.plugins.RxJavaPlugins; import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.time.Duration; +import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.metadata.ChatGenerationMetadata; +import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.metadata.DefaultUsage; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.Generation; import org.springframework.ai.chat.model.StreamingChatModel; import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.content.Media; +import org.springframework.core.io.ByteArrayResource; +import org.springframework.util.MimeType; import reactor.core.publisher.Flux; class SpringAITest { @@ -149,10 +163,10 @@ void testGenerateContentStreaming() { testObserver.awaitDone(5, TimeUnit.SECONDS); testObserver.assertComplete(); testObserver.assertNoErrors(); - testObserver.assertValueCount(3); + // Three partial chunks and the aggregated response. + testObserver.assertValueCount(4); List responses = testObserver.values(); - assertThat(responses).hasSize(3); // Verify each streaming response for (LlmResponse response : responses) { @@ -273,7 +287,293 @@ void testGenerateContentStreamingBackpressure() { testObserver.awaitDone(10, TimeUnit.SECONDS); testObserver.assertComplete(); testObserver.assertNoErrors(); - testObserver.assertValueCount(1000); + // 1000 partial chunks and the aggregated response. + testObserver.assertValueCount(1001); + } + + @Test + void generateContent_streaming_emitsPartialChunksThenAggregatedResponse() { + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn( + Flux.just( + createStreamingChatResponse("Hello there."), + createStreamingChatResponse(" How are"), + createStreamingChatResponse(" you?"))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + testObserver.assertComplete(); + List responses = testObserver.values(); + assertThat(responses).hasSize(4); + assertThat(responses.subList(0, 3)) + .allSatisfy( + response -> { + assertThat(response.partial()).hasValue(true); + assertThat(response.turnComplete()).isEmpty(); + }); + assertThat(responses.subList(0, 3)) + .extracting(response -> response.content().get().text()) + .containsExactly("Hello there.", " How are", " you?"); + LlmResponse aggregated = responses.get(3); + assertThat(aggregated.partial().orElse(false)).isFalse(); + assertThat(aggregated.turnComplete()).hasValue(true); + assertThat(aggregated.content().get().text()).isEqualTo("Hello there. How are you?"); + } + + @Test + void generateContent_streaming_aggregatedResponseCarriesLastUsage() { + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn( + Flux.just( + createStreamingChatResponse("Hi", new DefaultUsage(3, 1, 4)), + createStreamingChatResponse(" there", new DefaultUsage(7, 5, 12)))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + LlmResponse aggregated = testObserver.values().get(testObserver.values().size() - 1); + assertThat(aggregated.usageMetadata().get().totalTokenCount()).hasValue(12); + } + + @Test + void generateContent_streamingCutShort_aggregatedResponseCarriesFinishReasonAndError() { + ChatResponse lastChunk = + new ChatResponse( + List.of( + new Generation( + new AssistantMessage(" there"), + ChatGenerationMetadata.builder().finishReason("MAX_TOKENS").build()))); + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn( + Flux.just( + createStreamingChatResponse("Hi"), + lastChunk, + createStreamingChatResponse("", new DefaultUsage(3, 1, 4)))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + List responses = testObserver.values(); + assertThat(responses.subList(0, responses.size() - 1)) + .allSatisfy(response -> assertThat(response.errorCode()).isEmpty()); + LlmResponse aggregated = responses.get(responses.size() - 1); + assertThat(aggregated.partial().orElse(false)).isFalse(); + assertThat(aggregated.finishReason().map(FinishReason::knownEnum)) + .hasValue(FinishReason.Known.MAX_TOKENS); + assertThat(aggregated.errorCode().map(FinishReason::knownEnum)) + .hasValue(FinishReason.Known.MAX_TOKENS); + assertThat(aggregated.errorMessage()).hasValue("Maximum tokens reached"); + } + + @Test + void generateContent_streamingStopOnEmptyLastChunk_aggregatedResponseCarriesStopWithoutError() { + ChatResponse emptyStopChunk = + new ChatResponse( + List.of( + new Generation( + new AssistantMessage(""), + ChatGenerationMetadata.builder().finishReason("STOP").build()))); + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn(Flux.just(createStreamingChatResponse("Hi"), emptyStopChunk)); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + LlmResponse aggregated = testObserver.values().get(testObserver.values().size() - 1); + assertThat(aggregated.finishReason().map(FinishReason::knownEnum)) + .hasValue(FinishReason.Known.STOP); + assertThat(aggregated.errorCode()).isEmpty(); + } + + @Test + void generateContent_streamingThoughts_aggregatedResponseKeepsThoughtsApartFromText() { + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn( + Flux.just( + createThoughtChatResponse("Let me "), + createThoughtChatResponse("think."), + createStreamingChatResponse("Hello"), + createStreamingChatResponse(" there."))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + List parts = + testObserver.values().get(testObserver.values().size() - 1).content().get().parts().get(); + assertThat(parts) + .extracting(part -> part.text().get()) + .containsExactly("Let me think.", "Hello there."); + assertThat(parts).extracting(part -> part.thought().orElse(false)).containsExactly(true, false); + } + + @Test + void generateContent_streamingMedia_aggregatedResponseKeepsMediaParts() { + AssistantMessage imageMessage = + AssistantMessage.builder() + .content("") + .media( + List.of( + new Media( + MimeType.valueOf("image/png"), new ByteArrayResource(new byte[] {1, 2})))) + .build(); + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn( + Flux.just( + createStreamingChatResponse("Here:"), + new ChatResponse(List.of(new Generation(imageMessage))))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + List parts = + testObserver.values().get(testObserver.values().size() - 1).content().get().parts().get(); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).text()).hasValue("Here:"); + assertThat(parts.get(1).inlineData().flatMap(Blob::mimeType)).hasValue("image/png"); + } + + @Test + void generateContent_streaming_aggregatedResponseCarriesModelVersion() { + ChatResponse lastChunk = + new ChatResponse( + List.of(new Generation(new AssistantMessage(" there"))), + ChatResponseMetadata.builder().model("gemini-test-001").build()); + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn(Flux.just(createStreamingChatResponse("Hi"), lastChunk)); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + LlmResponse aggregated = testObserver.values().get(testObserver.values().size() - 1); + assertThat(aggregated.modelVersion()).hasValue("gemini-test-001"); + } + + @Test + void generateContent_streamingWithoutContent_endsWithResponseCarryingUsage() { + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn(Flux.just(createStreamingChatResponse("", new DefaultUsage(3, 0, 3)))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + testObserver.assertValueCount(1); + LlmResponse response = testObserver.values().get(0); + assertThat(response.partial().orElse(false)).isFalse(); + assertThat(response.turnComplete()).hasValue(true); + assertThat(response.content()).isEmpty(); + assertThat(response.usageMetadata().get().totalTokenCount()).hasValue(3); + } + + @Test + void generateContent_streamingWithoutChunks_emitsNothing() { + when(mockStreamingChatModel.stream(any(Prompt.class))).thenReturn(Flux.empty()); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + testObserver.assertComplete(); + testObserver.assertNoValues(); + } + + @Test + void generateContent_streamingFunctionCallWithoutId_aggregatedReusesPartialId() { + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn(Flux.just(createToolCallChatResponse("", "get_weather"))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + List responses = testObserver.values(); + assertThat(responses).hasSize(2); + FunctionCall partialCall = onlyFunctionCall(responses.get(0)); + FunctionCall aggregatedCall = onlyFunctionCall(responses.get(1)); + assertThat(aggregatedCall.name()).hasValue("get_weather"); + assertThat(aggregatedCall.id().get()).isNotEmpty().isEqualTo(partialCall.id().get()); + } + + @Test + void generateContent_streamingTextAndFunctionCalls_aggregatesTextThenCallsInOrder() { + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn( + Flux.just( + createStreamingChatResponse("Let me check."), + createToolCallChatResponse("call_1", "get_weather"), + createToolCallChatResponse("call_2", "get_time"))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + Content aggregated = testObserver.values().get(3).content().get(); + assertThat(aggregated.role()).hasValue("model"); + assertThat(aggregated.parts().get()) + .extracting( + part -> + part.text().orElseGet(() -> part.functionCall().flatMap(FunctionCall::id).get())) + .containsExactly("Let me check.", "call_1", "call_2"); + } + + @Test + void generateContent_streamingFunctionCallWithId_keepsProviderId() { + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn(Flux.just(createToolCallChatResponse("call_1", "get_weather"))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + assertThat(testObserver.values()) + .extracting(response -> onlyFunctionCall(response).id().get()) + .containsExactly("call_1", "call_1"); + } + + @Test + void generateContent_streamingError_isDeliveredOnce() { + List undeliverable = new ArrayList<>(); + Consumer previousHandler = RxJavaPlugins.getErrorHandler(); + RxJavaPlugins.setErrorHandler(undeliverable::add); + try { + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn(Flux.error(new IllegalStateException("boom"))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + TestSubscriber testObserver = springAI.generateContent(testRequest, true).test(); + + testObserver.awaitDone(5, TimeUnit.SECONDS); + testObserver.assertError(RuntimeException.class); + assertThat(undeliverable).isEmpty(); + } finally { + RxJavaPlugins.setErrorHandler(previousHandler); + } + } + + @Test + void generateContent_streamingCancelled_cancelsUnderlyingStream() { + AtomicBoolean cancelled = new AtomicBoolean(); + when(mockStreamingChatModel.stream(any(Prompt.class))) + .thenReturn(Flux.never().doOnCancel(() -> cancelled.set(true))); + SpringAI springAI = new SpringAI(mockStreamingChatModel); + + springAI.generateContent(testRequest, true).test().cancel(); + + assertThat(cancelled).isTrue(); + } + + private static FunctionCall onlyFunctionCall(LlmResponse response) { + List parts = response.content().get().parts().get(); + assertThat(parts).hasSize(1); + return parts.get(0).functionCall().get(); } private ChatResponse createStreamingChatResponse(String text) { @@ -281,4 +581,27 @@ private ChatResponse createStreamingChatResponse(String text) { Generation generation = new Generation(assistantMessage); return new ChatResponse(List.of(generation)); } + + private static ChatResponse createStreamingChatResponse(String text, DefaultUsage usage) { + return new ChatResponse( + List.of(new Generation(new AssistantMessage(text))), + ChatResponseMetadata.builder().usage(usage).build()); + } + + private static ChatResponse createThoughtChatResponse(String text) { + AssistantMessage message = + AssistantMessage.builder().content(text).properties(Map.of("isThought", true)).build(); + return new ChatResponse(List.of(new Generation(message))); + } + + private static ChatResponse createToolCallChatResponse(String id, String name) { + AssistantMessage message = + AssistantMessage.builder() + .content("") + .toolCalls( + List.of( + new AssistantMessage.ToolCall(id, "function", name, "{\"city\":\"Paris\"}"))) + .build(); + return new ChatResponse(List.of(new Generation(message))); + } }