diff --git a/koog-spring-ai-v2/koog-spring-ai-v2-starter-model-chat/src/main/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClient.kt b/koog-spring-ai-v2/koog-spring-ai-v2-starter-model-chat/src/main/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClient.kt index cf9d561652..1332715e67 100644 --- a/koog-spring-ai-v2/koog-spring-ai-v2-starter-model-chat/src/main/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClient.kt +++ b/koog-spring-ai-v2/koog-spring-ai-v2-starter-model-chat/src/main/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClient.kt @@ -185,6 +185,9 @@ public class SpringAiLLMClient( * [StreamFrame.ToolCallDelta] with a corresponding [StreamFrame.ToolCallComplete] and * emits [StreamFrame.TextComplete] / [StreamFrame.ReasoningComplete] boundaries. * + * The terminal [StreamFrame.End] carries the finish reason and the token usage collected over + * the whole stream, since a provider may report them in separate chunks. + * * All blocking I/O runs on the configured [dispatcher] (default [Dispatchers.IO]). */ override fun executeStreaming( @@ -201,10 +204,10 @@ public class SpringAiLLMClient( } catch (e: Exception) { throw LLMClientException(clientName, "ChatModel.stream() failed: ${e.message}", e) } - var lastChatResponse: ChatResponse? = null + var finishReason: String? = null + var metaInfo: ResponseMetaInfo? = null try { flux.asFlow().collect { chatResponse -> - lastChatResponse = chatResponse for ((generationIndex, generation) in chatResponse.results.withIndex()) { val assistantMessage = generation.output val text = assistantMessage.text @@ -222,6 +225,23 @@ public class SpringAiLLMClient( toolCallAssembler.accept(assistantMessage.toolCalls, generationIndex, this) } } + // Spring AI can report the finish reason and the token usage in different chunks: + // with stream-usage the last chunk carries the usage and no generations at all, so + // keep each as it arrives (see #2109). A chunk that has neither reports an empty + // finish reason and zero counts rather than nulls, so neither may overwrite a + // reported value. + chatResponse.results.firstOrNull()?.metadata?.finishReason + ?.takeIf { it.isNotEmpty() } + ?.let { finishReason = it } + val usage = chatResponse.metadata.usage + if (metaInfo == null || usage.totalTokens > 0) { + metaInfo = ResponseMetaInfo.create( + clock = clock, + totalTokensCount = usage.totalTokens, + inputTokensCount = usage.promptTokens, + outputTokensCount = usage.completionTokens + ) + } } toolCallAssembler.flush(this) } catch (e: CancellationException) { @@ -232,18 +252,6 @@ public class SpringAiLLMClient( throw LLMClientException(clientName, "ChatModel.stream() failed during collection: ${e.message}", e) } finally { // Always emit End frame so downstream consumers are not left hanging - val finishReason = lastChatResponse?.results?.firstOrNull()?.metadata?.finishReason - val usage = lastChatResponse?.metadata?.usage - val metaInfo = if (usage != null) { - ResponseMetaInfo.create( - clock = clock, - totalTokensCount = usage.totalTokens, - inputTokensCount = usage.promptTokens, - outputTokensCount = usage.completionTokens - ) - } else { - null - } emitEnd(finishReason = finishReason, metaInfo = metaInfo) } }.flowOn(dispatcher) diff --git a/koog-spring-ai-v2/koog-spring-ai-v2-starter-model-chat/src/test/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClientTest.kt b/koog-spring-ai-v2/koog-spring-ai-v2-starter-model-chat/src/test/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClientTest.kt index 7b785cf3df..5c5f5ed162 100644 --- a/koog-spring-ai-v2/koog-spring-ai-v2-starter-model-chat/src/test/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClientTest.kt +++ b/koog-spring-ai-v2/koog-spring-ai-v2-starter-model-chat/src/test/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClientTest.kt @@ -21,6 +21,7 @@ import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test import org.junit.jupiter.api.assertThrows 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.Usage import org.springframework.ai.chat.model.ChatModel @@ -45,6 +46,12 @@ class SpringAiLLMClientTest { private fun requestMeta() = RequestMetaInfo.create(KoogClock.System) + private fun stubUsage(promptTokens: Int, completionTokens: Int) = object : Usage { + override fun getPromptTokens(): Int = promptTokens + override fun getCompletionTokens(): Int = completionTokens + override fun getNativeUsage(): Any = emptyMap() + } + // ---- llmProvider ---- @Test @@ -153,12 +160,9 @@ class SpringAiLLMClientTest { @Test fun `execute maps usage metadata to response meta info`() = runBlocking { - val usage = object : Usage { - override fun getPromptTokens(): Int = 5 - override fun getCompletionTokens(): Int = 15 - override fun getNativeUsage(): Any = emptyMap() - } - val metadata = ChatResponseMetadata.builder().usage(usage).build() + val metadata = ChatResponseMetadata.builder() + .usage(stubUsage(promptTokens = 5, completionTokens = 15)) + .build() val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { override fun call(prompt: SpringPrompt) = ChatResponse(listOf(Generation(AssistantMessage("Done"))), metadata) @@ -340,6 +344,86 @@ class SpringAiLLMClientTest { assertEquals("Hi", textFrames[0].text) } + @Test + fun testExecuteStreamingKeepsFinishReasonWhenLastChunkCarriesOnlyUsage() = runBlocking { + val finishMetadata = ChatGenerationMetadata.builder().finishReason("STOP").build() + val usageMetadata = ChatResponseMetadata.builder() + .usage(stubUsage(promptTokens = 5, completionTokens = 10)) + .build() + val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { + override fun call(prompt: SpringPrompt) = throw UnsupportedOperationException() + override fun stream(prompt: SpringPrompt) = Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("Hello"), finishMetadata))), + // With stream-usage enabled the terminal chunk reports the usage and no generations + ChatResponse(emptyList(), usageMetadata) + ) + }).build() + val prompt = createPrompt(Message.User("Hi", requestMeta())) + val frames = client.executeStreaming(prompt, testModel, emptyList()).toList() + + val end = frames.filterIsInstance().single() + assertEquals("STOP", end.finishReason) + assertEquals(15, end.metaInfo.totalTokensCount) + } + + @Test + fun testExecuteStreamingKeepsTokenUsageWhenALaterChunkReportsNone() = runBlocking { + val usageMetadata = ChatResponseMetadata.builder() + .usage(stubUsage(promptTokens = 5, completionTokens = 10)) + .build() + val finishMetadata = ChatGenerationMetadata.builder().finishReason("STOP").build() + val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { + override fun call(prompt: SpringPrompt) = throw UnsupportedOperationException() + override fun stream(prompt: SpringPrompt) = Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("Hello"))), usageMetadata), + ChatResponse(listOf(Generation(AssistantMessage(" world"), finishMetadata))) + ) + }).build() + val prompt = createPrompt(Message.User("Hi", requestMeta())) + val frames = client.executeStreaming(prompt, testModel, emptyList()).toList() + + val end = frames.filterIsInstance().single() + assertEquals("STOP", end.finishReason) + assertEquals(5, end.metaInfo.inputTokensCount) + assertEquals(10, end.metaInfo.outputTokensCount) + assertEquals(15, end.metaInfo.totalTokensCount) + } + + @Test + fun testExecuteStreamingReportsZeroTokenCountsWhenNoChunkReportsUsage() = runBlocking { + val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { + override fun call(prompt: SpringPrompt) = throw UnsupportedOperationException() + override fun stream(prompt: SpringPrompt) = Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("Hello")))) + ) + }).build() + val prompt = createPrompt(Message.User("Hi", requestMeta())) + val frames = client.executeStreaming(prompt, testModel, emptyList()).toList() + + // A response without usage reports zero counts, not unknown ones, as it did before + val end = frames.filterIsInstance().single() + assertEquals(0, end.metaInfo.totalTokensCount) + } + + @Test + fun testExecuteStreamingIgnoresAnEmptyFinishReasonFromALaterChunk() = runBlocking { + val finishMetadata = ChatGenerationMetadata.builder().finishReason("STOP").build() + // Mistral AI and DeepSeek pad every streamed chunk with an empty finish reason + val paddedMetadata = ChatGenerationMetadata.builder().finishReason("").build() + val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { + override fun call(prompt: SpringPrompt) = throw UnsupportedOperationException() + override fun stream(prompt: SpringPrompt) = Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("Hello"), finishMetadata))), + ChatResponse(listOf(Generation(AssistantMessage(" world"), paddedMetadata))) + ) + }).build() + val prompt = createPrompt(Message.User("Hi", requestMeta())) + val frames = client.executeStreaming(prompt, testModel, emptyList()).toList() + + val end = frames.filterIsInstance().single() + assertEquals("STOP", end.finishReason) + } + // ---- moderate ---- @Test diff --git a/koog-spring-ai/koog-spring-ai-starter-model-chat/src/main/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClient.kt b/koog-spring-ai/koog-spring-ai-starter-model-chat/src/main/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClient.kt index e8c9b49fe7..817f0501c6 100644 --- a/koog-spring-ai/koog-spring-ai-starter-model-chat/src/main/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClient.kt +++ b/koog-spring-ai/koog-spring-ai-starter-model-chat/src/main/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClient.kt @@ -184,6 +184,9 @@ public class SpringAiLLMClient( * [StreamFrame.ToolCallDelta] with a corresponding [StreamFrame.ToolCallComplete] and * emits [StreamFrame.TextComplete] / [StreamFrame.ReasoningComplete] boundaries. * + * The terminal [StreamFrame.End] carries the finish reason and the token usage collected over + * the whole stream, since a provider may report them in separate chunks. + * * All blocking I/O runs on the configured [dispatcher] (default [Dispatchers.IO]). */ override fun executeStreaming( @@ -200,10 +203,10 @@ public class SpringAiLLMClient( } catch (e: Exception) { throw LLMClientException(clientName, "ChatModel.stream() failed: ${e.message}", e) } - var lastChatResponse: ChatResponse? = null + var finishReason: String? = null + var metaInfo: ResponseMetaInfo? = null try { flux.asFlow().collect { chatResponse -> - lastChatResponse = chatResponse for ((generationIndex, generation) in chatResponse.results.withIndex()) { val assistantMessage = generation.output val text = assistantMessage.text @@ -221,6 +224,23 @@ public class SpringAiLLMClient( toolCallAssembler.accept(assistantMessage.toolCalls, generationIndex, this) } } + // Spring AI can report the finish reason and the token usage in different chunks: + // with stream-usage the last chunk carries the usage and no generations at all, so + // keep each as it arrives (see #2109). A chunk that has neither reports an empty + // finish reason and zero counts rather than nulls, so neither may overwrite a + // reported value. + chatResponse.results.firstOrNull()?.metadata?.finishReason + ?.takeIf { it.isNotEmpty() } + ?.let { finishReason = it } + val usage = chatResponse.metadata.usage + if (metaInfo == null || usage.totalTokens > 0) { + metaInfo = ResponseMetaInfo.create( + clock = clock, + totalTokensCount = usage.totalTokens, + inputTokensCount = usage.promptTokens, + outputTokensCount = usage.completionTokens + ) + } } toolCallAssembler.flush(this) } catch (e: CancellationException) { @@ -231,18 +251,6 @@ public class SpringAiLLMClient( throw LLMClientException(clientName, "ChatModel.stream() failed during collection: ${e.message}", e) } finally { // Always emit End frame so downstream consumers are not left hanging - val finishReason = lastChatResponse?.results?.firstOrNull()?.metadata?.finishReason - val usage = lastChatResponse?.metadata?.usage - val metaInfo = if (usage != null) { - ResponseMetaInfo.create( - clock = clock, - totalTokensCount = usage.totalTokens, - inputTokensCount = usage.promptTokens, - outputTokensCount = usage.completionTokens - ) - } else { - null - } emitEnd(finishReason = finishReason, metaInfo = metaInfo) } }.flowOn(dispatcher) diff --git a/koog-spring-ai/koog-spring-ai-starter-model-chat/src/test/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClientTest.kt b/koog-spring-ai/koog-spring-ai-starter-model-chat/src/test/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClientTest.kt index 8633277bad..4ceafe2916 100644 --- a/koog-spring-ai/koog-spring-ai-starter-model-chat/src/test/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClientTest.kt +++ b/koog-spring-ai/koog-spring-ai-starter-model-chat/src/test/kotlin/ai/koog/spring/ai/chat/SpringAiLLMClientTest.kt @@ -21,6 +21,7 @@ import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test import org.junit.jupiter.api.assertThrows 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.Usage import org.springframework.ai.chat.model.ChatModel @@ -44,6 +45,12 @@ class SpringAiLLMClientTest { private fun requestMeta() = RequestMetaInfo.create(KoogClock.System) + private fun stubUsage(promptTokens: Int, completionTokens: Int) = object : Usage { + override fun getPromptTokens(): Int = promptTokens + override fun getCompletionTokens(): Int = completionTokens + override fun getNativeUsage(): Any = emptyMap() + } + // ---- llmProvider ---- @Test @@ -152,12 +159,9 @@ class SpringAiLLMClientTest { @Test fun `execute maps usage metadata to response meta info`() = runBlocking { - val usage = object : Usage { - override fun getPromptTokens(): Int = 5 - override fun getCompletionTokens(): Int = 15 - override fun getNativeUsage(): Any = emptyMap() - } - val metadata = ChatResponseMetadata.builder().usage(usage).build() + val metadata = ChatResponseMetadata.builder() + .usage(stubUsage(promptTokens = 5, completionTokens = 15)) + .build() val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { override fun call(prompt: SpringPrompt) = ChatResponse(listOf(Generation(AssistantMessage("Done"))), metadata) @@ -336,6 +340,86 @@ class SpringAiLLMClientTest { assertEquals("Hi", textFrames[0].text) } + @Test + fun testExecuteStreamingKeepsFinishReasonWhenLastChunkCarriesOnlyUsage() = runBlocking { + val finishMetadata = ChatGenerationMetadata.builder().finishReason("STOP").build() + val usageMetadata = ChatResponseMetadata.builder() + .usage(stubUsage(promptTokens = 5, completionTokens = 10)) + .build() + val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { + override fun call(prompt: SpringPrompt) = throw UnsupportedOperationException() + override fun stream(prompt: SpringPrompt) = Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("Hello"), finishMetadata))), + // With stream-usage enabled the terminal chunk reports the usage and no generations + ChatResponse(emptyList(), usageMetadata) + ) + }).build() + val prompt = createPrompt(Message.User("Hi", requestMeta())) + val frames = client.executeStreaming(prompt, testModel, emptyList()).toList() + + val end = frames.filterIsInstance().single() + assertEquals("STOP", end.finishReason) + assertEquals(15, end.metaInfo.totalTokensCount) + } + + @Test + fun testExecuteStreamingKeepsTokenUsageWhenALaterChunkReportsNone() = runBlocking { + val usageMetadata = ChatResponseMetadata.builder() + .usage(stubUsage(promptTokens = 5, completionTokens = 10)) + .build() + val finishMetadata = ChatGenerationMetadata.builder().finishReason("STOP").build() + val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { + override fun call(prompt: SpringPrompt) = throw UnsupportedOperationException() + override fun stream(prompt: SpringPrompt) = Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("Hello"))), usageMetadata), + ChatResponse(listOf(Generation(AssistantMessage(" world"), finishMetadata))) + ) + }).build() + val prompt = createPrompt(Message.User("Hi", requestMeta())) + val frames = client.executeStreaming(prompt, testModel, emptyList()).toList() + + val end = frames.filterIsInstance().single() + assertEquals("STOP", end.finishReason) + assertEquals(5, end.metaInfo.inputTokensCount) + assertEquals(10, end.metaInfo.outputTokensCount) + assertEquals(15, end.metaInfo.totalTokensCount) + } + + @Test + fun testExecuteStreamingReportsZeroTokenCountsWhenNoChunkReportsUsage() = runBlocking { + val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { + override fun call(prompt: SpringPrompt) = throw UnsupportedOperationException() + override fun stream(prompt: SpringPrompt) = Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("Hello")))) + ) + }).build() + val prompt = createPrompt(Message.User("Hi", requestMeta())) + val frames = client.executeStreaming(prompt, testModel, emptyList()).toList() + + // A response without usage reports zero counts, not unknown ones, as it did before + val end = frames.filterIsInstance().single() + assertEquals(0, end.metaInfo.totalTokensCount) + } + + @Test + fun testExecuteStreamingIgnoresAnEmptyFinishReasonFromALaterChunk() = runBlocking { + val finishMetadata = ChatGenerationMetadata.builder().finishReason("STOP").build() + // Mistral AI and DeepSeek pad every streamed chunk with an empty finish reason + val paddedMetadata = ChatGenerationMetadata.builder().finishReason("").build() + val client = SpringAiLLMClient.builder().chatModel(object : ChatModel { + override fun call(prompt: SpringPrompt) = throw UnsupportedOperationException() + override fun stream(prompt: SpringPrompt) = Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("Hello"), finishMetadata))), + ChatResponse(listOf(Generation(AssistantMessage(" world"), paddedMetadata))) + ) + }).build() + val prompt = createPrompt(Message.User("Hi", requestMeta())) + val frames = client.executeStreaming(prompt, testModel, emptyList()).toList() + + val end = frames.filterIsInstance().single() + assertEquals("STOP", end.finishReason) + } + // ---- moderate ---- @Test