diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/streaming/StreamingPromptRunner.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/streaming/StreamingPromptRunner.kt index 498e275ee..afe0c67b9 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/streaming/StreamingPromptRunner.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/streaming/StreamingPromptRunner.kt @@ -99,6 +99,17 @@ interface StreamingPromptRunner : PromptRunner { */ fun generateStream(): Flux + /** + * Generate a reactive stream containing text and thinking events. + * + * The default preserves compatibility for implementations that only + * provide text chunks. + * + * @return Flux emitting thinking and text events in response order + */ + fun generateStreamWithThinking(): Flux> = + generateStream().map { StreamingEvent.Object(it) } + } } diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreaming.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreaming.kt index f3df5b51c..e5dd1a4e7 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreaming.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreaming.kt @@ -38,6 +38,9 @@ internal data class DelegatingStreaming( return delegate.generateStream() } + override fun generateStreamWithThinking(): Flux> = + delegate.generateStreamWithThinking() + override fun createObjectStream(itemClass: Class): Flux = delegate.createObjectStream(itemClass) diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt index d353bb541..70871b8a5 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt @@ -341,6 +341,17 @@ internal data class OperationContextDelegate( ) } + override fun generateStreamWithThinking(): Flux> { + val streamingLlmOperations = streamingFactory().createStreamingOperations(llm) + + return streamingLlmOperations.generateStreamWithThinking( + messages = messages, + interaction = streamingInteraction(), + agentProcess = context.processContext.agentProcess, + action = action, + ) + } + override fun createObjectStream(itemClass: Class): Flux { val streamingLlmOperations = streamingFactory().createStreamingOperations(llm) diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt index 07005ffc3..9c46145f8 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt @@ -139,6 +139,9 @@ internal interface PromptExecutionDelegate : LlmUse { fun generateStream(): Flux + fun generateStreamWithThinking(): Flux> = + generateStream().map { StreamingEvent.Object(it) } + fun createObjectStream(itemClass: Class): Flux fun createObjectStreamWithThinking(itemClass: Class): Flux> diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/core/internal/streaming/StreamingLlmOperations.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/core/internal/streaming/StreamingLlmOperations.kt index 61d07702a..c114143a9 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/core/internal/streaming/StreamingLlmOperations.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/core/internal/streaming/StreamingLlmOperations.kt @@ -57,6 +57,21 @@ interface StreamingLlmOperations { action: Action?, ): Flux + /** + * Generate text and thinking events from messages. + * + * The default wraps the existing text-only stream so third-party + * implementations remain source compatible. + */ + fun generateStreamWithThinking( + messages: List, + interaction: LlmInteraction, + agentProcess: AgentProcess, + action: Action?, + ): Flux> = + generateStream(messages, interaction, agentProcess, action) + .map { StreamingEvent.Object(it) } + /** * Create a streaming list of objects from JSONL response in the context of an AgentProcess. * Each line in the LLM response should be a valid JSON object matching the output class. @@ -142,6 +157,21 @@ interface StreamingLlmOperations { action: Action? = null, ): Flux + /** + * Low-level text and thinking stream with optional platform context. + * + * The default wraps [doTransformStream] for source compatibility. + */ + fun doTransformStreamWithThinking( + messages: List, + interaction: LlmInteraction, + llmRequestEvent: LlmRequestEvent?, + agentProcess: AgentProcess? = null, + action: Action? = null, + ): Flux> = + doTransformStream(messages, interaction, llmRequestEvent, agentProcess, action) + .map { StreamingEvent.Object(it) } + /** * Low level object streaming transform with optional platform context. * Streams typed objects as they are parsed from JSONL response. diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperations.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperations.kt index 734edbe86..da200dc93 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperations.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperations.kt @@ -31,6 +31,7 @@ import com.embabel.agent.spi.support.springai.ChatClientLlmOperations import com.embabel.agent.spi.support.springai.SpringAiLlmService import com.embabel.agent.spi.support.springai.toSpringAiMessage import com.embabel.agent.spi.support.springai.toSpringToolCallbacks +import com.embabel.agent.spi.support.streaming.toTaggedThinkingEvents import com.embabel.chat.Message import com.embabel.common.ai.converters.streaming.StreamingJacksonOutputConverter import com.embabel.common.core.streaming.StreamingEvent @@ -113,6 +114,14 @@ internal class StreamingChatClientOperations( return doTransformStream(messages, interaction, null, agentProcess, action) } + override fun generateStreamWithThinking( + messages: List, + interaction: LlmInteraction, + agentProcess: AgentProcess, + action: Action?, + ): Flux> = + doTransformStreamWithThinking(messages, interaction, null, agentProcess, action) + override fun createObjectStream( messages: List, interaction: LlmInteraction, @@ -190,6 +199,16 @@ internal class StreamingChatClientOperations( ) } + override fun doTransformStreamWithThinking( + messages: List, + interaction: LlmInteraction, + llmRequestEvent: LlmRequestEvent?, + agentProcess: AgentProcess?, + action: Action?, + ): Flux> = + doTransformStream(messages, interaction, llmRequestEvent, agentProcess, action) + .toTaggedThinkingEvents() + /** * Creates a stream of typed objects from LLM JSONL responses, with thinking content suppressed. * diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImpl.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImpl.kt index 417544415..946709244 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImpl.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImpl.kt @@ -75,6 +75,14 @@ internal class StreamingLlmOperationsImpl( return doTransformStream(messages, interaction, null, agentProcess, action) } + override fun generateStreamWithThinking( + messages: List, + interaction: LlmInteraction, + agentProcess: AgentProcess, + action: Action?, + ): Flux> = + doTransformStreamWithThinking(messages, interaction, null, agentProcess, action) + override fun createObjectStream( messages: List, interaction: LlmInteraction, @@ -137,6 +145,16 @@ internal class StreamingLlmOperationsImpl( return messageStreamer.stream(messagesWithContributions, tools, interaction.toolCallInspectors) } + override fun doTransformStreamWithThinking( + messages: List, + interaction: LlmInteraction, + llmRequestEvent: LlmRequestEvent?, + agentProcess: AgentProcess?, + action: Action?, + ): Flux> = + doTransformStream(messages, interaction, llmRequestEvent, agentProcess, action) + .toTaggedThinkingEvents() + override fun doTransformObjectStream( messages: List, interaction: LlmInteraction, diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/ThinkingStreamSupport.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/ThinkingStreamSupport.kt new file mode 100644 index 000000000..4d52625c8 --- /dev/null +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/ThinkingStreamSupport.kt @@ -0,0 +1,176 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.agent.spi.support.streaming + +import com.embabel.common.core.streaming.StreamingEvent +import com.embabel.common.core.thinking.ThinkingTags +import reactor.core.publisher.Flux + +/** + * Extracts tagged thinking from text while preserving response order across + * arbitrary chunk boundaries. Parser state is scoped to each subscription. + */ +internal fun Flux.toTaggedThinkingEvents(): Flux> = + Flux.defer { + val parser = TaggedThinkingParser() + this@toTaggedThinkingEvents + .concatMap { Flux.fromIterable(parser.accept(it)) } + .concatWith(Flux.defer { Flux.fromIterable(parser.finish()) }) + } + +private class TaggedThinkingParser { + + private data class Tag( + val start: String, + val end: String, + ) + + private data class StartMatch( + val index: Int, + val tag: Tag?, + val legacy: Boolean, + ) + + private val tags = ThinkingTags.TAG_DEFINITIONS + .filterKeys { it != "legacy_prefix" && it != "no_prefix" } + .values + .map { Tag(it.first, it.second) } + + private val legacyPrefix = ThinkingTags.TAG_DEFINITIONS["legacy_prefix"]?.first.orEmpty() + private val buffer = StringBuilder() + private var activeTag: Tag? = null + private var atLineStart = true + + fun accept(text: String): List> { + buffer.append(text) + val events = mutableListOf>() + + while (buffer.isNotEmpty()) { + val tag = activeTag + if (tag != null) { + val endIndex = buffer.indexOf(tag.end) + if (endIndex < 0) break + val thinking = buffer.substring(0, endIndex).trim() + if (thinking.isNotEmpty()) events += StreamingEvent.Thinking(thinking) + buffer.delete(0, endIndex + tag.end.length) + activeTag = null + atLineStart = false + continue + } + + if (atLineStart && legacyPrefix.isNotEmpty() && buffer.startsWith(legacyPrefix)) { + val newlineIndex = buffer.indexOf("\n") + if (newlineIndex < 0) break + val thinking = buffer.substring(legacyPrefix.length, newlineIndex).trim() + if (thinking.isNotEmpty()) events += StreamingEvent.Thinking(thinking) + buffer.delete(0, newlineIndex + 1) + atLineStart = true + continue + } + + val match = findStart() + if (match != null) { + if (match.index > 0) { + emitText(events, buffer.substring(0, match.index)) + buffer.delete(0, match.index) + continue + } + if (match.legacy) break + activeTag = match.tag + buffer.delete(0, match.tag!!.start.length) + atLineStart = false + continue + } + + val heldSuffixLength = longestPartialStartSuffix() + val emitLength = buffer.length - heldSuffixLength + if (emitLength > 0) { + emitText(events, buffer.substring(0, emitLength)) + buffer.delete(0, emitLength) + } + break + } + + return events + } + + fun finish(): List> { + val events = mutableListOf>() + val tag = activeTag + if (tag != null) { + emitText(events, tag.start + buffer.toString()) + } else if (atLineStart && legacyPrefix.isNotEmpty() && buffer.startsWith(legacyPrefix)) { + val thinking = buffer.substring(legacyPrefix.length).trim() + if (thinking.isNotEmpty()) events += StreamingEvent.Thinking(thinking) + } else { + emitText(events, buffer.toString()) + } + buffer.clear() + activeTag = null + atLineStart = true + return events + } + + private fun findStart(): StartMatch? { + val tagMatch = tags + .mapNotNull { tag -> buffer.indexOf(tag.start).takeIf { it >= 0 }?.let { StartMatch(it, tag, false) } } + .minWithOrNull(compareBy { it.index }.thenByDescending { it.tag?.start?.length ?: 0 }) + val legacyMatch = if (legacyPrefix.isNotEmpty()) { + findLegacyStart()?.let { StartMatch(it, null, true) } + } else { + null + } + return listOfNotNull(tagMatch, legacyMatch).minByOrNull { it.index } + } + + private fun findLegacyStart(): Int? { + var fromIndex = 0 + while (fromIndex < buffer.length) { + val index = buffer.indexOf(legacyPrefix, fromIndex) + if (index < 0) return null + if ((index == 0 && atLineStart) || (index > 0 && buffer[index - 1] == '\n')) return index + fromIndex = index + 1 + } + return null + } + + private fun longestPartialStartSuffix(): Int { + val tagSuffixLength = tags.map { it.start }.maxOfOrNull { start -> + (1 until start.length) + .filter { length -> buffer.endsWith(start.substring(0, length)) } + .maxOrNull() ?: 0 + } ?: 0 + return maxOf(tagSuffixLength, longestPartialLegacySuffix()) + } + + private fun longestPartialLegacySuffix(): Int { + if (legacyPrefix.isEmpty()) return 0 + return (1 until legacyPrefix.length) + .filter { length -> + if (!buffer.endsWith(legacyPrefix.substring(0, length))) return@filter false + val startIndex = buffer.length - length + (startIndex == 0 && atLineStart) || (startIndex > 0 && buffer[startIndex - 1] == '\n') + } + .maxOrNull() ?: 0 + } + + private fun emitText(events: MutableList>, text: String) { + if (text.isNotEmpty()) { + events += StreamingEvent.Object(text) + atLineStart = text.endsWith("\n") + } + } +} diff --git a/embabel-agent-api/src/test/java/com/embabel/agent/api/streaming/StreamingThinkingJavaApiTest.java b/embabel-agent-api/src/test/java/com/embabel/agent/api/streaming/StreamingThinkingJavaApiTest.java new file mode 100644 index 000000000..41a16f34e --- /dev/null +++ b/embabel-agent-api/src/test/java/com/embabel/agent/api/streaming/StreamingThinkingJavaApiTest.java @@ -0,0 +1,38 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.agent.api.streaming; + +import com.embabel.agent.api.common.streaming.StreamingPromptRunner; +import com.embabel.common.core.streaming.StreamingEvent; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; + +import static org.assertj.core.api.Assertions.assertThat; + +class StreamingThinkingJavaApiTest { + + @Test + void exposesThinkingAwareTextStreamingToJava() throws NoSuchMethodException { + var method = StreamingPromptRunner.Streaming.class.getMethod("generateStreamWithThinking"); + + assertThat(method.getReturnType()).isEqualTo(Flux.class); + assertThat(method.isDefault()).isTrue(); + } + + Flux> callFromJava(StreamingPromptRunner.Streaming streaming) { + return streaming.generateStreamWithThinking(); + } +} diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingTest.kt index d23dca73c..189b3d8c5 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingTest.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingTest.kt @@ -98,6 +98,24 @@ class DelegatingStreamingTest { } } + @Nested + inner class GenerateStreamWithThinkingTest { + + @Test + fun `should delegate to delegate generateStreamWithThinking`() { + val mockStream = Flux.just( + StreamingEvent.Thinking("Reasoning"), + StreamingEvent.Object("Answer"), + ) + every { mockDelegate.generateStreamWithThinking() } returns mockStream + + val result = createStreamingOperations().generateStreamWithThinking() + + verify { mockDelegate.generateStreamWithThinking() } + assertEquals(mockStream.collectList().block(), result.collectList().block()) + } + } + @Nested inner class CreateObjectStreamTest { diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/streaming/StreamingImpl.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/streaming/StreamingImpl.kt index 9ba00b65a..c41e18474 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/streaming/StreamingImpl.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/streaming/StreamingImpl.kt @@ -69,6 +69,15 @@ internal class StreamingImpl( ) } + override fun generateStreamWithThinking(): Flux> { + return streamingLlmOperations.generateStreamWithThinking( + messages = messages, + interaction = interaction, + agentProcess = agentProcess, + action = action, + ) + } + override fun createObjectStream(itemClass: Class): Flux { return streamingLlmOperations.createObjectStream( messages = messages, diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/streaming/StreamingImplTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/streaming/StreamingImplTest.kt index fd1b2dc89..82012bc2c 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/streaming/StreamingImplTest.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/streaming/StreamingImplTest.kt @@ -104,6 +104,28 @@ class StreamingImplTest { assertEquals("test", firstItem) } + @Test + fun `should delegate generateStreamWithThinking to StreamingLlmOperations`() { + val mockStream = Flux.just( + StreamingEvent.Thinking("reasoning"), + StreamingEvent.Object("answer"), + ) + every { + mockStreamingLlmOperations.generateStreamWithThinking( + eq(initialMessages), any(), mockAgentProcess, mockAction + ) + } returns mockStream + + val result = streamingOperations.generateStreamWithThinking() + + verify { + mockStreamingLlmOperations.generateStreamWithThinking( + initialMessages, mockInteraction, mockAgentProcess, mockAction + ) + } + assertEquals(mockStream.collectList().block(), result.collectList().block()) + } + @Test fun `should delegate createObjectStream to StreamingLlmOperations`() { // Given diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperationsTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperationsTest.kt index 01155fbe8..7292d65ae 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperationsTest.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperationsTest.kt @@ -355,6 +355,25 @@ class StreamingChatClientOperationsTest { .verify(Duration.ofSeconds(1)) } + @Test + fun `generateStreamWithThinking extracts tagged content from the text stream`() { + mockChatClientForStreaming( + Flux.just("reasoninganswer") + ) + + val result = streamingOperations.generateStreamWithThinking( + messages = listOf(UserMessage("test")), + interaction = mockInteraction, + agentProcess = mockAgentProcess, + action = mockAction, + ) + + StepVerifier.create(result) + .expectNextMatches { it.isThinking() && it.getThinking() == "reasoning" } + .expectNextMatches { it.isObject() && it.getObject() == "answer" } + .verifyComplete() + } + @Test fun `should handle real streaming with reactive callbacks`() { // Given: Mixed content with multiple events diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsThinkingTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsThinkingTest.kt new file mode 100644 index 000000000..6361db9c5 --- /dev/null +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsThinkingTest.kt @@ -0,0 +1,149 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.agent.spi.support.streaming + +import com.embabel.agent.api.common.InteractionId +import com.embabel.agent.core.support.LlmInteraction +import com.embabel.agent.spi.LlmService +import com.embabel.agent.spi.ToolDecorator +import com.embabel.agent.spi.loop.streaming.LlmMessageStreamer +import com.embabel.chat.UserMessage +import com.embabel.common.core.streaming.StreamingEvent +import com.embabel.common.core.thinking.ThinkingTags +import io.mockk.every +import io.mockk.mockk +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test +import reactor.core.publisher.Flux +import tools.jackson.module.kotlin.jacksonObjectMapper + +class StreamingLlmOperationsThinkingTest { + + @Test + fun `extracts every supported tag at every chunk boundary`() { + ThinkingTags.TAG_DEFINITIONS + .filterKeys { it != "legacy_prefix" && it != "no_prefix" } + .forEach { (name, tags) -> + val source = "${tags.first}reasoning${tags.second}answer" + (0..source.length).forEach { split -> + val events = parse(source.substring(0, split), source.substring(split)) + assertThat(events.filterIsInstance().map { it.content }) + .describedAs("$name thinking split at $split") + .containsExactly("reasoning") + assertThat(events.text()) + .describedAs("$name text split at $split") + .isEqualTo("answer") + } + } + } + + @Test + fun `extracts legacy prefix at every chunk boundary`() { + val source = "//THINKING: reasoning\nanswer" + + (0..source.length).forEach { split -> + val events = parse(source.substring(0, split), source.substring(split)) + assertThat(events.filterIsInstance().map { it.content }) + .describedAs("legacy thinking split at $split") + .containsExactly("reasoning") + assertThat(events.text()) + .describedAs("legacy text split at $split") + .isEqualTo("answer") + } + } + + @Test + fun `extracts tagged thinking across chunk boundaries`() { + val events = collect("tagged rea", "soningfinal ", "answer") + + assertThat(events.filterIsInstance().map { it.content }) + .containsExactly("tagged reasoning") + assertThat(events.text()).isEqualTo("final answer") + } + + @Test + fun `extracts legacy thinking prefix across chunk boundaries`() { + val events = collect("//THINK", "ING: legacy rea", "soning\nanswer") + + assertThat(events).containsExactly( + StreamingEvent.Thinking("legacy reasoning"), + StreamingEvent.Object("answer"), + ) + } + + @Test + fun `keeps a legacy marker in the middle of a line as text`() { + val events = collect("answer //THINK", "ING: literal") + + assertThat(events.filterIsInstance()).isEmpty() + assertThat(events.text()).isEqualTo("answer //THINKING: literal") + } + + @Test + fun `keeps an unclosed thinking tag as text`() { + val events = collect("before unfinished reasoning") + + assertThat(events.filterIsInstance()).isEmpty() + assertThat(events.text()).isEqualTo("before unfinished reasoning") + } + + @Test + fun `keeps parser state isolated for repeated subscriptions`() { + val stream = operations("reasoninganswer") + .doTransformStreamWithThinking( + messages = listOf(UserMessage("question")), + interaction = interaction(), + llmRequestEvent = null, + ) + val expected = listOf( + StreamingEvent.Thinking("reasoning"), + StreamingEvent.Object("answer"), + ) + + assertThat(stream.collectList().block()).containsExactlyElementsOf(expected) + assertThat(stream.collectList().block()).containsExactlyElementsOf(expected) + } + + private fun collect(vararg chunks: String): List> = + operations(*chunks) + .doTransformStreamWithThinking( + messages = listOf(UserMessage("question")), + interaction = interaction(), + llmRequestEvent = null, + ) + .collectList() + .block()!! + + private fun parse(vararg chunks: String): List> = + Flux.fromArray(chunks).toTaggedThinkingEvents().collectList().block()!! + + private fun operations(vararg chunks: String): StreamingLlmOperationsImpl { + val llmService = mockk> { + every { promptContributors } returns emptyList() + } + return StreamingLlmOperationsImpl( + messageStreamer = LlmMessageStreamer { _, _, _ -> Flux.fromArray(chunks) }, + objectMapper = jacksonObjectMapper(), + llmService = llmService, + toolDecorator = mockk(), + ) + } + + private fun interaction() = LlmInteraction(id = InteractionId("thinking-stream")) + + private fun List>.text(): String = + filterIsInstance>().joinToString("") { it.item } +}