Rename parallel/reduce -> fork/merge

Author: tiginamaria

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

@@ -134,7 +134,7 @@ internal class AIAgentNode<Input, Output> internal constructor(
  *
  * @param Input The type of input data this node processes and produces as output.
  */
-public open class StartAIAgentNodeBase<Input>() : AIAgentNodeBase<Input, Input>() {
+public open class AIAgentStartNodeBase<Input>() : AIAgentNodeBase<Input, Input>() {
     /**
      * The name of the subgraph associated with the AI agent's starting node.
      *
@@ -160,7 +160,7 @@ public open class StartAIAgentNodeBase<Input>() : AIAgentNodeBase<Input, Input>(
  *
  * @param Output The type of data this node processes and produces.
  */
-public open class FinishAIAgentNodeBase<Output>() : AIAgentNodeBase<Output, Output>() {
+public open class AIAgentFinishNodeBase<Output>() : AIAgentNodeBase<Output, Output>() {
     /**
      * Stores the name of the subgraph associated with the node.
      *
@@ -196,7 +196,7 @@ public open class FinishAIAgentNodeBase<Output>() : AIAgentNodeBase<Output, Outp
  * This node effectively passes its input as-is to the next node in the execution
  * pipeline, allowing downstream nodes to transform or handle the data further.
  */
-internal class StartNode internal constructor() : StartAIAgentNodeBase<String>()
+internal class StartNode internal constructor() : AIAgentStartNodeBase<String>()
 
 /**
  * A specialized implementation of [FinishNode] that finalizes the execution of an AI agent subgraph.
@@ -210,5 +210,5 @@ internal class StartNode internal constructor() : StartAIAgentNodeBase<String>()
  *
  * This node is critical to denote the completion of localized processing within a subgraph context.
  */
-internal class FinishNode internal constructor() : FinishAIAgentNodeBase<String>()
+internal class FinishNode internal constructor() : AIAgentFinishNodeBase<String>()
 

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

@@ -15,8 +15,8 @@ import ai.koog.agents.core.utils.runCatchingCancellable
 @OptIn(InternalAgentsApi::class)
 public class AIAgentStrategy(
     override val name: String,
-    public val nodeStart: StartAIAgentNodeBase<String>,
-    public val nodeFinish: FinishAIAgentNodeBase<String>,
+    public val nodeStart: AIAgentStartNodeBase<String>,
+    public val nodeFinish: AIAgentFinishNodeBase<String>,
     toolSelectionStrategy: ToolSelectionStrategy
 ) : AIAgentSubgraph<String, String>(
     name, nodeStart, nodeFinish, toolSelectionStrategy

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

@@ -28,8 +28,8 @@ import kotlin.uuid.ExperimentalUuidApi
  */
 public open class AIAgentSubgraph<Input, Output>(
     override val name: String,
-    public val start: StartAIAgentNodeBase<Input>,
-    public val finish: FinishAIAgentNodeBase<Output>,
+    public val start: AIAgentStartNodeBase<Input>,
+    public val finish: AIAgentFinishNodeBase<Output>,
     private val toolSelectionStrategy: ToolSelectionStrategy,
 ) : AIAgentNodeBase<Input, Output>() {
     private companion object {
@@ -99,7 +99,7 @@ public open class AIAgentSubgraph<Input, Output>(
                     "Invalid finish node output type: ${currentInput?.let { it::class.simpleName }}"
                 )
             }
-            throw IllegalStateException("${FinishAIAgentNodeBase::class.simpleName} should always return String")
+            throw IllegalStateException("${AIAgentFinishNodeBase::class.simpleName} should always return String")
         }
     }
 

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/builder/AIAgentStrategyBuilder.kt

@@ -14,8 +14,8 @@ public class AIAgentStrategyBuilder(
     private val name: String,
     private val toolSelectionStrategy: ToolSelectionStrategy,
 ) : AIAgentSubgraphBuilderBase<String, String>(), BaseBuilder<AIAgentStrategy> {
-    public override val nodeStart: StartAIAgentNodeBase<String> = StartNode()
-    public override val nodeFinish: FinishAIAgentNodeBase<String> = FinishNode()
+    public override val nodeStart: AIAgentStartNodeBase<String> = StartNode()
+    public override val nodeFinish: AIAgentFinishNodeBase<String> = FinishNode()
 
     override fun build(): AIAgentStrategy {
         return AIAgentStrategy(

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/builder/AIAgentSubgraphBuilder.kt

@@ -27,7 +27,7 @@ public abstract class AIAgentSubgraphBuilderBase<Input, Output> {
      *
      * @param Input The type of input data that this starting node processes.
      */
-    public abstract val nodeStart: StartAIAgentNodeBase<Input>
+    public abstract val nodeStart: AIAgentStartNodeBase<Input>
 
     /**
      * Represents the "finish" node in the AI agent's subgraph structure. This node indicates
@@ -43,7 +43,7 @@ public abstract class AIAgentSubgraphBuilderBase<Input, Output> {
      *
      * @param Output The type of data processed and produced by this node.
      */
-    public abstract val nodeFinish: FinishAIAgentNodeBase<Output>
+    public abstract val nodeFinish: AIAgentFinishNodeBase<Output>
 
     /**
      * Defines a new node in the agent's stage, representing a unit of execution that takes an input and produces an output.
@@ -108,20 +108,20 @@ public abstract class AIAgentSubgraphBuilderBase<Input, Output> {
         vararg nodes: AIAgentNodeBase<Input, Output>,
         dispatcher: CoroutineDispatcher = Dispatchers.Default,
         name: String? = null,
-    ): AIAgentNodeDelegateBase<Input, List<ForkedNodeResult<Input, Output>>> {
-        return AIAgentNodeDelegate(name, ForkNodeBuilder(nodes.asList(), dispatcher))
+    ): AIAgentNodeDelegateBase<Input, List<ForkResult<Input, Output>>> {
+        return AIAgentNodeDelegate(name, AIAgentForkNodeBuilder(nodes.asList(), dispatcher))
     }
 
     /**
      * Creates a node that merges the results of the forked nodes.
-     * @param execute Function to merge the contexts after parallel execution
+     * @param execute Function to merge the contexts and outputs after parallel execution
      * @param name Optional node name
      */
     public fun <Input, Output> merge(
         name: String? = null,
-        execute: suspend AIAgentContextBase.(List<ForkedNodeResult<Input, Output>>) -> Pair<AIAgentContextBase, Output>,
-    ): AIAgentNodeDelegateBase<List<ForkedNodeResult<Input, Output>>, Output> {
-        return AIAgentNodeDelegate(name, MergeNodeBuilder(execute))
+        execute: suspend AIAgentContextBase.(List<ForkResult<Input, Output>>) -> Pair<AIAgentContextBase, Output>,
+    ): AIAgentNodeDelegateBase<List<ForkResult<Input, Output>>, Output> {
+        return AIAgentNodeDelegate(name, AIAgentMergeNodeBuilder(execute))
     }
 
     /**
@@ -136,11 +136,11 @@ public abstract class AIAgentSubgraphBuilderBase<Input, Output> {
     }
 
     /**
-     * Checks if finish node is reachable from start node.
+     * Checks if the finish node is reachable from start node.
      * @param start Starting node
-     * @return True if finish node is reachable
+     * @return True if the finish node is reachable
      */
-    protected fun isFinishReachable(start: StartAIAgentNodeBase<Input>): Boolean {
+    protected fun isFinishReachable(start: AIAgentStartNodeBase<Input>): Boolean {
         val visited = mutableSetOf<AIAgentNodeBase<*, *>>()
 
         fun visit(node: AIAgentNodeBase<*, *>): Boolean {
@@ -173,8 +173,8 @@ public class AIAgentSubgraphBuilder<Input, Output>(
     private val toolSelectionStrategy: ToolSelectionStrategy
 ) : AIAgentSubgraphBuilderBase<Input, Output>(),
     BaseBuilder<AIAgentSubgraphDelegate<Input, Output>> {
-    override val nodeStart: StartAIAgentNodeBase<Input> = StartAIAgentNodeBase()
-    override val nodeFinish: FinishAIAgentNodeBase<Output> = FinishAIAgentNodeBase()
+    override val nodeStart: AIAgentStartNodeBase<Input> = AIAgentStartNodeBase()
+    override val nodeFinish: AIAgentFinishNodeBase<Output> = AIAgentFinishNodeBase()
 
     override fun build(): AIAgentSubgraphDelegate<Input, Output> {
         require(isFinishReachable(nodeStart)) {
@@ -227,8 +227,8 @@ public interface AIAgentSubgraphDelegateBase<Input, Output> {
  */
 public open class AIAgentSubgraphDelegate<Input, Output> internal constructor(
     private val name: String?,
-    public val nodeStart: StartAIAgentNodeBase<Input>,
-    public val nodeFinish: FinishAIAgentNodeBase<Output>,
+    public val nodeStart: AIAgentStartNodeBase<Input>,
+    public val nodeFinish: AIAgentFinishNodeBase<Output>,
     private val toolSelectionStrategy: ToolSelectionStrategy
 ) : AIAgentSubgraphDelegateBase<Input, Output> {
     private var subgraph: AIAgentSubgraph<Input, Output>? = null
@@ -258,7 +258,7 @@ public open class AIAgentSubgraphDelegate<Input, Output> internal constructor(
  * @property context Context of the node on the node termination state
  * @property output Output of the node
  */
-public data class ForkedNodeResult<Input, Output>(
+public data class ForkResult<Input, Output>(
     val nodeName: String,
     val input: Input,
     val context: AIAgentContextBase,
@@ -272,10 +272,10 @@ public data class ForkedNodeResult<Input, Output>(
  * @param dispatcher Coroutine dispatcher to use for parallel execution
  */
 @OptIn(ExperimentalUuidApi::class)
-public class ForkNodeBuilder<Input, Output> internal constructor(
+public class AIAgentForkNodeBuilder<Input, Output> internal constructor(
     private val nodes: List<AIAgentNodeBase<Input, Output>>,
     private val dispatcher: CoroutineDispatcher
-) : AIAgentNodeBuilder<Input, List<ForkedNodeResult<Input, Output>>>(
+) : AIAgentNodeBuilder<Input, List<ForkResult<Input, Output>>>(
     execute = { input ->
         val initialContext: AIAgentContextBase = this
         val mapResults = supervisorScope {
@@ -284,7 +284,7 @@ public class ForkNodeBuilder<Input, Output> internal constructor(
                     val nodeContext = (initialContext as? ai.koog.agents.core.agent.context.AIAgentContext)?.fork()
                         ?: initialContext.fork()
                     val result = node.execute(nodeContext, input)
-                    ForkedNodeResult(node.name, input, nodeContext, result)
+                    ForkResult(node.name, input, nodeContext, result)
                 }
             }
         }
@@ -299,9 +299,9 @@ public class ForkNodeBuilder<Input, Output> internal constructor(
  * @param execute Function to merge the contexts after parallel execution
  */
 @OptIn(ExperimentalUuidApi::class)
-public class MergeNodeBuilder<Input, Output> internal constructor(
-    private val execute: suspend AIAgentContextBase.(List<ForkedNodeResult<Input, Output>>) -> Pair<AIAgentContextBase, Output>,
-) : AIAgentNodeBuilder<List<ForkedNodeResult<Input, Output>>, Output>(
+public class AIAgentMergeNodeBuilder<Input, Output> internal constructor(
+    private val execute: suspend AIAgentContextBase.(List<ForkResult<Input, Output>>) -> Pair<AIAgentContextBase, Output>,
+) : AIAgentNodeBuilder<List<ForkResult<Input, Output>>, Output>(
     execute = { input ->
         val (context, output) = execute(input)
         this.replace(context)

agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/dsl/extension/ParallelNodesTest.kt

@@ -3,7 +3,7 @@ package ai.koog.agents.core.dsl.extension
 import ai.koog.agents.core.agent.AIAgent
 import ai.koog.agents.core.agent.config.AIAgentConfig
 import ai.koog.agents.core.agent.entity.AIAgentStorageKey
-import ai.koog.agents.core.dsl.builder.ForkedNodeResult
+import ai.koog.agents.core.dsl.builder.ForkResult
 import ai.koog.agents.core.dsl.builder.forwardTo
 import ai.koog.agents.core.dsl.builder.strategy
 import ai.koog.agents.core.tools.ToolRegistry
@@ -148,7 +148,7 @@ class ParallelNodesTest {
             )
 
             // Create nodes to verify the context isolation during parallel execution
-            val verifyNode by node<List<ForkedNodeResult<Unit, String>>, String>("verifyNode") { results ->
+            val verifyNode by node<List<ForkResult<Unit, String>>, String>("verifyNode") { results ->
                 results.map {
                     // This node should only see the changes from node1
                     val value1 = it.context.storage.get(testKey1)

agents/agents-test/src/commonMain/kotlin/ai/koog/agents/testing/feature/TestingFeature.kt

@@ -108,7 +108,7 @@ public sealed class NodeReference<Input, Output> {
         override fun resolve(subgraph: AIAgentSubgraph<*, *>): AIAgentNodeBase<Input, Output> {
             val visited = mutableSetOf<String>()
             fun visit(node: AIAgentNodeBase<*, *>): AIAgentNodeBase<Input, Output>? {
-                if (node is FinishAIAgentNodeBase) return null
+                if (node is AIAgentFinishNodeBase) return null
                 if (visited.contains(node.name)) return null
                 visited.add(node.name)
                 if (node.name == name) return node as? AIAgentNodeBase<Input, Output>