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 aefdd863d..e70ec27c2 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 @@ -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, + interaction: LlmInteraction, + agentProcess: AgentProcess, + action: Action?, + ): Flux> = + 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. "aaaaannnnn" → contains opening tag but no closing tag; + // does not start with so not detected as START + // → ThinkingState.CONTINUATION + // → onNext: StreamingEvent.Thinking("aaaaannnnn", CONTINUATION) + // + // c. "nnnnn" → ends with closing tag, no opening tag + // → ThinkingState.END + // → onNext: StreamingEvent.Thinking("nnnnn", END) + // (tags not stripped — extractThinkingContent only strips + // complete … pairs found on a single line) + // + // d. "xyz" → complete thinking block on one line + // → ThinkingState.BOTH + // → onNext: StreamingEvent.Thinking("xyz", BOTH) [tags stripped] + .concatMap { line -> StreamingLineClassifier.classify(line) } + override fun createObjectStream( messages: List, interaction: LlmInteraction, diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImplTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImplTest.kt new file mode 100644 index 000000000..4c77d92b1 --- /dev/null +++ b/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> = + 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`() { + // "\n" arrives as three separate chunks, none of which is a complete line alone + val events = run(implWith("\n", "reasoning\n", "\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("reasoning\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("\n", "mid\n", "\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", "reasoning\n", "```\n")) + assertEquals(1, events.size) + assertEquals(ThinkingState.BOTH, (events[0] as StreamingEvent.Thinking).state) + } + } + } +} diff --git a/embabel-agent-common/embabel-agent-ai/src/main/kotlin/com/embabel/common/ai/converters/streaming/StreamingLineClassifier.kt b/embabel-agent-common/embabel-agent-ai/src/main/kotlin/com/embabel/common/ai/converters/streaming/StreamingLineClassifier.kt new file mode 100644 index 000000000..81317f3e7 --- /dev/null +++ b/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 `...` 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> { + // 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() + } + } +} diff --git a/embabel-agent-common/embabel-agent-ai/src/test/kotlin/com/embabel/common/ai/converters/streaming/StreamingLineClassifierTest.kt b/embabel-agent-common/embabel-agent-ai/src/test/kotlin/com/embabel/common/ai/converters/streaming/StreamingLineClassifierTest.kt new file mode 100644 index 000000000..24e341e02 --- /dev/null +++ b/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> = + 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`() { + // "aaaaannnnn" does not start with , so not detected as START + val events = classify("aaaaannnnn") + assertEquals(1, events.size) + val event = events[0] as StreamingEvent.Thinking + assertEquals("aaaaannnnn", event.content) + assertEquals(ThinkingState.CONTINUATION, event.state) + } + } + + @Nested + inner class ThinkingTags { + + @Test + fun `complete think block strips tags and emits BOTH`() { + val events = classify("xyz") + 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("") + 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("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("") + 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 … pairs on one line + val events = classify("nnnnn") + assertEquals(1, events.size) + val event = events[0] as StreamingEvent.Thinking + assertEquals("nnnnn", 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()) + } + } +}