Refactor streaming api to support tool calls

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/AgentLoopContextExt.kt

@@ -11,6 +11,7 @@ import ai.koog.agents.core.tools.ToolArgs
 import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.agents.core.tools.ToolResult
 import ai.koog.prompt.message.Message
+import ai.koog.prompt.streaming.StreamFrame
 import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredDataDefinition
 import ai.koog.prompt.structure.StructuredResponse
@@ -171,12 +172,12 @@ public suspend inline fun <reified T> AIAgentFunctionalContext.requestLLMStructu
  *
  * @param message The content of the message to be sent to the LLM.
  * @param structureDefinition Optional structure to guide the LLM response.
- * @return A flow of string chunks from the LLM response.
+ * @return A flow of [StreamFrame] objects from the LLM response.
  */
 public suspend fun AIAgentFunctionalContext.requestLLMStreaming(
     message: String,
     structureDefinition: StructuredDataDefinition? = null
-): Flow<String> {
+): Flow<StreamFrame> {
     return llm.writeSession {
         updatePrompt {
             user(message)

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMWriteSession.kt

@@ -16,6 +16,7 @@ import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.params.LLMParams
+import ai.koog.prompt.streaming.StreamFrame
 import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredDataDefinition
 import ai.koog.prompt.structure.StructuredOutputConfig
@@ -477,9 +478,9 @@ public class AIAgentLLMWriteSession internal constructor(
      *
      * @param definition an optional parameter to define a structured data format. When provided, it will be used
      * in constructing the prompt for the language model request.
-     * @return a flow of strings that streams the responses from the language model.
+     * @return a flow of `StreamingFrame` objects that streams the responses from the language model.
      */
-    public fun requestLLMStreaming(definition: StructuredDataDefinition? = null): Flow<String> {
+    public fun requestLLMStreaming(definition: StructuredDataDefinition? = null): Flow<StreamFrame> {
         if (definition != null) {
             val prompt = prompt(prompt, clock) {
                 user {
@@ -488,7 +489,6 @@ public class AIAgentLLMWriteSession internal constructor(
             }
             this.prompt = prompt
         }
-
-        return executor.executeStreaming(prompt, model)
+        return executor.executeStreaming(prompt, model, tools)
     }
 }

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentNodes.kt

@@ -17,11 +17,14 @@ import ai.koog.prompt.dsl.PromptBuilder
 import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
+import ai.koog.prompt.streaming.StreamFrame
+import ai.koog.prompt.streaming.toMessageResponses
 import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredDataDefinition
 import ai.koog.prompt.structure.StructuredOutputConfig
 import ai.koog.prompt.structure.StructuredResponse
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.toList
 
 /**
  * A pass-through node that does nothing and returns input as output
@@ -240,7 +243,7 @@ public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStr
 public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStreaming(
     name: String? = null,
     structureDefinition: StructuredDataDefinition? = null,
-    transformStreamData: suspend (Flow<String>) -> Flow<T>
+    transformStreamData: suspend (Flow<StreamFrame>) -> Flow<T>
 ): AIAgentNodeDelegate<String, Flow<T>> =
     node(name) { message ->
         llm.writeSession {
@@ -264,7 +267,7 @@ public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStreaming(
 public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStreaming(
     name: String? = null,
     structureDefinition: StructuredDataDefinition? = null,
-): AIAgentNodeDelegate<String, Flow<String>> = nodeLLMRequestStreaming(name, structureDefinition) { it }
+): AIAgentNodeDelegate<String, Flow<StreamFrame>> = nodeLLMRequestStreaming(name, structureDefinition) { it }
 
 /**
  * A node that appends a user message to the LLM prompt and gets multiple LLM responses with tool calls enabled.
@@ -315,6 +318,43 @@ public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMCompressHi
     input
 }
 
+/**
+ * A node that performs LLM streaming, collects all stream frames, converts them to response messages,
+ * and updates the prompt with the results.
+ *
+ * This node is useful when you want to:
+ * - Stream responses from the LLM for real-time feedback
+ * - Collect the complete streamed response as messages
+ * - Automatically update the conversation history with the streamed responses
+ *
+ * The node will:
+ * 1. Initiate a streaming request to the LLM
+ * 2. Collect all stream frames (text, tool calls, etc.)
+ * 3. Convert the collected frames into proper Message.Response objects
+ * 4. Update the prompt with these messages for conversation continuity
+ * 5. Return the collected messages
+ *
+ * @param T The type of input this node accepts (passed through without modification)
+ * @param name Optional node name for identification in the agent graph
+ * @param structureDefinition Optional structure definition to guide the LLM's response format
+ * @return A node delegate that accepts input of type T and returns a list of response messages
+ *
+ * @see nodeLLMRequestStreaming for streaming without automatic prompt updates
+ * @see ai.koog.agents.core.agent.session.AIAgentLLMWriteSession.requestLLMStreaming for the underlying streaming functionality
+ */
+@AIAgentBuilderDslMarker
+public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStreamingAndSendResults(
+    name: String? = null,
+    structureDefinition: StructuredDataDefinition? = null
+): AIAgentNodeDelegate<T, List<Message.Response>> = node(name) { input ->
+    llm.writeSession {
+        requestLLMStreaming(structureDefinition)
+            .toList()
+            .toMessageResponses()
+            .also { updatePrompt { messages(it) } }
+    }
+}
+
 // ==========
 // Tool nodes
 // ==========

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/AIAgentPipeline.kt

@@ -9,6 +9,8 @@ import ai.koog.agents.core.environment.AIAgentEnvironment
 import ai.koog.agents.core.feature.config.FeatureConfig
 import ai.koog.agents.core.feature.handler.AfterLLMCallContext
 import ai.koog.agents.core.feature.handler.AfterLLMCallHandler
+import ai.koog.agents.core.feature.handler.AfterStreamContext
+import ai.koog.agents.core.feature.handler.AfterStreamHandler
 import ai.koog.agents.core.feature.handler.AgentBeforeCloseContext
 import ai.koog.agents.core.feature.handler.AgentBeforeCloseHandler
 import ai.koog.agents.core.feature.handler.AgentContextHandler
@@ -23,13 +25,20 @@ import ai.koog.agents.core.feature.handler.AgentTransformEnvironmentContext
 import ai.koog.agents.core.feature.handler.BeforeAgentStartedHandler
 import ai.koog.agents.core.feature.handler.BeforeLLMCallContext
 import ai.koog.agents.core.feature.handler.BeforeLLMCallHandler
+import ai.koog.agents.core.feature.handler.BeforeStreamContext
+import ai.koog.agents.core.feature.handler.BeforeStreamHandler
 import ai.koog.agents.core.feature.handler.ExecuteLLMHandler
 import ai.koog.agents.core.feature.handler.ExecuteToolHandler
 import ai.koog.agents.core.feature.handler.StrategyFinishContext
 import ai.koog.agents.core.feature.handler.StrategyFinishedHandler
 import ai.koog.agents.core.feature.handler.StrategyHandler
 import ai.koog.agents.core.feature.handler.StrategyStartContext
 import ai.koog.agents.core.feature.handler.StrategyStartedHandler
+import ai.koog.agents.core.feature.handler.StreamErrorContext
+import ai.koog.agents.core.feature.handler.StreamErrorHandler
+import ai.koog.agents.core.feature.handler.StreamFrameContext
+import ai.koog.agents.core.feature.handler.StreamFrameHandler
+import ai.koog.agents.core.feature.handler.StreamHandler
 import ai.koog.agents.core.feature.handler.ToolCallContext
 import ai.koog.agents.core.feature.handler.ToolCallFailureContext
 import ai.koog.agents.core.feature.handler.ToolCallFailureHandler
@@ -46,6 +55,7 @@ import ai.koog.prompt.dsl.ModerationResult
 import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
+import ai.koog.prompt.streaming.StreamFrame
 import io.github.oshai.kotlinlogging.KotlinLogging
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -117,6 +127,12 @@ public abstract class AIAgentPipeline {
      */
     protected val executeLLMHandlers: MutableMap<AIAgentStorageKey<*>, ExecuteLLMHandler> = mutableMapOf()
 
+    /**
+     * Map of feature storage keys to their stream handlers.
+     * These handlers manage the streaming lifecycle events (before, during, and after streaming).
+     */
+    protected val streamHandlers: MutableMap<AIAgentStorageKey<*>, StreamHandler> = mutableMapOf()
+
     internal suspend fun prepareFeatures() {
         withContext(featurePrepareDispatcher) {
             registeredFeatures.values.forEach { featureConfig ->
@@ -202,6 +218,71 @@ public abstract class AIAgentPipeline {
         agentHandlers.values.forEach { handler -> handler.agentRunErrorHandler.handle(eventContext) }
     }
 
+    /**
+     * Invoked before streaming from a language model begins.
+     *
+     * This method notifies all registered stream handlers that streaming is about to start,
+     * allowing them to perform preprocessing or logging operations.
+     *
+     * @param runId The unique identifier for this streaming session
+     * @param prompt The prompt being sent to the language model
+     * @param model The language model being used for streaming
+     * @param tools The list of available tool descriptors for this streaming session
+     */
+    public suspend fun onBeforeStream(runId: String, prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>) {
+        val eventContext = BeforeStreamContext(runId, prompt, model, tools)
+        streamHandlers.values.forEach { handler -> handler.beforeStreamHandler.handle(eventContext) }
+    }
+
+    /**
+     * Invoked when a stream frame is received during the streaming process.
+     *
+     * This method notifies all registered stream handlers about each incoming stream frame,
+     * allowing them to process, transform, or aggregate the streaming content in real-time.
+     *
+     * @param runId The unique identifier for this streaming session
+     * @param streamFrame The individual stream frame containing partial response data
+     */
+    public suspend fun onStreamFrame(runId: String, streamFrame: StreamFrame) {
+        val eventContext = StreamFrameContext(runId, streamFrame)
+        streamHandlers.values.forEach { handler -> handler.streamFrameHandler.handle(eventContext) }
+    }
+
+    /**
+     * Invoked if an error occurs during the streaming process.
+     *
+     * This method notifies all registered stream handlers about the streaming error,
+     * allowing them to handle or log the error.
+     *
+     * @param runId The unique identifier for this streaming session
+     * @param throwable The exception that occurred during streaming, if applicable
+     */
+    public suspend fun onStreamError(runId: String, throwable: Throwable) {
+        val eventContext = StreamErrorContext(runId, throwable)
+        streamHandlers.values.forEach { handler -> handler.streamErrorHandler.handle(eventContext) }
+    }
+
+    /**
+     * Invoked after streaming from a language model completes.
+     *
+     * This method notifies all registered stream handlers that streaming has finished,
+     * allowing them to perform post-processing, cleanup, or final logging operations.
+     *
+     * @param runId The unique identifier for this streaming session
+     * @param prompt The prompt that was sent to the language model
+     * @param model The language model that was used for streaming
+     * @param tools The list of tool descriptors that were available for this streaming session
+     */
+    public suspend fun onAfterStream(
+        runId: String,
+        prompt: Prompt,
+        model: LLModel,
+        tools: List<ToolDescriptor>
+    ) {
+        val eventContext = AfterStreamContext(runId, prompt, model, tools)
+        streamHandlers.values.forEach { handler -> handler.afterStreamHandler.handle(eventContext) }
+    }
+
     /**
      * Invoked before an agent is closed to perform necessary pre-closure operations.
      *
@@ -661,6 +742,104 @@ public abstract class AIAgentPipeline {
         }
     }
 
+    /**
+     * Intercepts streaming operations before they begin to modify or log the streaming request.
+     *
+     * This method allows features to hook into the streaming pipeline before streaming starts,
+     * enabling preprocessing, validation, or logging of streaming requests.
+     *
+     * @param interceptContext The context containing the feature and its implementation
+     * @param handle The handler that processes before-stream events
+     *
+     * Example:
+     * ```
+     * pipeline.interceptBeforeStream(InterceptContext) { eventContext ->
+     *     logger.info("About to start streaming with prompt: ${eventContext.prompt.messages.last().content}")
+     * }
+     * ```
+     */
+    public fun <TFeature : Any> interceptBeforeStream(
+        interceptContext: InterceptContext<TFeature>,
+        handle: suspend TFeature.(eventContext: BeforeStreamContext) -> Unit
+    ) {
+        val existingHandler = streamHandlers.getOrPut(interceptContext.feature.key) { StreamHandler() }
+
+        existingHandler.beforeStreamHandler = BeforeStreamHandler { eventContext: BeforeStreamContext ->
+            with(interceptContext.featureImpl) { handle(eventContext) }
+        }
+    }
+
+    /**
+     * Intercepts stream frames as they are received during the streaming process.
+     *
+     * This method allows features to process individual stream frames in real-time,
+     * enabling monitoring, transformation, or aggregation of streaming content.
+     *
+     * @param interceptContext The context containing the feature and its implementation
+     * @param handle The handler that processes stream frame events
+     *
+     * Example:
+     * ```
+     * pipeline.interceptOnStreamFrame(InterceptContext) { eventContext ->
+     *     logger.debug("Received stream frame: ${eventContext.streamFrame}")
+     * }
+     * ```
+     */
+    public fun <TFeature : Any> interceptOnStreamFrame(
+        interceptContext: InterceptContext<TFeature>,
+        handle: suspend TFeature.(eventContext: StreamFrameContext) -> Unit
+    ) {
+        val existingHandler = streamHandlers.getOrPut(interceptContext.feature.key) { StreamHandler() }
+
+        existingHandler.streamFrameHandler = StreamFrameHandler { eventContext: StreamFrameContext ->
+            with(interceptContext.featureImpl) { handle(eventContext) }
+        }
+    }
+
+    /**
+     * Intercepts errors during the streaming process.
+     *
+     * @param interceptContext The context containing the feature and its implementation
+     * @param handle The handler that processes stream errors
+     */
+    public fun <TFeature : Any> interceptOnStreamError(
+        interceptContext: InterceptContext<TFeature>,
+        handle: suspend TFeature.(eventContext: StreamErrorContext) -> Unit
+    ) {
+        val existingHandler = streamHandlers.getOrPut(interceptContext.feature.key) { StreamHandler() }
+
+        existingHandler.streamErrorHandler = StreamErrorHandler { eventContext: StreamErrorContext ->
+            with(interceptContext.featureImpl) { handle(eventContext) }
+        }
+    }
+
+    /**
+     * Intercepts streaming operations after they complete to perform post-processing or cleanup.
+     *
+     * This method allows features to hook into the streaming pipeline after streaming finishes,
+     * enabling post-processing, cleanup, or final logging of the streaming session.
+     *
+     * @param interceptContext The context containing the feature and its implementation
+     * @param handle The handler that processes after-stream events
+     *
+     * Example:
+     * ```
+     * pipeline.interceptAfterStream(InterceptContext) { eventContext ->
+     *     logger.info("Streaming completed for run: ${eventContext.runId}")
+     * }
+     * ```
+     */
+    public fun <TFeature : Any> interceptAfterStream(
+        interceptContext: InterceptContext<TFeature>,
+        handle: suspend TFeature.(eventContext: AfterStreamContext) -> Unit
+    ) {
+        val existingHandler = streamHandlers.getOrPut(interceptContext.feature.key) { StreamHandler() }
+
+        existingHandler.afterStreamHandler = AfterStreamHandler { eventContext: AfterStreamContext ->
+            with(interceptContext.featureImpl) { handle(eventContext) }
+        }
+    }
+
     /**
      * Intercepts and handles tool calls for the specified feature and its implementation.
      * Updates the tool call handler for the given feature key with a custom handler.