Generate Stream with Thinking

embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImpl.kt

@@ -37,6 +37,7 @@ import com.embabel.agent.spi.support.guardrails.validateUserInput
 import com.embabel.chat.Message
 import com.embabel.chat.UserMessage
 import com.embabel.common.ai.converters.streaming.StreamingJacksonOutputConverter
+import com.embabel.common.ai.converters.streaming.StreamingLineClassifier
 import com.embabel.common.core.streaming.StreamingEvent
 import tools.jackson.databind.ObjectMapper
 import org.slf4j.LoggerFactory
@@ -80,6 +81,43 @@ internal class StreamingLlmOperationsImpl(
         return doTransformStream(messages, interaction, null, agentProcess, action)
     }
 
+    fun generateStreamWithThinking(
+        messages: List<Message>,
+        interaction: LlmInteraction,
+        agentProcess: AgentProcess,
+        action: Action?,
+    ): Flux<StreamingEvent<String>> =
+        doTransformStream(messages, interaction, null, agentProcess, action)
+            // Buffer raw LLM chunks into complete newline-delimited lines before classifying.
+            // The LLM streams arbitrary byte chunks; thinking tags and JSON objects only make
+            // sense as whole lines, so we must reassemble them first.
+            .transform { rawChunksToLines(it) }
+            // Classify each line as either thinking content or a dropped line.
+            // Each item arriving in onNext is StreamingEvent.Thinking(content, state).
+            // Non-thinking lines (plain JSON artifacts) are silently dropped.
+            //
+            // Examples:
+            //
+            //   a. "aaaaaa"               → plain text, not JSON, no tags
+            //                              → ThinkingState.CONTINUATION
+            //                              → onNext: StreamingEvent.Thinking("aaaaaa", CONTINUATION)
+            //
+            //   b. "aaaaa<think>nnnnn"    → contains opening tag but no closing tag;
+            //                                does not start with <think> so not detected as START
+            //                              → ThinkingState.CONTINUATION
+            //                              → onNext: StreamingEvent.Thinking("aaaaa<think>nnnnn", CONTINUATION)
+            //
+            //   c. "nnnnn</think>"        → ends with closing tag, no opening tag
+            //                              → ThinkingState.END
+            //                              → onNext: StreamingEvent.Thinking("nnnnn</think>", END)
+            //                                (tags not stripped — extractThinkingContent only strips
+            //                                 complete <think>…</think> pairs found on a single line)
+            //
+            //   d. "<think>xyz</think>"   → complete thinking block on one line
+            //                              → ThinkingState.BOTH
+            //                              → onNext: StreamingEvent.Thinking("xyz", BOTH)  [tags stripped]
+            .concatMap { line -> StreamingLineClassifier.classify(line) }
+
     override fun <O> createObjectStream(
         messages: List<Message>,
         interaction: LlmInteraction,

embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImplTest.kt

@@ -0,0 +1,141 @@
+/*
+ * 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.AgentProcess
+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.streaming.ThinkingState
+import io.mockk.every
+import io.mockk.mockk
+import tools.jackson.module.kotlin.jacksonObjectMapper
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Nested
+import org.junit.jupiter.api.Test
+import reactor.core.publisher.Flux
+
+class StreamingLlmOperationsImplTest {
+
+    private lateinit var llmService: LlmService<*>
+    private lateinit var toolDecorator: ToolDecorator
+    private lateinit var agentProcess: AgentProcess
+
+    private val interaction = LlmInteraction(id = InteractionId("test"))
+    private val messages = listOf(UserMessage("hello"))
+
+    @BeforeEach
+    fun setUp() {
+        llmService = mockk(relaxed = true)
+        toolDecorator = mockk(relaxed = true)
+        agentProcess = mockk(relaxed = true)
+        every { llmService.promptContributors } returns emptyList()
+    }
+
+    private fun implWith(vararg chunks: String): StreamingLlmOperationsImpl {
+        val streamer = LlmMessageStreamer { _, _, _ -> Flux.fromArray(chunks) }
+        return StreamingLlmOperationsImpl(
+            messageStreamer = streamer,
+            objectMapper = jacksonObjectMapper(),
+            llmService = llmService,
+            toolDecorator = toolDecorator,
+        )
+    }
+
+    private fun run(impl: StreamingLlmOperationsImpl): List<StreamingEvent<String>> =
+        impl.generateStreamWithThinking(messages, interaction, agentProcess, null)
+            .collectList().block()!!
+
+    @Nested
+    inner class GenerateStreamWithThinking {
+
+        @Nested
+        inner class LineBuffering {
+
+            @Test
+            fun `chunks split across emissions are reassembled before classification`() {
+                // "<think>\n" arrives as three separate chunks, none of which is a complete line alone
+                val events = run(implWith("<th", "ink", ">\n", "reasoning\n", "</think>\n"))
+                assertEquals(3, events.size)
+                assertEquals(ThinkingState.START, (events[0] as StreamingEvent.Thinking).state)
+                assertEquals(ThinkingState.CONTINUATION, (events[1] as StreamingEvent.Thinking).state)
+                assertEquals(ThinkingState.END, (events[2] as StreamingEvent.Thinking).state)
+            }
+
+            @Test
+            fun `multiple newlines in one chunk emit multiple events`() {
+                val events = run(implWith("line one\nline two\n"))
+                assertEquals(2, events.size)
+            }
+
+            @Test
+            fun `trailing content without newline is flushed at stream end`() {
+                // no trailing \n — rawChunksToLines flushes the buffer on complete
+                val events = run(implWith("hello world"))
+                assertEquals(1, events.size)
+                assertEquals("hello world", (events[0] as StreamingEvent.Thinking).content)
+            }
+        }
+
+        @Nested
+        inner class ThinkingClassification {
+
+            @Test
+            fun `complete think block emits BOTH with tags stripped`() {
+                val events = run(implWith("<think>reasoning</think>\n"))
+                assertEquals(1, events.size)
+                val event = events[0] as StreamingEvent.Thinking
+                assertEquals("reasoning", event.content)
+                assertEquals(ThinkingState.BOTH, event.state)
+            }
+
+            @Test
+            fun `multi-line think block emits START then CONTINUATION then END`() {
+                val events = run(implWith("<think>\n", "mid\n", "</think>\n"))
+                assertEquals(3, events.size)
+                assertEquals(ThinkingState.START, (events[0] as StreamingEvent.Thinking).state)
+                assertEquals(ThinkingState.CONTINUATION, (events[1] as StreamingEvent.Thinking).state)
+                assertEquals(ThinkingState.END, (events[2] as StreamingEvent.Thinking).state)
+            }
+
+            @Test
+            fun `plain prose lines emit as CONTINUATION`() {
+                val events = run(implWith("line one\n", "line two\n"))
+                assertEquals(2, events.size)
+                events.forEach { assertEquals(ThinkingState.CONTINUATION, (it as StreamingEvent.Thinking).state) }
+            }
+
+            @Test
+            fun `JSON-shaped lines are dropped`() {
+                val events = run(implWith("{\"key\":\"value\"}\n"))
+                assertTrue(events.isEmpty())
+            }
+
+            @Test
+            fun `code fence lines are dropped`() {
+                val events = run(implWith("```json\n", "<think>reasoning</think>\n", "```\n"))
+                assertEquals(1, events.size)
+                assertEquals(ThinkingState.BOTH, (events[0] as StreamingEvent.Thinking).state)
+            }
+        }
+    }
+}

embabel-agent-common/embabel-agent-ai/src/main/kotlin/com/embabel/common/ai/converters/streaming/StreamingLineClassifier.kt

@@ -0,0 +1,75 @@
+/*
+ * 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.common.ai.converters.streaming
+
+import com.embabel.common.ai.converters.streaming.support.ThinkingDetector
+import com.embabel.common.core.streaming.StreamingEvent
+import com.embabel.common.core.streaming.ThinkingState
+import reactor.core.publisher.Flux
+
+/**
+ * Routes a single newline-delimited line from an LLM stream to the appropriate [StreamingEvent].
+ *
+ * Separates Reactor event-routing from pure thinking detection ([ThinkingDetector]),
+ * so raw-string streaming shares the same dispatch logic in one place.
+ */
+object StreamingLineClassifier {
+
+    // Matches bare markdown code-fence lines such as ```json or ``` that LLMs emit
+    // as formatting artifacts between thinking blocks and JSON. These carry no content
+    // and must be dropped before the thinking path sees them.
+    // Aligns with the equivalent inline check in [StreamingJacksonOutputConverter.convertStreamWithThinking].
+    private val codeFencePattern = Regex("^```\\w*$")
+
+    /**
+     * Classify a single [line] from a newline-delimited LLM stream into zero or one [StreamingEvent.Thinking].
+     *
+     * Every line arriving from the LLM is either:
+     *  - **Thinking content** — wrapped in a tag such as `<think>...</think>` or a partial
+     *    multi-line variant. These become [StreamingEvent.Thinking] events carrying the extracted
+     *    text and a [ThinkingState] that tells the consumer whether this is a complete block,
+     *    the start, a continuation, or the end of a multi-line block.
+     *  - **A bare code fence** (e.g. ` ```json ` or ` ``` `) — a formatting artifact emitted
+     *    by some models between thinking blocks and output. Always dropped.
+     *  - **Anything else** — dropped; [StreamingEvent.Object] is not emitted in the
+     *    raw-string streaming context this classifier serves.
+     *
+     * Detection is delegated to [ThinkingDetector]; this class only owns the Reactor mapping.
+     *
+     * @param line a complete newline-delimited line from the LLM stream (no trailing newline)
+     * @return a [Flux] of at most one [StreamingEvent.Thinking], or empty when the line is dropped
+     */
+    fun classify(line: String): Flux<StreamingEvent<String>> {
+        // Ask ThinkingDetector to classify the line. NONE means non-thinking content (dropped);
+        // anything else (BOTH, START, CONTINUATION, END) means the line contains thinking markup.
+        val state = ThinkingDetector.detectThinkingState(line)
+
+        return when (state) {
+            // Non-thinking line (e.g. a stray JSON line) — not expected in raw-string streaming.
+            ThinkingState.NONE -> Flux.empty()
+
+            // Thinking content line — but first filter out bare code fences (``` / ```json)
+            // which are formatting artifacts that must not leak into thinking events.
+            else -> if (!line.trim().matches(codeFencePattern))
+                // Extract the actual thinking text (strips surrounding tags when present)
+                // and emit a Thinking event with the detected state for multi-line tracking.
+                Flux.just(StreamingEvent.Thinking(ThinkingDetector.extractThinkingContent(line), state))
+            else
+                // Bare code fence — drop silently.
+                Flux.empty()
+        }
+    }
+}

embabel-agent-common/embabel-agent-ai/src/test/kotlin/com/embabel/common/ai/converters/streaming/StreamingLineClassifierTest.kt

@@ -0,0 +1,118 @@
+/*
+ * 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.common.ai.converters.streaming
+
+import com.embabel.common.core.streaming.StreamingEvent
+import com.embabel.common.core.streaming.ThinkingState
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Nested
+import org.junit.jupiter.api.Test
+
+class StreamingLineClassifierTest {
+
+    private fun classify(line: String): List<StreamingEvent<String>> =
+        StreamingLineClassifier.classify(line).collectList().block()!!
+
+    @Nested
+    inner class PlainText {
+
+        @Test
+        fun `plain text emits Thinking with CONTINUATION state`() {
+            val events = classify("aaaaaa")
+            assertEquals(1, events.size)
+            val event = events[0] as StreamingEvent.Thinking
+            assertEquals("aaaaaa", event.content)
+            assertEquals(ThinkingState.CONTINUATION, event.state)
+        }
+
+        @Test
+        fun `text with embedded opening tag not at start emits CONTINUATION`() {
+            // "aaaaa<think>nnnnn" does not start with <think>, so not detected as START
+            val events = classify("aaaaa<think>nnnnn")
+            assertEquals(1, events.size)
+            val event = events[0] as StreamingEvent.Thinking
+            assertEquals("aaaaa<think>nnnnn", event.content)
+            assertEquals(ThinkingState.CONTINUATION, event.state)
+        }
+    }
+
+    @Nested
+    inner class ThinkingTags {
+
+        @Test
+        fun `complete think block strips tags and emits BOTH`() {
+            val events = classify("<think>xyz</think>")
+            assertEquals(1, events.size)
+            val event = events[0] as StreamingEvent.Thinking
+            assertEquals("xyz", event.content)
+            assertEquals(ThinkingState.BOTH, event.state)
+        }
+
+        @Test
+        fun `standalone opening tag emits START`() {
+            val events = classify("<think>")
+            assertEquals(1, events.size)
+            val event = events[0] as StreamingEvent.Thinking
+            assertEquals(ThinkingState.START, event.state)
+        }
+
+        @Test
+        fun `line starting with opening tag and content emits START`() {
+            val events = classify("<think>start of reasoning")
+            assertEquals(1, events.size)
+            val event = events[0] as StreamingEvent.Thinking
+            assertEquals(ThinkingState.START, event.state)
+        }
+
+        @Test
+        fun `standalone closing tag emits END`() {
+            val events = classify("</think>")
+            assertEquals(1, events.size)
+            val event = events[0] as StreamingEvent.Thinking
+            assertEquals(ThinkingState.END, event.state)
+        }
+
+        @Test
+        fun `text ending with closing tag emits END with tags not stripped`() {
+            // extractThinkingContent only strips complete <think>…</think> pairs on one line
+            val events = classify("nnnnn</think>")
+            assertEquals(1, events.size)
+            val event = events[0] as StreamingEvent.Thinking
+            assertEquals("nnnnn</think>", event.content)
+            assertEquals(ThinkingState.END, event.state)
+        }
+    }
+
+    @Nested
+    inner class DroppedLines {
+
+        @Test
+        fun `code fence backtick-json is dropped`() {
+            assertTrue(classify("```json").isEmpty())
+        }
+
+        @Test
+        fun `bare code fence is dropped`() {
+            assertTrue(classify("```").isEmpty())
+        }
+
+        @Test
+        fun `JSON-shaped line is dropped`() {
+            assertTrue(classify("""{"key":"value"}""").isEmpty())
+        }
+    }
+}