Introduced AIAgentStrategy - generic one

Author: Rizzen

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

@@ -6,10 +6,33 @@ import ai.koog.agents.core.agent.context.AIAgentLLMContext
 import ai.koog.agents.core.agent.entity.AIAgentStateManager
 import ai.koog.agents.core.agent.entity.AIAgentStorage
 import ai.koog.agents.core.agent.entity.AIAgentStorageKey
+import ai.koog.agents.core.annotation.InternalAgentsApi
 import ai.koog.agents.core.environment.AIAgentEnvironment
 import ai.koog.agents.core.feature.AIAgentFeature
+import ai.koog.agents.core.feature.AIAgentNonGraphPipeline
 import ai.koog.prompt.message.Message
 
+/**
+ * AIAgentLoopContext represents the execution context for an AI agent operating in a loop.
+ * It provides access to critical components such as the environment, configuration, large language model (LLM) context,
+ * state management, and storage. Additionally, it enables the agent to store, retrieve, and manage context-specific data
+ * during its execution lifecycle.
+ *
+ * @property environment The environment interface allowing the agent to interact with the external world,
+ * including executing tools and reporting problems.
+ * @property agentId A unique identifier for the agent, differentiating it from other agents in the system.
+ * @property runId A unique identifier for the current run or instance of the agent's operation.
+ * @property agentInput The input data passed to the agent, which can be of any type, depending on the agent's context.
+ * @property config The configuration settings for the agent, including its prompt and model details,
+ * as well as operational constraints like iteration limits.
+ * @property llm The context for interacting with the large language model used by the agent, enabling message history
+ * retrieval and processing.
+ * @property stateManager The state management component responsible for tracking and updating the agent's state during its execution.
+ * @property storage A storage interface providing persistent storage capabilities for the agent's data.
+ * @property strategyName The name of the agent's strategic approach or operational method, determining its behavior
+ * during execution.
+ */
+@Suppress("UNCHECKED_CAST")
 public class AIAgentLoopContext(
     override val environment: AIAgentEnvironment,
     override val agentId: String,
@@ -19,30 +42,32 @@ public class AIAgentLoopContext(
     override val llm: AIAgentLLMContext,
     override val stateManager: AIAgentStateManager,
     override val storage: AIAgentStorage,
-    override val strategyName: String
+    override val strategyName: String,
+    public val pipeline: AIAgentNonGraphPipeline
 ) : AIAgentContext {
+
+    private val storeMap: MutableMap<AIAgentStorageKey<*>, Any> = mutableMapOf()
+
     override fun store(key: AIAgentStorageKey<*>, value: Any) {
-        TODO("Not yet implemented")
+        storeMap[key] = value
     }
 
-    override fun <T> get(key: AIAgentStorageKey<*>): T? {
-        TODO("Not yet implemented")
-    }
+    override fun <T> get(key: AIAgentStorageKey<*>): T? = storeMap[key] as T?
 
-    override fun remove(key: AIAgentStorageKey<*>): Boolean {
-        TODO("Not yet implemented")
-    }
+    override fun remove(key: AIAgentStorageKey<*>): Boolean = storeMap.remove(key) != null
+
+    @OptIn(InternalAgentsApi::class)
+    private val features: Map<AIAgentStorageKey<*>, Any> =
+        pipeline.getAgentFeatures(this)
 
     override fun <Feature : Any> feature(key: AIAgentStorageKey<Feature>): Feature? {
-        TODO("Not yet implemented")
+        @Suppress("UNCHECKED_CAST")
+        return features[key] as Feature?
     }
 
-    override fun <Feature : Any> feature(feature: AIAgentFeature<*, Feature>): Feature? {
-        TODO("Not yet implemented")
-    }
+    override fun <Feature : Any> feature(feature: AIAgentFeature<*, Feature>): Feature? = feature(feature.key)
 
     override suspend fun getHistory(): List<Message> {
-        TODO("Not yet implemented")
+        return llm.readSession { prompt.messages }
     }
-
-}
+}

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

@@ -1,13 +1,18 @@
 package ai.koog.agents.core.agent
 
+import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
 import ai.koog.agents.core.agent.entity.AIAgentStrategy
 
 /**
  *
  * */
 public class AIAgentLoopStrategy<Input, Output>(
     override val name: String
-): AIAgentStrategy<Input, Output, AIAgentLoopContext> {
-
+) : AIAgentStrategy<Input, Output, AIAgentLoopContext> {
+    override suspend fun execute(
+        context: AIAgentGraphContextBase,
+        input: Input
+    ): Output? {
+        TODO("Not yet implemented")
+    }
 }
-

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

@@ -2,15 +2,13 @@
 
 package ai.koog.agents.core.agent.entity
 
-import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
 import ai.koog.agents.core.annotation.InternalAgentsApi
 import ai.koog.agents.core.utils.runCatchingCancellable
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.serializer
 
-
 /**
  * Represents a strategy for managing and executing AI agent workflows built as subgraphs of interconnected nodes.
  *

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

@@ -15,7 +15,7 @@ import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
  * @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<Input, Output, TContext: AIAgentContext> {
+public interface AIAgentStrategy<Input, Output, TContext : AIAgentContext> {
     /**
      * The name of the AI agent strategy.
      *

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

@@ -212,7 +212,7 @@ public abstract class AIAgentPipeline {
      * @param context The context of the strategy execution
      */
     @OptIn(InternalAgentsApi::class)
-    public suspend fun onStrategyStarted(strategy: AIAgentStrategy, context: AIAgentContext) {
+    public suspend fun onStrategyStarted(strategy: AIAgentStrategy<*, *, *>, context: AIAgentContext) {
         strategyHandlers.values.forEach { handler ->
             val eventContext =
                 StrategyStartContext(runId = context.runId, strategy = strategy, feature = handler.feature)
@@ -229,7 +229,7 @@ public abstract class AIAgentPipeline {
      */
     @OptIn(InternalAgentsApi::class)
     public suspend fun onStrategyFinished(
-        strategy: AIAgentStrategy,
+        strategy: AIAgentStrategy<*, *, *>,
         context: AIAgentContext,
         result: Any?,
         resultType: KType,

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

@@ -20,7 +20,7 @@ public interface StrategyEventHandlerContext : EventHandlerContext
  */
 public class StrategyStartContext<TFeature>(
     public val runId: String,
-    public val strategy: AIAgentStrategy,
+    public val strategy: AIAgentStrategy<*, *, *>,
     public val feature: TFeature
 ) : StrategyEventHandlerContext
 
@@ -36,7 +36,7 @@ public class StrategyStartContext<TFeature>(
  */
 public class StrategyFinishContext<TFeature>(
     public val runId: String,
-    public val strategy: AIAgentStrategy,
+    public val strategy: AIAgentStrategy<*, *, *>,
     public val feature: TFeature,
     public val result: Any?,
     public val resultType: KType,