diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/callback/ToolCallInspector.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/callback/ToolCallInspector.kt index f09e4c3d9..3b22420dc 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/callback/ToolCallInspector.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/callback/ToolCallInspector.kt @@ -22,8 +22,7 @@ import com.embabel.chat.ToolCall * Read-only observer for individual tool call events. * * Provides observation of tool execution without access to conversation history - * or iteration state. Works in both streaming mode (where the framework manages - * the tool loop internally) and non-streaming mode (as a lightweight alternative + * or iteration state. Works in both streaming and non-streaming modes as a lightweight alternative * to [ToolLoopInspector] when history/iteration context is not needed). * * @see ToolLoopInspector for tool loop-level inspection with full conversation context @@ -49,10 +48,9 @@ interface ToolCallInspector { * Lightweight context without conversation history or iteration state. * Used in both streaming and non-streaming modes. * - * **Note:** In streaming mode (Spring AI), `toolCall.id` is a generated UUID - * because the underlying framework does not expose LLM-assigned call IDs during - * tool execution. The ID is unique per call and can be used to correlate - * before/after events for the same tool execution. + * In streaming mode, adapters preserve the provider-assigned tool-call ID when the + * provider supplies one. The same ID is passed to both callbacks so inspectors can + * correlate the execution without depending on a specific provider. * * @property toolCall The tool call about to be executed */ diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/loop/streaming/LlmMessageStreamer.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/loop/streaming/LlmMessageStreamer.kt index 19dfe6be0..27cb2a1da 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/loop/streaming/LlmMessageStreamer.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/loop/streaming/LlmMessageStreamer.kt @@ -17,38 +17,57 @@ package com.embabel.agent.spi.loop.streaming import com.embabel.agent.api.tool.Tool import com.embabel.agent.api.tool.callback.ToolCallInspector +import com.embabel.agent.core.Usage import com.embabel.agent.spi.loop.LlmMessageSender import com.embabel.agent.spi.loop.ToolLoop +import com.embabel.chat.AssistantMessage import com.embabel.chat.Message import reactor.core.publisher.Flux +/** + * Provider-neutral events emitted by one streaming LLM inference call. + * + * [Content] events are forwarded as soon as they arrive. [Complete] is emitted once, + * after provider-specific fragments (including partial tool calls) have been assembled. + * + * This SPI type is deliberately distinct from + * [com.embabel.common.core.streaming.StreamingEvent], which is the public, converted + * application stream of thinking and typed objects. These events exist below that layer: + * [DefaultStreamingToolLoop] consumes [Complete] internally and forwards only content. + */ +sealed interface LlmInferenceStreamEvent { + data class Content(val text: String) : LlmInferenceStreamEvent + + data class Complete( + val message: Message, + val usage: Usage? = null, + ) : LlmInferenceStreamEvent +} + /** * Framework-agnostic interface for streaming LLM inference. * * Streaming counterpart of [LlmMessageSender]. Implementations handle the actual - * LLM communication (Spring AI, LangChain4j, etc.) and return a reactive stream - * of raw content chunks. + * LLM communication (Spring AI, LangChain4j, etc.). The original [stream] contract + * remains the single abstract method for source and binary compatibility. The new + * [streamInference] method exposes the assembled terminal message required by an + * Embabel-owned streaming tool loop. * * **Key Differences from Non-Streaming:** - * - Returns `Flux` instead of `LlmMessageResponse` - * - Tool execution is managed by the underlying framework (e.g., Spring AI) - * since the streaming API is opaque - we cannot inject a custom [ToolLoop] - * - Only observation of tool execution is possible via [ToolCallInspector] + * - Emits content incrementally, followed by the assembled assistant message + * - Does not execute tools; [ToolLoop] remains under Embabel's control + * - Provider adapters are responsible for assembling partial tool-call fragments * * @see LlmMessageSender for non-streaming equivalent - * @see ToolCallInspector for tool execution observation */ fun interface LlmMessageStreamer { /** - * Stream raw content chunks from the LLM. - * - * The returned Flux emits content as it arrives from the LLM. - * Tool calls are handled internally by the underlying framework. + * Stream raw content chunks using the original provider-managed tool contract. * * @param messages The conversation history * @param tools Available tools for the LLM to invoke during streaming - * @param toolCallInspectors Inspectors to observe tool call events + * @param toolCallInspectors Inspectors for provider-managed tool execution * @return Flux of raw content chunks */ fun stream( @@ -56,4 +75,27 @@ fun interface LlmMessageStreamer { tools: List, toolCallInspectors: List, ): Flux + + /** + * Stream one inference without delegating tool execution to the provider. + * + * Provider adapters should override this method to emit an assembled [Complete] + * event containing any requested tool calls. The default bridge keeps existing + * [LlmMessageStreamer] implementations usable: their raw stream is replayed as + * [Content] and summarized as a terminal assistant message. Such legacy adapters + * retain their provider-managed tool behavior and cannot expose assembled tool calls + * to [DefaultStreamingToolLoop] until they override this method. + */ + fun streamInference( + messages: List, + tools: List, + ): Flux = Flux.defer { + val content = stream(messages, tools, emptyList()).cache() + Flux.concat( + content.map { LlmInferenceStreamEvent.Content(it) }, + content.collectList().map { + LlmInferenceStreamEvent.Complete(AssistantMessage(it.joinToString(""))) + }, + ) + } } diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/loop/streaming/StreamingToolLoop.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/loop/streaming/StreamingToolLoop.kt new file mode 100644 index 000000000..f9e1a5f74 --- /dev/null +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/loop/streaming/StreamingToolLoop.kt @@ -0,0 +1,230 @@ +/* + * 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.loop.streaming + +import com.embabel.agent.api.tool.Tool +import com.embabel.agent.api.tool.ToolCallContext +import com.embabel.agent.api.tool.callback.ToolCallInspector +import com.embabel.agent.spi.loop.AutoCorrectionPolicy +import com.embabel.agent.spi.loop.MaxIterationsExceededException +import com.embabel.agent.spi.loop.ToolInjectionStrategy +import com.embabel.agent.spi.loop.ToolNotFoundAction +import com.embabel.agent.spi.loop.ToolNotFoundPolicy +import com.embabel.agent.spi.loop.support.applyToolInjection +import com.embabel.agent.spi.loop.support.executeTool +import com.embabel.chat.AssistantMessageWithToolCalls +import com.embabel.chat.Message +import com.embabel.chat.ToolCall +import com.embabel.chat.ToolResultMessage +import org.slf4j.LoggerFactory +import reactor.core.publisher.Flux +import tools.jackson.databind.ObjectMapper + +/** + * Provider-neutral counterpart of [com.embabel.agent.spi.loop.ToolLoop]. + * + * Implementations manage tool execution and conversation continuation while delegating + * each individual streaming inference to an [LlmMessageStreamer]. The returned [Flux] + * contains content from every inference turn in arrival order; provider-level completion + * events remain internal to the loop. + */ +fun interface StreamingToolLoop { + fun execute(initialMessages: List, initialTools: List): Flux + + companion object { + /** + * Create the default provider-neutral streaming tool loop while keeping its + * implementation internal to Embabel's SPI surface. + */ + @JvmStatic + @JvmOverloads + fun create( + messageStreamer: LlmMessageStreamer, + objectMapper: ObjectMapper, + injectionStrategy: ToolInjectionStrategy = ToolInjectionStrategy.NONE, + maxIterations: Int = 20, + toolDecorator: ((Tool) -> Tool)? = null, + toolCallInspectors: List = emptyList(), + toolCallContext: ToolCallContext = ToolCallContext.EMPTY, + toolNotFoundPolicy: ToolNotFoundPolicy = AutoCorrectionPolicy(), + ): StreamingToolLoop = DefaultStreamingToolLoop( + messageStreamer = messageStreamer, + objectMapper = objectMapper, + injectionStrategy = injectionStrategy, + maxIterations = maxIterations, + toolDecorator = toolDecorator, + toolCallInspectors = toolCallInspectors, + toolCallContext = toolCallContext, + toolNotFoundPolicy = toolNotFoundPolicy, + ) + } +} + +/** + * Default Embabel-owned streaming tool loop. + * + * Code flow for each subscription: + * 1. Create isolated mutable conversation and tool state. + * 2. Ask [LlmMessageStreamer] to perform one inference. + * 3. Forward [LlmInferenceStreamEvent.Content] immediately. + * 4. Use [LlmInferenceStreamEvent.Complete] to append the assembled assistant message. + * 5. If the assistant requested tools, execute them, append results, apply + * [ToolInjectionStrategy], and recursively start the next inference. + * 6. Complete when the assistant requests no more tools, or emit direct tool content + * when a tool has `returnDirect=true`. + * + * `Flux.defer` is deliberately used at the entry point and for every turn. It creates + * state at subscription time and prevents conversation history, iteration counters, or + * dynamically injected tools from leaking between multiple subscriptions to the same Flux. + * `Flux.concatMap` preserves the ordering between content chunks, terminal processing, + * tool execution, and the subsequent inference. + * + * Multiple tool calls in one assistant message are executed sequentially, matching + * [com.embabel.agent.spi.loop.support.DefaultToolLoop]. This preserves declared order + * and deterministic `returnDirect`, inspector, and dynamic-injection semantics. Parallel + * execution would require a separate contract for those ordering-sensitive behaviours. + * + * @param maxIterations Upper bound for inference turns. Production wiring supplies + * [com.embabel.agent.core.support.LlmInteraction.maxToolIterations]; `20` is only the + * direct-construction default. + */ +internal class DefaultStreamingToolLoop( + private val messageStreamer: LlmMessageStreamer, + private val objectMapper: ObjectMapper, + private val injectionStrategy: ToolInjectionStrategy = ToolInjectionStrategy.NONE, + private val maxIterations: Int = 20, + private val toolDecorator: ((Tool) -> Tool)? = null, + private val toolCallInspectors: List = emptyList(), + private val toolCallContext: ToolCallContext = ToolCallContext.EMPTY, + private val toolNotFoundPolicy: ToolNotFoundPolicy = AutoCorrectionPolicy(), +) : StreamingToolLoop { + + private val logger = LoggerFactory.getLogger(javaClass) + + override fun execute( + initialMessages: List, + initialTools: List, + ): Flux = Flux.defer { + streamTurn( + State( + conversationHistory = initialMessages.toMutableList(), + availableTools = initialTools.toMutableList(), + ) + ) + } + + private fun streamTurn(state: State): Flux = Flux.defer { + if (state.iterations >= maxIterations) { + return@defer Flux.error(MaxIterationsExceededException(maxIterations)) + } + state.iterations++ + logger.debug( + "Streaming tool loop iteration {} with {} available tools", + state.iterations, + state.availableTools.size, + ) + + messageStreamer.streamInference( + messages = state.conversationHistory.toList(), + tools = state.availableTools.toList(), + ).concatMap { event -> + when (event) { + is LlmInferenceStreamEvent.Content -> Flux.just(event.text) + is LlmInferenceStreamEvent.Complete -> continueFrom(event, state) + } + } + } + + /** + * Processes the assembled terminal message for one inference turn. + * + * Content was already forwarded as [LlmInferenceStreamEvent.Content], so an assistant + * response without tool calls contributes no additional item and completes with + * [Flux.empty]. Tool calls append their results to the conversation before the next + * inference starts. A direct-return tool instead emits its result and ends the loop. + */ + private fun continueFrom( + event: LlmInferenceStreamEvent.Complete, + state: State, + ): Flux { + state.conversationHistory.add(event.message) + val assistant = event.message as? AssistantMessageWithToolCalls + ?: return Flux.empty() + if (assistant.toolCalls.isEmpty()) return Flux.empty() + + for (toolCall in assistant.toolCalls) { + val directResult = executeToolCall(toolCall, state) + if (directResult != null) return Flux.just(directResult) + } + return streamTurn(state) + } + + /** + * Executes one tool call and returns direct content when the tool short-circuits the loop. + */ + private fun executeToolCall(toolCall: ToolCall, state: State): String? { + val tool = state.availableTools.find { it.definition.name == toolCall.name } + if (tool == null) { + when (val action = toolNotFoundPolicy.handle(toolCall.name, state.availableTools)) { + is ToolNotFoundAction.Throw -> throw action.exception + is ToolNotFoundAction.FeedbackToModel -> { + logger.warn(action.message) + state.conversationHistory.add( + ToolResultMessage(toolCall.id, toolCall.name, action.message) + ) + return null + } + } + } + + toolNotFoundPolicy.onToolFound() + logger.debug("Executing streaming tool: {}", toolCall.name) + val executed = executeTool( + tool = tool, + toolCall = toolCall, + toolCallContext = toolCallContext, + toolCallInspectors = toolCallInspectors, + ) + + applyInjection(toolCall, executed.content, state) + state.conversationHistory.add(ToolResultMessage(toolCall.id, toolCall.name, executed.content)) + if (tool.metadata.returnDirect) { + logger.info("Tool '{}' has returnDirect=true — completing streaming loop", toolCall.name) + return executed.content + } + return null + } + + private fun applyInjection(toolCall: ToolCall, resultContent: String, state: State) { + applyToolInjection( + toolCall = toolCall, + resultContent = resultContent, + conversationHistory = state.conversationHistory, + availableTools = state.availableTools, + iteration = state.iterations, + injectionStrategy = injectionStrategy, + objectMapper = objectMapper, + toolDecorator = toolDecorator, + logger = logger, + ) + } + + private data class State( + val conversationHistory: MutableList, + val availableTools: MutableList, + var iterations: Int = 0, + ) +} diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/SpringAiLlmService.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/SpringAiLlmService.kt index 3444f0539..c2e21cd8e 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/SpringAiLlmService.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/SpringAiLlmService.kt @@ -24,7 +24,6 @@ import com.embabel.common.ai.model.* import com.embabel.common.ai.prompt.KnowledgeCutoffDate import com.embabel.common.ai.prompt.PromptContributor import tools.jackson.databind.annotation.JsonSerialize -import org.springframework.ai.chat.client.ChatClient import org.springframework.ai.chat.messages.UserMessage import org.springframework.ai.chat.model.ChatModel import org.springframework.ai.chat.model.ChatResponse @@ -141,8 +140,11 @@ data class SpringAiLlmService @JvmOverloads constructor( } override fun createMessageStreamer(options: LlmOptions): LlmMessageStreamer { - val chatClient = ChatClient.create(chatModel) - return SpringAiLlmMessageStreamer(chatClient, convertOptions(options)) + return SpringAiLlmMessageStreamer( + chatModel = chatModel, + chatOptions = convertOptions(options), + toolResponseContentAdapter = toolResponseContentAdapter, + ) } override fun supportsStreaming(): Boolean = StreamingCapabilityVerifier.supportsStreaming(chatModel) diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/SpringAiLlmMessageStreamer.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/SpringAiLlmMessageStreamer.kt index dd4a31af2..5c612370c 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/SpringAiLlmMessageStreamer.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/SpringAiLlmMessageStreamer.kt @@ -18,44 +18,56 @@ package com.embabel.agent.spi.support.springai.streaming import com.embabel.agent.api.tool.Tool import com.embabel.agent.api.tool.callback.ToolCallInspector import com.embabel.agent.spi.loop.streaming.LlmMessageStreamer +import com.embabel.agent.spi.loop.streaming.LlmInferenceStreamEvent import com.embabel.agent.spi.support.springai.SpringAiLlmMessageSender import com.embabel.agent.spi.support.springai.toSpringAiMessage +import com.embabel.agent.spi.support.springai.toEmbabelMessage +import com.embabel.agent.spi.support.springai.toEmbabelUsage import com.embabel.agent.spi.support.springai.toSpringToolCallbacks -import com.embabel.agent.spi.support.springai.withInspectors +import com.embabel.agent.spi.support.springai.ToolResponseContentAdapter import com.embabel.chat.Message -import org.springframework.ai.chat.client.ChatClient +import org.springframework.ai.chat.model.ChatModel +import org.springframework.ai.chat.model.ChatResponse +import org.springframework.ai.chat.model.MessageAggregator import org.springframework.ai.chat.prompt.ChatOptions import org.springframework.ai.chat.prompt.Prompt import reactor.core.publisher.Flux +import reactor.core.publisher.Sinks /** * Spring AI implementation of [LlmMessageStreamer]. * - * Streams raw content chunks from the LLM using Spring AI's ChatClient. - * Tool execution is handled internally by Spring AI during streaming. - * Tool call events can be observed via [ToolCallInspector]. + * Streams one inference directly through Spring AI's [ChatModel]. Calling the model + * directly advertises tools without installing `ToolCallingAdvisor`, leaving execution + * and continuation to Embabel's streaming tool loop. Consequently, tool-call inspectors + * are invoked by [com.embabel.agent.spi.loop.streaming.DefaultStreamingToolLoop], around + * the actual tool execution, rather than by this single-inference adapter. * - * @param chatClient The Spring AI ChatClient for streaming + * @param chatModel The Spring AI model for streaming * @param chatOptions Options for the LLM call (temperature, model, etc.) * @see SpringAiLlmMessageSender for non-streaming equivalent - * @see ToolCallInspector for tool call observation */ internal class SpringAiLlmMessageStreamer( - private val chatClient: ChatClient, + private val chatModel: ChatModel, private val chatOptions: ChatOptions, + private val toolResponseContentAdapter: ToolResponseContentAdapter = ToolResponseContentAdapter.PASSTHROUGH, ) : LlmMessageStreamer { override fun stream( messages: List, tools: List, toolCallInspectors: List, - ): Flux { - val springAiMessages = messages.map { it.toSpringAiMessage() } - val toolCallbacks = tools.toSpringToolCallbacks().withInspectors(toolCallInspectors) - // Spring AI 2.0: bake toolCallbacks into ToolCallingChatOptions AND pass them on the - // request spec. The former preserves the ToolCallingChatOptions subtype through the - // chatModel-defaults merge; the latter survives the merge that would otherwise reset - // prompt.options.toolCallbacks to the model's empty default. + ): Flux = streamInference(messages, tools) + .ofType(LlmInferenceStreamEvent.Content::class.java) + .map { it.text } + + override fun streamInference( + messages: List, + tools: List, + ): Flux = Flux.defer { + val springAiMessages = messages + .map { it.toSpringAiMessage(toolResponseContentAdapter) } + val toolCallbacks = tools.toSpringToolCallbacks() val effectiveOptions = if (chatOptions is org.springframework.ai.model.tool.ToolCallingChatOptions) { chatOptions.mutate() .toolCallbacks(toolCallbacks) @@ -64,11 +76,34 @@ internal class SpringAiLlmMessageStreamer( chatOptions } val prompt = Prompt(springAiMessages, effectiveOptions) + // Bridges MessageAggregator's terminal callback back into the reactive sequence. + // It emits one assembled response for this inference turn, after its partial + // content/tool fragments and before Embabel executes any requested tools. + val aggregate = Sinks.one() - return chatClient - .prompt(prompt) - .tools(toolCallbacks) - .stream() - .content() + // Spring providers may split a tool-call ID, name, and JSON arguments across + // multiple ChatResponse chunks. MessageAggregator leaves those chunks in the + // returned Flux (so content stays live) and invokes this callback once with the + // assembled assistant response. The aggregate becomes the terminal SPI event; + // this adapter deliberately performs no tool execution. + chatModel.stream(prompt).publish { responses -> + val content = responses + .flatMapIterable { response -> listOfNotNull(response.result?.output?.text) } + .filter { it.isNotEmpty() } + .map { LlmInferenceStreamEvent.Content(it) } + val completion = MessageAggregator() + .aggregate(responses, aggregate::tryEmitValue) + .then(aggregate.asMono()) + .map { response -> + val output = response.result?.output + ?: throw IllegalStateException("Streaming response contained no assistant output") + LlmInferenceStreamEvent.Complete( + message = output.toEmbabelMessage(), + usage = response.metadata?.usage?.toEmbabelUsage(), + ) + } + + Flux.mergeSequential(content, completion) + } } } 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 eb96c7052..529ecd971 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 @@ -491,10 +491,20 @@ internal class StreamingChatClientOperations( ): Flux { return if (useMessageStreamer) { val streamerMessages = buildMessagesWithContributions(messages, promptContributions) - SpringAiLlmMessageStreamer(chatClient, chatOptions).stream(streamerMessages, tools, toolCallInspectors) + val toolCallbacks = tools.toSpringToolCallbacks() + val effectiveOptions = if (chatOptions is org.springframework.ai.model.tool.ToolCallingChatOptions) { + chatOptions.mutate().toolCallbacks(toolCallbacks).build() + } else { + chatOptions + } + chatClient + .prompt(Prompt(streamerMessages.map { it.toSpringAiMessage() }, effectiveOptions)) + .tools(toolCallbacks) + .stream() + .content() } else { // Spring AI 2.0: bake toolCallbacks into ToolCallingChatOptions AND pass them via - // .toolCallbacks() on the request spec. The former preserves the ToolCallingChatOptions + // .tools() on the request spec. The former preserves the ToolCallingChatOptions // subtype through the chatModel-defaults merge; the latter survives the merge that // would otherwise reset prompt.options.toolCallbacks to the model's empty default. val springAiToolCallbacks = tools.toSpringToolCallbacks() 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..aefdd863d 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 @@ -17,12 +17,17 @@ package com.embabel.agent.spi.support.streaming import com.embabel.agent.api.event.LlmRequestEvent import com.embabel.agent.api.tool.Tool +import com.embabel.agent.api.tool.ToolCallContext import com.embabel.agent.core.Action 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.agent.spi.loop.streaming.DefaultStreamingToolLoop +import com.embabel.agent.spi.loop.AutoCorrectionPolicy +import com.embabel.agent.spi.loop.ChainedToolInjectionStrategy +import com.embabel.agent.spi.loop.ToolInjectionStrategy import com.embabel.agent.core.internal.streaming.StreamingLlmOperations import com.embabel.agent.spi.support.PROMPT_ELEMENT_SEPARATOR import com.embabel.agent.spi.support.ToolResolutionHelper @@ -134,7 +139,7 @@ internal class StreamingLlmOperationsImpl( val messagesWithContributions = buildMessagesWithContributions(messages, promptContributions) // Stream raw chunks from LLM - return messageStreamer.stream(messagesWithContributions, tools, interaction.toolCallInspectors) + return streamWithToolLoop(messagesWithContributions, tools, interaction, agentProcess, action) } override fun doTransformObjectStream( @@ -230,7 +235,13 @@ internal class StreamingLlmOperationsImpl( val messagesWithContributions = buildMessagesWithContributions(messages, fullPromptContributions) // Step 1: Raw chunk stream from LLM - val rawChunkFlux: Flux = messageStreamer.stream(messagesWithContributions, tools, interaction.toolCallInspectors) + val rawChunkFlux: Flux = streamWithToolLoop( + messagesWithContributions, + tools, + interaction, + agentProcess, + action, + ) .filter { it.isNotEmpty() } .doOnNext { chunk -> logger.trace("RAW CHUNK: '${chunk.replace("\n", "\\n")}'") } @@ -273,6 +284,53 @@ internal class StreamingLlmOperationsImpl( return ToolResolutionHelper.resolveAndDecorate(interaction, agentProcess, action, toolDecorator) } + /** + * Stream through Embabel's provider-neutral tool loop. The loop is created per + * invocation so conversation, injection, and retry state cannot leak between subscribers. + */ + private fun streamWithToolLoop( + messages: List, + tools: List, + interaction: LlmInteraction, + agentProcess: AgentProcess?, + action: Action?, + ): Flux { + val injectionStrategy = if (interaction.additionalInjectionStrategies.isEmpty()) { + ToolInjectionStrategy.DEFAULT + } else { + ChainedToolInjectionStrategy( + listOf(ToolInjectionStrategy.DEFAULT) + interaction.additionalInjectionStrategies + ) + } + val processToolCallContext = agentProcess + ?.processContext + ?.processOptions + ?.toolCallContext + ?: ToolCallContext.EMPTY + val toolCallContext = processToolCallContext.merge(interaction.toolCallContext) + val injectedToolDecorator: ((Tool) -> Tool)? = agentProcess?.let { process -> + { tool -> + toolDecorator.decorate( + tool = tool, + agentProcess = process, + action = action, + llmOptions = interaction.llm, + ) + } + } + + return DefaultStreamingToolLoop( + messageStreamer = messageStreamer, + objectMapper = objectMapper, + injectionStrategy = injectionStrategy, + maxIterations = interaction.maxToolIterations, + toolDecorator = injectedToolDecorator, + toolCallInspectors = interaction.toolCallInspectors, + toolCallContext = toolCallContext, + toolNotFoundPolicy = interaction.toolNotFoundPolicy ?: AutoCorrectionPolicy(), + ).execute(messages, tools) + } + /** * Convert raw streaming chunks to NDJSON lines. * Handles all cases: multiple \n in one chunk, no \n in chunk, line spanning many chunks. diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/loop/streaming/StreamingToolLoopTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/loop/streaming/StreamingToolLoopTest.kt new file mode 100644 index 000000000..742c5a733 --- /dev/null +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/loop/streaming/StreamingToolLoopTest.kt @@ -0,0 +1,220 @@ +/* + * 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.loop.streaming + +import com.embabel.agent.api.tool.Tool +import com.embabel.agent.api.tool.callback.AfterToolCallContext +import com.embabel.agent.api.tool.callback.BeforeToolCallContext +import com.embabel.agent.api.tool.callback.ToolCallInspector +import com.embabel.agent.spi.loop.ToolInjectionContext +import com.embabel.agent.spi.loop.ToolInjectionResult +import com.embabel.agent.spi.loop.ToolInjectionStrategy +import com.embabel.chat.AssistantMessage +import com.embabel.chat.AssistantMessageWithToolCalls +import com.embabel.chat.Message +import com.embabel.chat.ToolCall +import com.embabel.chat.ToolResultMessage +import com.embabel.chat.UserMessage +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Test +import reactor.core.publisher.Flux +import reactor.test.StepVerifier +import tools.jackson.module.kotlin.jacksonObjectMapper + +class StreamingToolLoopTest { + + @Test + fun `preserves content across tool turns and executes tools`() { + val histories = mutableListOf>() + var calls = 0 + val streamer = inferenceStreamer { messages, _ -> + histories += messages + calls++ + if (calls == 1) { + Flux.just( + LlmInferenceStreamEvent.Content("checking"), + LlmInferenceStreamEvent.Complete( + AssistantMessageWithToolCalls( + content = "checking", + toolCalls = listOf(ToolCall("call-1", "lookup", "{}")), + ) + ), + ) + } else { + Flux.just( + LlmInferenceStreamEvent.Content("final"), + LlmInferenceStreamEvent.Complete(AssistantMessage("final")), + ) + } + } + val lookup = Tool.of("lookup", "Lookup") { Tool.Result.text("tool-result") } + + val result = DefaultStreamingToolLoop(streamer, jacksonObjectMapper()) + .execute(listOf(UserMessage("question")), listOf(lookup)) + + StepVerifier.create(result) + .expectNext("checking", "final") + .verifyComplete() + assertEquals(2, histories.size) + assertEquals("tool-result", histories[1].filterIsInstance().single().content) + } + + @Test + fun `applies injection strategy before the next turn`() { + val child = Tool.of("child", "Child tool") { Tool.Result.text("done") } + val facade = Tool.of("facade", "Reveal child") { Tool.Result.text("revealed") } + val toolsByTurn = mutableListOf>() + var calls = 0 + val streamer = inferenceStreamer { _, tools -> + toolsByTurn += tools.map { it.definition.name } + calls++ + when (calls) { + 1 -> toolCall("facade", "call-1") + 2 -> toolCall("child", "call-2") + else -> Flux.just(LlmInferenceStreamEvent.Complete(AssistantMessage("complete"))) + } + } + val injection = object : ToolInjectionStrategy { + override fun evaluate(context: ToolInjectionContext): ToolInjectionResult = + if (context.lastToolCall.toolName == "facade") { + ToolInjectionResult.replace(facade, listOf(child)) + } else { + ToolInjectionResult.noChange() + } + } + + StepVerifier.create( + DefaultStreamingToolLoop(streamer, jacksonObjectMapper(), injection) + .execute(listOf(UserMessage("question")), listOf(facade)) + ).verifyComplete() + + assertEquals(listOf("facade"), toolsByTurn[0]) + assertEquals(listOf("child"), toolsByTurn[1]) + } + + @Test + fun `notifies inspectors from the provider-neutral loop`() { + val observed = mutableListOf() + val inspector = object : ToolCallInspector { + override fun beforeToolCall(context: BeforeToolCallContext) { + observed += "before:${context.toolCall.id}" + } + + override fun afterToolCall(context: AfterToolCallContext) { + observed += "after:${context.toolCall.id}:${context.resultAsString}" + } + } + var calls = 0 + val streamer = inferenceStreamer { _, _ -> + calls++ + if (calls == 1) toolCall("lookup", "call-1") + else Flux.just(LlmInferenceStreamEvent.Complete(AssistantMessage("complete"))) + } + val lookup = Tool.of("lookup", "Lookup") { Tool.Result.text("result") } + + StepVerifier.create( + DefaultStreamingToolLoop( + messageStreamer = streamer, + objectMapper = jacksonObjectMapper(), + toolCallInspectors = listOf(inspector), + ).execute(listOf(UserMessage("question")), listOf(lookup)) + ).verifyComplete() + + assertEquals(listOf("before:call-1", "after:call-1:result"), observed) + } + + @Test + fun `executes multiple tool calls sequentially in declared order`() { + val executionOrder = mutableListOf() + val first = Tool.of("first", "First") { + executionOrder += "first" + Tool.Result.text("one") + } + val second = Tool.of("second", "Second") { + executionOrder += "second" + Tool.Result.text("two") + } + var turns = 0 + val streamer = inferenceStreamer { _, _ -> + turns++ + if (turns == 1) { + Flux.just( + LlmInferenceStreamEvent.Complete( + AssistantMessageWithToolCalls( + content = "", + toolCalls = listOf( + ToolCall("call-1", "first", "{}"), + ToolCall("call-2", "second", "{}"), + ), + ) + ) + ) + } else { + Flux.just(LlmInferenceStreamEvent.Complete(AssistantMessage("complete"))) + } + } + + StepVerifier.create( + DefaultStreamingToolLoop(streamer, jacksonObjectMapper()) + .execute(listOf(UserMessage("question")), listOf(first, second)) + ).verifyComplete() + + assertEquals(listOf("first", "second"), executionOrder) + } + + @Test + fun `does not duplicate terminal assistant content`() { + val streamer = inferenceStreamer { _, _ -> + Flux.just( + LlmInferenceStreamEvent.Content("answer"), + LlmInferenceStreamEvent.Complete(AssistantMessage("answer")), + ) + } + + StepVerifier.create( + DefaultStreamingToolLoop(streamer, jacksonObjectMapper()) + .execute(listOf(UserMessage("question")), emptyList()) + ) + .expectNext("answer") + .verifyComplete() + } + + private fun toolCall(name: String, id: String): Flux = Flux.just( + LlmInferenceStreamEvent.Complete( + AssistantMessageWithToolCalls( + content = "", + toolCalls = listOf(ToolCall(id, name, "{}")), + ) + ) + ) + + private fun inferenceStreamer( + block: (List, List) -> Flux, + ): LlmMessageStreamer = object : LlmMessageStreamer { + override fun stream( + messages: List, + tools: List, + toolCallInspectors: List, + ): Flux = block(messages, tools) + .ofType(LlmInferenceStreamEvent.Content::class.java) + .map { it.text } + + override fun streamInference( + messages: List, + tools: List, + ): Flux = block(messages, tools) + } +} diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/springai/streaming/SpringAiLlmMessageStreamerTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/springai/streaming/SpringAiLlmMessageStreamerTest.kt index 09090df73..75836719c 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/springai/streaming/SpringAiLlmMessageStreamerTest.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/springai/streaming/SpringAiLlmMessageStreamerTest.kt @@ -16,230 +16,96 @@ package com.embabel.agent.spi.support.springai.streaming import com.embabel.agent.api.tool.Tool -import com.embabel.agent.api.tool.callback.AfterToolCallContext -import com.embabel.agent.api.tool.callback.BeforeToolCallContext -import com.embabel.agent.api.tool.callback.ToolCallInspector -import com.embabel.chat.AssistantMessage -import com.embabel.chat.SystemMessage +import com.embabel.agent.spi.loop.streaming.LlmInferenceStreamEvent +import com.embabel.chat.AssistantMessageWithToolCalls import com.embabel.chat.UserMessage import io.mockk.every import io.mockk.mockk import io.mockk.slot -import io.mockk.verify -import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertInstanceOf import org.junit.jupiter.api.BeforeEach -import org.junit.jupiter.api.Nested import org.junit.jupiter.api.Test -import org.springframework.ai.chat.client.ChatClient +import org.springframework.ai.chat.messages.AssistantMessage +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.prompt.Prompt import org.springframework.ai.model.tool.ToolCallingChatOptions -import org.springframework.ai.tool.ToolCallback import reactor.core.publisher.Flux import reactor.test.StepVerifier -import java.time.Duration -/** - * Unit tests for [SpringAiLlmMessageStreamer]. - * - * Verifies that the streamer correctly: - * - Converts Embabel messages to Spring AI messages - * - Converts Embabel tools to Spring AI tool callbacks - * - Calls the ChatClient streaming chain correctly - */ class SpringAiLlmMessageStreamerTest { - private lateinit var mockChatClient: ChatClient - // Use a real ToolCallingChatOptions instance — the Builder is self-bounded generic - // (Builder>) and cannot be cleanly mocked. The real implementation - // exercises the production .mutate().toolCallbacks().build() chain end-to-end. + private lateinit var chatModel: ChatModel private lateinit var chatOptions: ToolCallingChatOptions - private lateinit var mockRequestSpec: ChatClient.ChatClientRequestSpec - private lateinit var mockStreamSpec: ChatClient.StreamResponseSpec private val capturedPrompt = slot() @BeforeEach fun setUp() { - mockChatClient = mockk(relaxed = true) + chatModel = mockk() chatOptions = ToolCallingChatOptions.builder().build() - mockRequestSpec = mockk(relaxed = true) - mockStreamSpec = mockk(relaxed = true) capturedPrompt.clear() - - // Spring AI 2.0: production bakes toolCallbacks into the ToolCallingChatOptions - // AND also passes them via .toolCallbacks() on the request spec — the spec call - // survives Spring AI's options merge that would otherwise reset the toolCallbacks - // list to the chat model's empty default. - every { mockChatClient.prompt(capture(capturedPrompt)) } returns mockRequestSpec - every { mockRequestSpec.tools(any>()) } returns mockRequestSpec - every { mockRequestSpec.stream() } returns mockStreamSpec } @Test - fun `stream calls ChatClient with correct chain`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - val messages = listOf(UserMessage("Hello")) - val tools = emptyList() - every { mockStreamSpec.content() } returns Flux.just("response") - - // When - streamer.stream(messages, tools, emptyList()) - - // Then - verify { mockChatClient.prompt(any()) } - verify { mockRequestSpec.tools(any>()) } - verify { mockRequestSpec.stream() } - // Spring AI 2.0: prompt.options is the rebuilt ToolCallingChatOptions. - assertTrue(capturedPrompt.captured.options is ToolCallingChatOptions) - } - - @Test - fun `stream returns content from ChatClient`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - val messages = listOf(UserMessage("Hello")) - val expectedContent = Flux.just("chunk1", "chunk2", "chunk3") - every { mockStreamSpec.content() } returns expectedContent - - // When - val result = streamer.stream(messages, emptyList(), emptyList()) + fun `streams content and emits assembled completion`() { + every { chatModel.stream(capture(capturedPrompt)) } returns Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("hello")))) + ) + val streamer = SpringAiLlmMessageStreamer(chatModel, chatOptions) - // Then - StepVerifier.create(result) - .expectNext("chunk1") - .expectNext("chunk2") - .expectNext("chunk3") + StepVerifier.create(streamer.streamInference(listOf(UserMessage("Hi")), emptyList())) + .expectNext(LlmInferenceStreamEvent.Content("hello")) + .assertNext { event -> + val completion = assertInstanceOf(LlmInferenceStreamEvent.Complete::class.java, event) + assertEquals("hello", completion.message.content) + } .verifyComplete() } @Test - fun `stream handles empty message list`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - every { mockStreamSpec.content() } returns Flux.just("response") - - // When - val result = streamer.stream(emptyList(), emptyList(), emptyList()) + fun `preserves the original raw content streaming signature`() { + every { chatModel.stream(capture(capturedPrompt)) } returns Flux.just( + ChatResponse(listOf(Generation(AssistantMessage("hello")))) + ) + val streamer = SpringAiLlmMessageStreamer(chatModel, chatOptions) - // Then - StepVerifier.create(result) - .expectNext("response") + StepVerifier.create(streamer.stream(listOf(UserMessage("Hi")), emptyList(), emptyList())) + .expectNext("hello") .verifyComplete() } @Test - fun `stream handles multiple message types`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - val messages = listOf( - SystemMessage("You are helpful"), - UserMessage("Hello"), - AssistantMessage("Hi there"), - UserMessage("How are you?") + fun `assembles tool calls without executing them`() { + val call = AssistantMessage.ToolCall("call-1", "function", "lookup", "{\"id\":42}") + val output = AssistantMessage.builder().content("").toolCalls(listOf(call)).build() + every { chatModel.stream(capture(capturedPrompt)) } returns Flux.just( + ChatResponse(listOf(Generation(output))) ) - every { mockStreamSpec.content() } returns Flux.just("response") - - // When - val result = streamer.stream(messages, emptyList(), emptyList()) - - // Then - should complete without error - StepVerifier.create(result) - .expectNext("response") + val tool = Tool.of("lookup", "Lookup an item") { Tool.Result.text("should not execute") } + val streamer = SpringAiLlmMessageStreamer(chatModel, chatOptions) + + StepVerifier.create(streamer.streamInference(listOf(UserMessage("Lookup 42")), listOf(tool))) + .assertNext { event -> + val completion = assertInstanceOf(LlmInferenceStreamEvent.Complete::class.java, event) + val message = assertInstanceOf(AssistantMessageWithToolCalls::class.java, completion.message) + assertEquals("lookup", message.toolCalls.single().name) + assertEquals("{\"id\":42}", message.toolCalls.single().arguments) + } .verifyComplete() - // Verify prompt was built with all messages - verify { mockChatClient.prompt(any()) } + val options = capturedPrompt.captured.options as ToolCallingChatOptions + assertEquals(listOf("lookup"), options.toolCallbacks.map { it.toolDefinition.name() }) } @Test - fun `stream handles empty flux response`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - every { mockStreamSpec.content() } returns Flux.empty() - - // When - val result = streamer.stream(listOf(UserMessage("test")), emptyList(), emptyList()) - - // Then - StepVerifier.create(result) - .verifyComplete() - } - - @Test - fun `stream propagates errors from ChatClient`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - val expectedError = RuntimeException("LLM error") - every { mockStreamSpec.content() } returns Flux.error(expectedError) - - // When - val result = streamer.stream(listOf(UserMessage("test")), emptyList(), emptyList()) - - // Then - StepVerifier.create(result) - .expectError(RuntimeException::class.java) - .verify(Duration.ofSeconds(1)) - } - - @Nested - inner class InspectorIntegrationTests { - - @Test - fun `stream wraps tools with inspectors when provided`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - val tools = listOf( - Tool.of("tool1", "First tool") { _ -> Tool.Result.text("result1") } - ) - val inspector = object : ToolCallInspector { - override fun beforeToolCall(context: BeforeToolCallContext) {} - override fun afterToolCall(context: AfterToolCallContext) {} - } - every { mockStreamSpec.content() } returns Flux.just("response") - - // When - streamer.stream(listOf(UserMessage("test")), tools, listOf(inspector)) - - // Then - verify callbacks were baked into the Prompt's ToolCallingChatOptions - val capturedOptions = capturedPrompt.captured.options as ToolCallingChatOptions - assertTrue(capturedOptions.toolCallbacks.isNotEmpty(), "Callbacks should not be empty") - } - - @Test - fun `stream with multiple inspectors wraps tools correctly`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - val tools = listOf( - Tool.of("tool1", "First tool") { _ -> Tool.Result.text("result1") } - ) - val inspector1 = object : ToolCallInspector {} - val inspector2 = object : ToolCallInspector {} - every { mockStreamSpec.content() } returns Flux.just("response") - - // When - streamer.stream(listOf(UserMessage("test")), tools, listOf(inspector1, inspector2)) - - // Then - val capturedOptions = capturedPrompt.captured.options as ToolCallingChatOptions - assertTrue(capturedOptions.toolCallbacks.isNotEmpty(), "Callbacks should not be empty") - } - - @Test - fun `stream with empty inspectors list does not wrap tools`() { - // Given - val streamer = SpringAiLlmMessageStreamer(mockChatClient, chatOptions) - val tools = listOf( - Tool.of("tool1", "First tool") { _ -> Tool.Result.text("result1") } - ) - every { mockStreamSpec.content() } returns Flux.just("response") - - // When - streamer.stream(listOf(UserMessage("test")), tools, emptyList()) + fun `propagates model streaming errors`() { + every { chatModel.stream(any()) } returns Flux.error(IllegalStateException("failure")) + val streamer = SpringAiLlmMessageStreamer(chatModel, chatOptions) - // Then - callbacks still baked into options - val capturedOptions = capturedPrompt.captured.options as ToolCallingChatOptions - assertTrue(capturedOptions.toolCallbacks.isNotEmpty(), "Callbacks should not be empty") - verify { mockStreamSpec.content() } - } + StepVerifier.create(streamer.streamInference(listOf(UserMessage("Hi")), emptyList())) + .expectErrorMessage("failure") + .verify() } } diff --git a/embabel-agent-autoconfigure/models/embabel-agent-ollama-autoconfigure/pom.xml b/embabel-agent-autoconfigure/models/embabel-agent-ollama-autoconfigure/pom.xml index c628234c3..252340baa 100644 --- a/embabel-agent-autoconfigure/models/embabel-agent-ollama-autoconfigure/pom.xml +++ b/embabel-agent-autoconfigure/models/embabel-agent-ollama-autoconfigure/pom.xml @@ -13,6 +13,10 @@ Ollama Models for Embabel Agent API https://github.com/embabel/embabel-agent + + 1.14.0 + + https://github.com/embabel/embabel-agent scm:git:https://github.com/embabel/embabel-agent.git @@ -36,6 +40,20 @@ embabel-agent-test-internal test + + + dev.langchain4j + langchain4j + ${langchain4j.version} + test + + + + dev.langchain4j + langchain4j-ollama + ${langchain4j.version} + test + diff --git a/embabel-agent-autoconfigure/models/embabel-agent-ollama-autoconfigure/src/test/kotlin/com/embabel/agent/config/models/ollama/LangChain4jStreamingToolLoopIT.kt b/embabel-agent-autoconfigure/models/embabel-agent-ollama-autoconfigure/src/test/kotlin/com/embabel/agent/config/models/ollama/LangChain4jStreamingToolLoopIT.kt new file mode 100644 index 000000000..8d01eb093 --- /dev/null +++ b/embabel-agent-autoconfigure/models/embabel-agent-ollama-autoconfigure/src/test/kotlin/com/embabel/agent/config/models/ollama/LangChain4jStreamingToolLoopIT.kt @@ -0,0 +1,240 @@ +/* + * 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.config.models.ollama + +import com.embabel.agent.api.tool.Tool +import com.embabel.agent.api.tool.callback.ToolCallInspector +import com.embabel.agent.api.tool.progressive.UnfoldingTool +import com.embabel.agent.core.Usage +import com.embabel.agent.spi.loop.ToolInjectionStrategy +import com.embabel.agent.spi.loop.streaming.LlmInferenceStreamEvent +import com.embabel.agent.spi.loop.streaming.LlmMessageStreamer +import com.embabel.agent.spi.loop.streaming.StreamingToolLoop +import com.embabel.chat.AssistantMessage +import com.embabel.chat.AssistantMessageWithToolCalls +import com.embabel.chat.Message +import com.embabel.chat.SystemMessage +import com.embabel.chat.ToolCall +import com.embabel.chat.ToolResultMessage +import com.embabel.chat.UserMessage +import dev.langchain4j.agent.tool.ToolExecutionRequest +import dev.langchain4j.agent.tool.ToolSpecification +import dev.langchain4j.data.message.AiMessage +import dev.langchain4j.data.message.ChatMessage +import dev.langchain4j.data.message.ToolExecutionResultMessage +import dev.langchain4j.model.chat.StreamingChatModel +import dev.langchain4j.model.chat.request.ChatRequest +import dev.langchain4j.model.chat.response.ChatResponse +import dev.langchain4j.model.chat.response.StreamingChatResponseHandler +import dev.langchain4j.model.ollama.OllamaStreamingChatModel +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable +import org.springframework.core.env.Environment +import org.springframework.core.env.StandardEnvironment +import reactor.core.publisher.Flux +import tools.jackson.module.kotlin.jacksonObjectMapper +import java.time.Duration +import java.util.concurrent.atomic.AtomicInteger + +/** + * Proves that the provider-neutral streaming loop can be driven by LangChain4j. + * + * Tool execution remains in [StreamingToolLoop]. The adapter only translates + * messages and tool definitions and assembles LangChain4j's terminal response. + */ +@EnabledIfEnvironmentVariable( + named = "OLLAMA_BASE_URL", + matches = ".+", + disabledReason = "Integration test requires OLLAMA_BASE_URL", +) +class LangChain4jStreamingToolLoopIT { + + private val environment: Environment = StandardEnvironment() + + @Test + fun `executes a tool and continues streaming through LangChain4j`() { + val model = OllamaStreamingChatModel.builder() + .baseUrl(environment.getRequiredProperty("ollama.base-url")) + .modelName(environment.getProperty("ollama.model", "qwen2.5:3b")) + .temperature(0.0) + .numCtx(4096) + .numPredict(128) + .timeout(Duration.ofMinutes(2)) + .build() + val calls = AtomicInteger() + val status = Tool.of("current_status", "Return the current system status") { + calls.incrementAndGet() + Tool.Result.text("GREEN") + } + + val chunks = StreamingToolLoop.create( + messageStreamer = LangChain4jMessageStreamer(model), + objectMapper = jacksonObjectMapper(), + ).execute( + initialMessages = listOf( + SystemMessage("You must call current_status before answering."), + UserMessage("Call current_status now, then report its result in one short sentence."), + ), + initialTools = listOf(status), + ).collectList().block(Duration.ofMinutes(2)) + + assertThat(calls.get()).isEqualTo(1) + assertThat(chunks).isNotNull.isNotEmpty + assertThat(chunks!!.joinToString("")).containsIgnoringCase("green") + } + + @Test + fun `unfolds a child tool for the next LangChain4j streaming turn`() { + val model = OllamaStreamingChatModel.builder() + .baseUrl(environment.getRequiredProperty("ollama.base-url")) + .modelName(environment.getProperty("ollama.model", "qwen2.5:3b")) + .temperature(0.0) + .numCtx(4096) + .numPredict(192) + .timeout(Duration.ofMinutes(2)) + .build() + val childCalls = AtomicInteger() + val child = Tool.of("read_secret_code", "Return the secret code") { + childCalls.incrementAndGet() + Tool.Result.text("ORANGE-42") + } + val facade = UnfoldingTool.of( + name = "secret_tools", + description = "Reveal the tool needed to read the secret code", + innerTools = listOf(child), + childToolUsageNotes = "Call read_secret_code next, then report its result.", + ) + val toolsByTurn = mutableListOf>() + + val chunks = StreamingToolLoop.create( + messageStreamer = LangChain4jMessageStreamer(model) { tools -> + toolsByTurn += tools.map { it.definition.name } + }, + objectMapper = jacksonObjectMapper(), + injectionStrategy = ToolInjectionStrategy.DEFAULT, + ).execute( + initialMessages = listOf( + SystemMessage( + "First call secret_tools. After it reveals a child tool, call that child tool. " + + "Only then report the secret code." + ), + UserMessage("Use the available tools to obtain the secret code."), + ), + initialTools = listOf(facade), + ).collectList().block(Duration.ofMinutes(2)) + + assertThat(toolsByTurn).hasSizeGreaterThanOrEqualTo(3) + assertThat(toolsByTurn[0]).containsExactly("secret_tools") + assertThat(toolsByTurn[1]).containsExactly("read_secret_code") + assertThat(childCalls.get()).isEqualTo(1) + assertThat(chunks).isNotNull.isNotEmpty + assertThat(chunks!!.joinToString("")).containsIgnoringCase("orange-42") + } +} + +private class LangChain4jMessageStreamer( + private val model: StreamingChatModel, + private val onTools: (List) -> Unit = {}, +) : LlmMessageStreamer { + + override fun stream( + messages: List, + tools: List, + toolCallInspectors: List, + ): Flux = streamInference(messages, tools) + .ofType(LlmInferenceStreamEvent.Content::class.java) + .map { it.text } + + override fun streamInference( + messages: List, + tools: List, + ): Flux = Flux.create { sink -> + onTools(tools) + val request = ChatRequest.builder() + .messages(messages.map(::toLangChain4jMessage)) + .toolSpecifications(tools.map(::toLangChain4jTool)) + .build() + + model.chat(request, object : StreamingChatResponseHandler { + override fun onPartialResponse(partialResponse: String) { + sink.next(LlmInferenceStreamEvent.Content(partialResponse)) + } + + override fun onCompleteResponse(completeResponse: ChatResponse) { + sink.next( + LlmInferenceStreamEvent.Complete( + message = toEmbabelMessage(completeResponse.aiMessage()), + usage = completeResponse.tokenUsage()?.let { + Usage(it.inputTokenCount(), it.outputTokenCount(), it) + }, + ) + ) + sink.complete() + } + + override fun onError(error: Throwable) { + sink.error(error) + } + }) + } + + private fun toLangChain4jMessage(message: Message): ChatMessage = when (message) { + is SystemMessage -> dev.langchain4j.data.message.SystemMessage.from(message.content) + is UserMessage -> dev.langchain4j.data.message.UserMessage.from(message.content) + is AssistantMessageWithToolCalls -> AiMessage.from( + message.content, + message.toolCalls.map { + ToolExecutionRequest.builder() + .id(it.id) + .name(it.name) + .arguments(it.arguments) + .build() + }, + ) + + is ToolResultMessage -> ToolExecutionResultMessage.from( + message.toolCallId, + message.toolName, + message.content, + ) + + is AssistantMessage -> AiMessage.from(message.content) + else -> error("Unsupported message type: ${message::class.qualifiedName}") + } + + private fun toLangChain4jTool(tool: Tool): ToolSpecification = + ToolSpecification.builder() + .name(tool.definition.name) + .description(tool.definition.description) + .build() + + private fun toEmbabelMessage(message: AiMessage): Message = + if (message.hasToolExecutionRequests()) { + AssistantMessageWithToolCalls( + content = message.text().orEmpty(), + toolCalls = message.toolExecutionRequests().map { + ToolCall( + id = it.id(), + name = it.name(), + arguments = it.arguments(), + ) + }, + ) + } else { + AssistantMessage(message.text().orEmpty()) + } +} diff --git a/embabel-agent-docs/src/main/asciidoc/reference/streaming/page.adoc b/embabel-agent-docs/src/main/asciidoc/reference/streaming/page.adoc index 5fdc0ed69..f69b5c278 100644 --- a/embabel-agent-docs/src/main/asciidoc/reference/streaming/page.adoc +++ b/embabel-agent-docs/src/main/asciidoc/reference/streaming/page.adoc @@ -10,9 +10,33 @@ This feature is well aligned with Embabel focus on object-oriented programming m - ``StreamingEvent`` - wraps Thinking or user Object - ```StreamingPromptRunnerBuilder``` - runner with streaming capabilities -- Spring Reactive Programming Support for Spring AI ChatClient as underlying infrastructure +- ``LlmMessageStreamer`` - vendor-neutral SPI for one streaming LLM inference +- ``StreamingToolLoop`` - vendor-neutral counterpart of the blocking `ToolLoop` abstraction +- ``DefaultStreamingToolLoop`` - executes tools and updates the conversation between inference turns - All reactive callbacks, such as _doOnNext_, _doOnComplete_, etc. are at developer's disposal +==== Streaming with Tools + +Streaming tool calls use the same provider-neutral execution model as blocking calls. +The `LlmMessageStreamer.streamInference` method advertises the currently available tools and streams +one LLM inference, but it does not execute tools. Embabel assembles the assistant response, executes +requested tools, adds their results to conversation history, and starts the next streaming inference. + +The original three-argument `LlmMessageStreamer.stream` method remains available for compatibility +with provider-managed streaming integrations. Existing implementations can continue returning raw +content chunks. Integrations that participate in Embabel's streaming tool loop override +`streamInference` to expose the assembled assistant message and tool calls. + +Content from every inference turn remains in the returned stream. This includes thinking content +emitted before or between tool calls. Because the tool set is evaluated between turns, +`ToolInjectionStrategy` implementations such as `UnfoldingToolInjectionStrategy` can replace a +facade with its child tools before the next request. + +Provider integrations override `LlmMessageStreamer.streamInference` by emitting incremental +`LlmInferenceStreamEvent.Content` events followed by one `LlmInferenceStreamEvent.Complete` event containing +the assembled assistant message. The completed message must include provider-assigned tool-call IDs, +names, and complete arguments. Provider adapters must not execute those calls themselves. + ==== Example - Simple Thinking and Object Streaming with Callbacks [source,java] ---- @@ -90,7 +114,3 @@ This feature is well aligned with Embabel focus on object-oriented programming m }) .blockLast(Duration.ofSeconds(6000)); ---- - - - -