Introduced AIAgentStrategy

Author: Rizzen

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

@@ -1,7 +1,7 @@
 package ai.koog.agents.core.agent
 
 import ai.koog.agents.core.agent.config.AIAgentConfigBase
-import ai.koog.agents.core.agent.context.AIAgentContextBase
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.context.AIAgentLLMContext
 import ai.koog.agents.core.agent.entity.AIAgentStateManager
 import ai.koog.agents.core.agent.entity.AIAgentStorage
@@ -22,7 +22,7 @@ public class EssentialAIAgentContext(
     override val stateManager: AIAgentStateManager,
     override val storage: AIAgentStorage,
     override val strategyName: String
-) : AIAgentContextBase {
+) : AIAgentContext {
 
     private val storeMap: MutableMap<AIAgentStorageKey<*>, Any> = mutableMapOf()
 

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/context/AIAgentContextBase.kt renamed to agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/context/AIAgentContext.kt

@@ -9,12 +9,12 @@ import ai.koog.agents.core.feature.AIAgentFeature
 import ai.koog.prompt.message.Message
 
 /**
- * The [AIAgentContextBase] interface represents the context of an AI agent in the lifecycle.
+ * The [AIAgentContext] interface represents the context of an AI agent in the lifecycle.
  * It provides access to the environment, configuration, LLM context, state management, storage, and other
  * metadata necessary for the operation of the agent.
  * Additionally, it supports features for custom workflows and extensibility.
  */
-public interface AIAgentContextBase {
+public interface AIAgentContext {
 
     /**
      * Represents the environment in which the agent operates.
@@ -153,9 +153,9 @@ public interface AIAgentContextBase {
 }
 
 /**
- * Utility function to get [AIAgentContextBase.agentInput] and try to cast it to some expected type.
+ * Utility function to get [AIAgentContext.agentInput] and try to cast it to some expected type.
  *
  * @throws ClassCastException If agent input can't be cast to [T]
  */
-public inline fun <reified T> AIAgentContextBase.agentInput(): T =
+public inline fun <reified T> AIAgentContext.agentInput(): T =
     agentInput as? T ?: throw ClassCastException("Can't cast agent input to ${T::class}. Agent input: $agentInput")

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

@@ -29,7 +29,7 @@ import kotlin.reflect.KType
  * to handle graph-specific logic, such as node traversal, input/output management,
  * and handling complex dependencies between graph nodes.
  */
-public interface AIAgentGraphContextBase : AIAgentContextBase {
+public interface AIAgentGraphContextBase : AIAgentContext {
     /**
      * Represents the execution pipeline associated with the AI agent operating
      * in a graph-based context.
@@ -97,7 +97,7 @@ public interface AIAgentGraphContextBase : AIAgentContextBase {
      *
      * @param context The context to replace the current context with.
      */
-    public suspend fun replace(context: AIAgentContextBase)
+    public suspend fun replace(context: AIAgentContext)
 }
 
 /**
@@ -232,14 +232,14 @@ public class AIAgentGraphContext(
     }
 
     /**
-     * Creates a new instance of [AIAgentContextBase] with an updated list of tools, replacing the current tools
+     * Creates a new instance of [AIAgentContext] with an updated list of tools, replacing the current tools
      * in the LLM context with the provided list.
      *
      * @param tools The new list of tools to be used in the LLM context, represented as [ToolDescriptor] objects.
-     * @return A new instance of [AIAgentContextBase] with the updated tools configuration.
+     * @return A new instance of [AIAgentContext] with the updated tools configuration.
      */
     @InternalAgentsApi
-    public fun copyWithTools(tools: List<ToolDescriptor>): AIAgentContextBase {
+    public fun copyWithTools(tools: List<ToolDescriptor>): AIAgentContext {
         return this.copy(llm = llm.copy(tools = tools))
     }
 
@@ -249,7 +249,7 @@ public class AIAgentGraphContext(
         stateManager = this.stateManager.copy(),
     )
 
-    override suspend fun replace(context: AIAgentContextBase) {
+    override suspend fun replace(context: AIAgentContext) {
         mutableAIAgentContext.replace(
             context.llm,
             context.stateManager,
@@ -279,7 +279,7 @@ public val agentContextDataAdditionalKey: AIAgentStorageKey<AgentContextData> =
  * @param data The context-specific data to be stored for later retrieval or use within the agent context.
  */
 @InternalAgentsApi
-public fun AIAgentContextBase.store(data: AgentContextData) {
+public fun AIAgentContext.store(data: AgentContextData) {
     this.store(agentContextDataAdditionalKey, data)
 }
 
@@ -295,7 +295,7 @@ public fun AIAgentContextBase.store(data: AgentContextData) {
  * @return The agent context data, or null if no context data is associated.
  */
 @InternalAgentsApi
-public fun AIAgentContextBase.getAgentContextData(): AgentContextData? {
+public fun AIAgentContext.getAgentContextData(): AgentContextData? {
     return this.get(agentContextDataAdditionalKey)
 }
 
@@ -307,6 +307,6 @@ public fun AIAgentContextBase.getAgentContextData(): AgentContextData? {
  * @return `true` if the agent context data was successfully removed, or `false` if no data was found to remove.
  */
 @OptIn(InternalAgentsApi::class)
-public fun AIAgentContextBase.removeAgentContextData(): Boolean {
+public fun AIAgentContext.removeAgentContextData(): Boolean {
     return this.remove(agentContextDataAdditionalKey)
 }

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

@@ -1,6 +1,6 @@
 package ai.koog.agents.core.agent.entity
 
-import ai.koog.agents.core.agent.context.AIAgentContextBase
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.utils.Option
 
 /**
@@ -17,9 +17,9 @@ import ai.koog.agents.core.utils.Option
  */
 public class AIAgentEdge<IncomingOutput, OutgoingInput> internal constructor(
     public val toNode: AIAgentNodeBase<OutgoingInput, *>,
-    internal val forwardOutput: suspend (context: AIAgentContextBase, output: IncomingOutput) -> Option<OutgoingInput>,
+    internal val forwardOutput: suspend (context: AIAgentContext, output: IncomingOutput) -> Option<OutgoingInput>,
 ) {
     @Suppress("UNCHECKED_CAST")
-    internal suspend fun forwardOutputUnsafe(output: Any?, context: AIAgentContextBase): Option<OutgoingInput> =
+    internal suspend fun forwardOutputUnsafe(output: Any?, context: AIAgentContext): Option<OutgoingInput> =
         forwardOutput(context, output as IncomingOutput)
 }

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

@@ -27,7 +27,8 @@ public class AIAgentGraphStrategy<Input, Output>(
     nodeStart,
     nodeFinish,
     toolSelectionStrategy
-) {
+),
+    AIAgentStrategy {
     /**
      * Represents the metadata of the subgraph associated with the AI agent strategy.
      *

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

@@ -0,0 +1,26 @@
+package ai.koog.agents.core.agent.entity
+
+import ai.koog.agents.core.agent.context.AIAgentContext
+
+/**
+ * An interface representing the execution strategy of an AI agent.
+ *
+ * This strategy defines how the AI agent processes input to produce
+ * an output within the context of its operation. It serves as the primary
+ * mechanism to encapsulate various decision-making, learning, or processing
+ * approaches used by the agent in a flexible manner.
+ *
+ * @param Input The type of input data that the strategy will process.
+ * @param Output The type of output data that the strategy will generate.
+ * @param TContext The type of context in which the strategy is executed, extending [AIAgentContext].
+ */
+public interface AIAgentStrategy {
+    /**
+     * The name of the AI agent strategy.
+     *
+     * This property provides a human-readable identifier for the strategy,
+     * which can be used for logging, debugging, or distinguishing between
+     * multiple strategies within the system.
+     */
+    public val name: String
+}

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

@@ -2,7 +2,7 @@ package ai.koog.agents.core.agent.entity
 
 import ai.koog.agents.core.agent.AIAgentMaxNumberOfIterationsReachedException
 import ai.koog.agents.core.agent.AIAgentStuckInTheNodeException
-import ai.koog.agents.core.agent.context.AIAgentContextBase
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
 import ai.koog.agents.core.agent.context.DetachedPromptExecutorAPI
 import ai.koog.agents.core.agent.context.getAgentContextData
@@ -84,7 +84,7 @@ public open class AIAgentSubgraph<Input, Output>(
     )
 
     @OptIn(DetachedPromptExecutorAPI::class)
-    private suspend fun selectTools(context: AIAgentContextBase) = when (toolSelectionStrategy) {
+    private suspend fun selectTools(context: AIAgentContext) = when (toolSelectionStrategy) {
         is ToolSelectionStrategy.ALL -> context.llm.tools
         is ToolSelectionStrategy.NONE -> emptyList()
         is ToolSelectionStrategy.Tools -> toolSelectionStrategy.tools
@@ -222,7 +222,7 @@ public open class AIAgentSubgraph<Input, Output>(
         return result
     }
 
-    private fun formatLog(context: AIAgentContextBase, message: String): String =
+    private fun formatLog(context: AIAgentContext, message: String): String =
         "$message [$name, ${context.strategyName}, ${context.runId}]"
 }
 

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

@@ -1,6 +1,6 @@
 package ai.koog.agents.core.dsl.builder
 
-import ai.koog.agents.core.agent.context.AIAgentContextBase
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.entity.AIAgentEdge
 import ai.koog.agents.core.agent.entity.AIAgentNodeBase
 import ai.koog.agents.core.utils.Option
@@ -57,7 +57,7 @@ public class AIAgentEdgeBuilder<IncomingOutput, OutgoingInput, CompatibleOutput
 public class AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput, OutgoingInput> internal constructor(
     internal val fromNode: AIAgentNodeBase<*, IncomingOutput>,
     internal val toNode: AIAgentNodeBase<OutgoingInput, *>,
-    internal val forwardOutputComposition: suspend (AIAgentContextBase, IncomingOutput) -> Option<IntermediateOutput>
+    internal val forwardOutputComposition: suspend (AIAgentContext, IncomingOutput) -> Option<IntermediateOutput>
 ) {
     /**
      * Filters the intermediate outputs of the [ai.koog.agents.core.agent.entity.AIAgentNode] based on a specified condition.
@@ -69,7 +69,7 @@ public class AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput,
      */
     @EdgeTransformationDslMarker
     public infix fun onCondition(
-        block: suspend AIAgentContextBase.(output: IntermediateOutput) -> Boolean
+        block: suspend AIAgentContext.(output: IntermediateOutput) -> Boolean
     ): AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput, OutgoingInput> {
         return AIAgentEdgeBuilderIntermediate(
             fromNode = fromNode,
@@ -90,7 +90,7 @@ public class AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput,
      */
     @EdgeTransformationDslMarker
     public infix fun <NewIntermediateOutput> transformed(
-        block: suspend AIAgentContextBase.(IntermediateOutput) -> NewIntermediateOutput
+        block: suspend AIAgentContext.(IntermediateOutput) -> NewIntermediateOutput
     ): AIAgentEdgeBuilderIntermediate<IncomingOutput, NewIntermediateOutput, OutgoingInput> {
         return AIAgentEdgeBuilderIntermediate(
             fromNode = fromNode,

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

@@ -1,6 +1,6 @@
 package ai.koog.agents.core.dsl.builder
 
-import ai.koog.agents.core.agent.context.AIAgentContextBase
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
 import ai.koog.agents.core.agent.entity.AIAgentNode
 import ai.koog.agents.core.agent.entity.AIAgentNodeBase
@@ -21,7 +21,7 @@ import kotlin.reflect.KType
  * @param outputType [KType] of the [Output]
  * @constructor Used internally to create a new builder with the provided execution logic.
  * @param execute A suspending function to define the execution logic of the node. This function will be called
- * in the scope of [AIAgentContextBase], where it has access to the AI agent's context and tools relevant
+ * in the scope of [AIAgentContext], where it has access to the AI agent's context and tools relevant
  * to its operation.
  */
 public open class AIAgentNodeBuilder<Input, Output>(

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

@@ -1,7 +1,7 @@
 package ai.koog.agents.core.dsl.builder
 
 import ai.koog.agents.core.agent.config.AIAgentConfigBase
-import ai.koog.agents.core.agent.context.AIAgentContextBase
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
 import ai.koog.agents.core.agent.context.AIAgentLLMContext
 import ai.koog.agents.core.agent.entity.AIAgentStateManager
@@ -109,7 +109,7 @@ public class AIAgentParallelNodesMergeContext<Input, Output>(
      * @param context The new context to replace the current one in the underlying context base.
      * @return Unit
      */
-    override suspend fun replace(context: AIAgentContextBase): Unit = underlyingContextBase.replace(context)
+    override suspend fun replace(context: AIAgentContext): Unit = underlyingContextBase.replace(context)
 
     /**
      * Selects a result based on a predicate.