[agents] Introduce planner agent type. Implement GOAP (#1232)

Introduce new "planner" agent type, in addition to existing "graph" and
"functional". These agents can work iteratively by building, executing
and updating the plan. Two planner strategies are provided out of the
box - simple LLM and GOAP. New agent type and built-in planners are
implemented as a separate optional dependency (module) -
`agents-planner`

---------

Co-authored-by: Andrey Bragin <[email protected]>

Author: antoniibelyshev

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

@@ -4,19 +4,28 @@ import ai.koog.agents.core.utils.ActiveProperty
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 
+/**
+ * Represents the state of an AI agent.
+ */
 @OptIn(ExperimentalStdlibApi::class)
-internal class AIAgentState internal constructor(
+public class AIAgentState(
     iterations: Int = 0,
 ) : AutoCloseable {
-    var iterations: Int by ActiveProperty(iterations) { isActive }
+    /**
+     * The number of iterations that have been completed since the agent was created.
+     */
+    public var iterations: Int by ActiveProperty(iterations) { isActive }
 
     private var isActive = true
 
     override fun close() {
         isActive = false
     }
 
-    internal fun copy(): AIAgentState {
+    /**
+     * Creates a copy of the current state.
+     */
+    public fun copy(): AIAgentState {
         return AIAgentState(
             iterations = iterations
         )
@@ -33,12 +42,16 @@ internal class AIAgentState internal constructor(
  * @constructor Creates a new instance of AIAgentStateManager with the initial state,
  * defaulting to a new `AIAgentState` if not provided.
  */
-public class AIAgentStateManager internal constructor(
+public class AIAgentStateManager(
     private var state: AIAgentState = AIAgentState()
 ) {
     private val mutex = Mutex()
 
-    internal suspend fun <T> withStateLock(block: suspend (AIAgentState) -> T): T = mutex.withLock {
+    /**
+     * Executes the provided suspending [block] of code with exclusive access to the current state.
+     * @return The result of [block].
+     */
+    public suspend fun <T> withStateLock(block: suspend (AIAgentState) -> T): T = mutex.withLock {
         val result = block(state)
         val newState = AIAgentState(
             iterations = state.iterations

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

@@ -21,15 +21,15 @@ public class AIAgentStorageKey<T : Any>(public val name: String) {
  * @param name The name of the storage key, used to uniquely identify it.
  * @return A new instance of [AIAgentStorageKey] for the specified type.
  */
-public inline fun <reified T : Any> createStorageKey(name: String): AIAgentStorageKey<T> = AIAgentStorageKey<T>(name)
+public fun <T : Any> createStorageKey(name: String): AIAgentStorageKey<T> = AIAgentStorageKey(name)
 
 /**
  * Concurrent-safe key-value storage for an agent.
  * You can create typed keys for your data using the [createStorageKey] function and
  * set and retrieve data using it by calling [set] and [get].
  *
  */
-public class AIAgentStorage internal constructor() {
+public class AIAgentStorage {
     private val mutex = Mutex()
     private val storage = mutableMapOf<AIAgentStorageKey<*>, Any>()
 

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

@@ -222,7 +222,7 @@ public open class AIAgentSubgraph<TInput, TOutput>(
 
     @OptIn(InternalAgentsApi::class)
     private suspend fun executeWithInnerContext(context: AIAgentGraphContextBase, initialInput: TInput): TOutput? {
-        logger.info { formatLog(context, "Executing subgraph '$name'") }
+        logger.debug { formatLog(context, "Executing subgraph '$name'") }
 
         var currentNode: AIAgentNodeBase<*, *> = start
         var currentInput: Any? = initialInput

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

@@ -27,7 +27,7 @@ public open class AIAgentException(problem: String, throwable: Throwable? = null
  * @param node The node in which the agent becomes stuck.
  * @param output The output produced by the node that doesn't match any edge conditions.
  */
-internal class AIAgentStuckInTheNodeException(node: AIAgentNodeBase<*, *>, output: Any?) :
+public class AIAgentStuckInTheNodeException(node: AIAgentNodeBase<*, *>, output: Any?) :
     AIAgentException(
         "When executing agent graph, stuck in node ${node.name} " +
             "because output $output doesn't match any condition on available edges."
@@ -57,5 +57,5 @@ public class AIAgentMaxNumberOfIterationsReachedException(maxNumberOfIterations:
  *
  * @param message A descriptive message explaining the reason for termination.
  */
-internal class AIAgentTerminationByClientException(message: String) :
+public class AIAgentTerminationByClientException(message: String) :
     AIAgentException("Agent was canceled by the client ($message)")

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/environment/ContextualAgentEnvironment.kt

@@ -7,12 +7,13 @@ import io.github.oshai.kotlinlogging.KotlinLogging
 import kotlin.uuid.ExperimentalUuidApi
 import kotlin.uuid.Uuid
 
-internal class ContextualAgentEnvironment(
+@Suppress("MissingKDocForPublicAPI")
+public class ContextualAgentEnvironment(
     private val environment: AIAgentEnvironment,
     private val context: AIAgentContext,
 ) : AIAgentEnvironment {
 
-    companion object {
+    private companion object {
         private val logger = KotlinLogging.logger { }
     }
 

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/environment/GenericAgentEnvironment.kt

@@ -8,7 +8,10 @@ import ai.koog.agents.core.tools.annotations.InternalAgentToolsApi
 import ai.koog.prompt.message.Message
 import io.github.oshai.kotlinlogging.KLogger
 
-internal class GenericAgentEnvironment(
+/**
+ * Represents base agent environment with generic abstractions.
+ */
+public class GenericAgentEnvironment(
     private val agentId: String,
     private val logger: KLogger,
     private val toolRegistry: ToolRegistry,

agents/agents-planner/Module.md

@@ -0,0 +1,10 @@
+# Module agents-planner
+
+Module for implementing planning capabilities in AI agents.
+
+## Overview
+
+The agents-planner module provides components for creating AI agents that can plan and execute multi-step tasks through iterative planning cycles.
+Each planning cycle builds or updates a plan based on the current state, executes a step from the plan, and checks if the goal has been achieved.
+
+This module also provides GOAP (Goal-Oriented Action Planning) that use search algorithms to find optimal action sequences based on predefined goals and actions.

agents/agents-planner/build.gradle.kts

@@ -0,0 +1,44 @@
+import ai.koog.gradle.publish.maven.Publishing.publishToMaven
+
+group = rootProject.group
+version = rootProject.version
+
+plugins {
+    id("ai.kotlin.multiplatform")
+    alias(libs.plugins.kotlin.serialization)
+}
+
+kotlin {
+    sourceSets {
+        commonMain {
+            dependencies {
+                api(project(":agents:agents-core"))
+                api(project(":agents:agents-ext"))
+                api(project(":agents:agents-features:agents-features-memory"))
+                api(project(":agents:agents-utils"))
+                api(project(":prompt:prompt-executor:prompt-executor-model"))
+                api(project(":prompt:prompt-structure"))
+
+                api(project(":prompt:prompt-markdown"))
+            }
+        }
+
+        commonTest {
+            dependencies {
+                implementation(kotlin("test"))
+                implementation(libs.kotlinx.coroutines.test)
+                implementation(project(":agents:agents-test"))
+            }
+        }
+
+        jvmTest {
+            dependencies {
+                implementation(kotlin("test-junit5"))
+            }
+        }
+    }
+
+    explicitApi()
+}
+
+publishToMaven()

agents/agents-planner/src/commonMain/kotlin/ai/koog/agents/planner/AIAgentPlanner.kt

@@ -0,0 +1,100 @@
+package ai.koog.agents.planner
+
+import ai.koog.agents.core.agent.context.AIAgentContext
+import ai.koog.agents.core.agent.context.AIAgentFunctionalContext
+import ai.koog.agents.core.agent.exception.AIAgentMaxNumberOfIterationsReachedException
+import io.github.oshai.kotlinlogging.KotlinLogging
+import kotlin.reflect.KType
+
+/**
+ * An abstract base planner component, which can be used to implement different types of AI agent planner execution flows.
+ *
+ * An entry point is an [execute] method, which accepts an initial arbitrary [State] and returns the final [State] after the execution.
+ *
+ * Planner flow works as follows:
+ * 1. Build a plan: [buildPlan]
+ * 2. Execute a step in the plan: [executeStep]
+ * 3. Repeat steps 1 and 2 until the plan is considered completed. Then the final [State] is returned.
+ *
+ * @property stateType [KType] of the [State].
+ */
+public abstract class AIAgentPlanner<State, Plan>(
+    public val stateType: KType,
+) {
+    private companion object {
+        private val logger = KotlinLogging.logger { }
+    }
+
+    /**
+     * Builds a plan
+     */
+    protected abstract suspend fun buildPlan(
+        context: AIAgentFunctionalContext,
+        state: State,
+        plan: Plan?
+    ): Plan
+
+    /**
+     * Executes a step in the plan.
+     */
+    protected abstract suspend fun executeStep(
+        context: AIAgentFunctionalContext,
+        state: State,
+        plan: Plan
+    ): State
+
+    /**
+     * Checks if the plan is completed.
+     */
+    protected abstract suspend fun isPlanCompleted(
+        context: AIAgentFunctionalContext,
+        state: State,
+        plan: Plan
+    ): Boolean
+
+    /**
+     * Executes the main loop for the planner, which involves building and executing plans iteratively until
+     * the plan is considered successfully completed or a max number of iterations is reached.
+     *
+     * @param context AI Agent's context
+     * @param input The initial state to be used as the starting point for the execution process.
+     * @return The final state after the execution of the plans.
+     * @throws AIAgentMaxNumberOfIterationsReachedException If the maximum number of iterations defined in the agent's
+     * configuration is exceeded.
+     */
+    public suspend fun execute(
+        context: AIAgentFunctionalContext,
+        input: State
+    ): State {
+        logger.debug { formatLog(context, "Starting planner execution") }
+        var state = input
+        var plan: Plan = buildPlan(context, state, null)
+
+        while (!isPlanCompleted(context, state, plan)) {
+            val iterations = context.stateManager.withStateLock { state ->
+                if (++state.iterations > context.config.maxAgentIterations) {
+                    logger.error {
+                        formatLog(
+                            context,
+                            "Max iterations limit (${context.config.maxAgentIterations}) reached"
+                        )
+                    }
+                    throw AIAgentMaxNumberOfIterationsReachedException(context.config.maxAgentIterations)
+                }
+
+                state.iterations
+            }
+
+            logger.debug { formatLog(context, "Executing plan step #$iterations") }
+            state = executeStep(context, state, plan)
+            plan = buildPlan(context, state, plan)
+            logger.debug { formatLog(context, "Finished executing plan step #$iterations") }
+        }
+
+        logger.debug { formatLog(context, "Finished planner execution") }
+        return state
+    }
+
+    private fun formatLog(context: AIAgentContext, message: String): String =
+        "$message [${context.strategyName}, ${context.runId}]"
+}

agents/agents-planner/src/commonMain/kotlin/ai/koog/agents/planner/AIAgentPlannerFeature.kt

@@ -0,0 +1,18 @@
+package ai.koog.agents.planner
+
+import ai.koog.agents.core.feature.AIAgentFeature
+import ai.koog.agents.core.feature.config.FeatureConfig
+
+/**
+ * Represents a planner-specific AI agent feature that can be installed into an [AIAgentPlannerPipeline].
+ *
+ * @param TConfig The type of configuration required for the feature, extending [FeatureConfig].
+ * @param TFeatureImpl The type representing the concrete implementation of the feature.
+ */
+public interface AIAgentPlannerFeature<TConfig : FeatureConfig, TFeatureImpl : Any> : AIAgentFeature<TConfig, TFeatureImpl> {
+    /**
+     * Installs the feature into the specified [pipeline].
+     * @return The implementation of the feature.
+     */
+    public fun install(config: TConfig, pipeline: AIAgentPlannerPipeline): TFeatureImpl
+}