[agents] Introduce planner agent type. Implement GOAP

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

@@ -12,7 +12,7 @@
  * @param name The string identifier that uniquely represents the storage key.
  */
 public class AIAgentStorageKey<T : Any>(public val name: String) {
     override fun toString(): String = "${super.toString()}(name=$name)"
 }
 
 /**
@@ -21,15 +21,15 @@
  * @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 @@
  * @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 @@
  *
  * @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
+}

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

@@ -0,0 +1,34 @@
+package ai.koog.agents.planner
+
+import ai.koog.agents.core.feature.config.FeatureConfig
+import ai.koog.agents.core.feature.pipeline.AIAgentPipeline
+import kotlinx.datetime.Clock
+
+/**
+ * Represents a specific implementation of an AI agent pipeline that uses a planner approach.
+ *
+ * @property clock The clock used for time-based operations within the pipeline
+ */
+public class AIAgentPlannerPipeline(clock: Clock = Clock.System) : AIAgentPipeline(clock) {
+
+    /**
+     * Installs a non-graph feature into the pipeline with the provided configuration.
+     *
+     * @param TConfig The type of the feature configuration
+     * @param TFeature The type of the feature being installed
+     * @param feature The feature implementation to be installed
+     * @param configure A lambda to customize the feature configuration
+     */
+    public fun <TConfig : FeatureConfig, TFeature : Any> install(
+        feature: AIAgentPlannerFeature<TConfig, TFeature>,
+        configure: TConfig.() -> Unit,
+    ) {
+        val featureConfig = feature.createInitialConfig().apply { configure() }
+        val featureImpl = feature.install(
+            config = featureConfig,
+            pipeline = this,
+        )
+
+        super.install(feature.key, featureConfig, featureImpl)
+    }
+}

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

@@ -0,0 +1,38 @@
+package ai.koog.agents.planner
+
+import ai.koog.agents.core.agent.context.AIAgentFunctionalContext
+import ai.koog.agents.core.agent.context.with
+import ai.koog.agents.core.agent.entity.AIAgentStrategy
+import kotlin.coroutines.cancellation.CancellationException
+
+/**
+ * Represents a planner strategy for an AI agent.
+ * @param State The type of the state.
+ * @param Plan The type of the plan.
+ * @param name The name of the strategy.
+ * @param planner The instance of the planner defining the exact planner strategy.
+ */
+public class AIAgentPlannerStrategy<State, Plan>(
+    override val name: String,
+    private val planner: AIAgentPlanner<State, Plan>,
+) : AIAgentStrategy<State, State, AIAgentFunctionalContext> {
+    override suspend fun execute(
+        context: AIAgentFunctionalContext,
+        input: State
+    ): State {
+        return try {
+            context.with(partName = name) { executionInfo, eventId ->
+                context.pipeline.onStrategyStarting(eventId, executionInfo, this, context)
+                val result = planner.execute(context, input)
+                context.pipeline.onStrategyCompleted(eventId, executionInfo, this, context, result, planner.stateType)
+
+                result
+            }
+        } catch (e: CancellationException) {
+            throw e
+        } catch (e: Exception) {
+            context.environment.reportProblem(e)
+            throw e
+        }
+    }
+}