[WIP] KG-178. Introduce hierarchy for agent execution flow

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

@@ -1,13 +1,17 @@
 package ai.koog.agents.core.agent
 
 import ai.koog.agents.core.agent.context.AIAgentFunctionalContext
+import ai.koog.agents.core.agent.context.withParent
 import ai.koog.agents.core.agent.entity.AIAgentStrategy
+import io.github.oshai.kotlinlogging.KotlinLogging
+import kotlin.reflect.typeOf
+import kotlin.uuid.ExperimentalUuidApi
 
 /**
  * A strategy for implementing AI agent behavior that operates in a loop-based manner.
  *
  * The [AIAgentFunctionalStrategy] class allows for the definition of a custom looping logic
- * that processes input and produces output by utilizing an [ai.koog.agents.core.agent.context.AIAgentFunctionalContext]. This strategy
+ * that processes input and produces output by using an [ai.koog.agents.core.agent.context.AIAgentFunctionalContext]. This strategy
  * can be used to define iterative decision-making or execution processes for AI agents.
  *
  * @param TInput The type of input data processed by the strategy.
@@ -20,10 +24,24 @@ public class AIAgentFunctionalStrategy<TInput, TOutput>(
     override val name: String,
     public val func: suspend AIAgentFunctionalContext.(TInput) -> TOutput
 ) : AIAgentStrategy<TInput, TOutput, AIAgentFunctionalContext> {
+
+    private companion object {
+        private val logger = KotlinLogging.logger { }
+    }
+
+    @OptIn(ExperimentalUuidApi::class)
     override suspend fun execute(
         context: AIAgentFunctionalContext,
         input: TInput
-    ): TOutput = context.func(input)
+    ): TOutput = withParent(context = context, partName = name) { executionInfo ->
+        context.pipeline.onStrategyStarting(executionInfo, this@AIAgentFunctionalStrategy, context)
+        val result = context.func(input)
+        context.pipeline.onStrategyCompleted(executionInfo, this@AIAgentFunctionalStrategy, context, result, typeOf<Any?>())
+
+        logger.debug { "Finished executing strategy (name: $name) with result: $result" }
+        result
+    }
+
 }
 
 /**

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

@@ -2,6 +2,7 @@ package ai.koog.agents.core.agent
 
 import ai.koog.agents.core.agent.GraphAIAgent.FeatureContext
 import ai.koog.agents.core.agent.config.AIAgentConfig
+import ai.koog.agents.core.agent.context.withParent
 import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
 import ai.koog.agents.core.annotation.InternalAgentsApi
 import ai.koog.agents.core.tools.Tool

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

@@ -3,10 +3,14 @@ package ai.koog.agents.core.agent
 import ai.koog.agents.core.agent.config.AIAgentConfig
 import ai.koog.agents.core.agent.context.AIAgentFunctionalContext
 import ai.koog.agents.core.agent.context.AIAgentLLMContext
+import ai.koog.agents.core.agent.context.AgentExecutionInfo
+import ai.koog.agents.core.agent.context.AgentExecutionPath
 import ai.koog.agents.core.agent.entity.AIAgentStateManager
 import ai.koog.agents.core.agent.entity.AIAgentStorage
 import ai.koog.agents.core.annotation.InternalAgentsApi
+import ai.koog.agents.core.environment.AIAgentEnvironment
 import ai.koog.agents.core.environment.GenericAgentEnvironment
+import ai.koog.agents.core.environment.GenericAgentEnvironmentProxy
 import ai.koog.agents.core.feature.AIAgentFeature
 import ai.koog.agents.core.feature.AIAgentFunctionalFeature
 import ai.koog.agents.core.feature.PromptExecutorProxy
@@ -52,14 +56,6 @@ public class FunctionalAIAgent<Input, Output>(
 
     override val pipeline: AIAgentFunctionalPipeline = AIAgentFunctionalPipeline(clock)
 
-    private val environment = GenericAgentEnvironment(
-        agentId = this.id,
-        strategyId = strategy.name,
-        logger = logger,
-        toolRegistry = toolRegistry,
-        pipeline = pipeline
-    )
-
     /**
      * Represents a context for managing and configuring features in an AI agent.
      * Provides functionality to install and configure features into a specific instance of an AI agent.
@@ -84,32 +80,81 @@ public class FunctionalAIAgent<Input, Output>(
     }
 
     override suspend fun prepareContext(agentInput: Input, runId: String): AIAgentFunctionalContext {
-        val llm = AIAgentLLMContext(
+
+        val environment = GenericAgentEnvironment(
+            agentId = id,
+            logger = logger,
+            toolRegistry = toolRegistry,
+        )
+
+        val initialAgentLLMContext = AIAgentLLMContext(
             tools = toolRegistry.tools.map { it.descriptor },
             toolRegistry = toolRegistry,
             prompt = agentConfig.prompt,
             model = agentConfig.model,
-            promptExecutor = PromptExecutorProxy(
-                executor = promptExecutor,
-                pipeline = pipeline,
-                runId = runId
-            ),
+            promptExecutor = promptExecutor,
             environment = environment,
             config = agentConfig,
             clock = clock
         )
 
-        return AIAgentFunctionalContext(
-            environment = environment,
+
+        val executionPath = AgentExecutionPath(id)
+        val executionInfo = AgentExecutionInfo(id = id, parent = null, path = executionPath)
+        val preparedEnvironment = prepareEnvironment()
+
+        // Context
+        val agentContext = AIAgentFunctionalContext(
+            environment = preparedEnvironment,
             agentId = id,
             runId = runId,
             agentInput = agentInput,
             config = agentConfig,
-            llm = llm,
+            llm = initialAgentLLMContext,
             stateManager = AIAgentStateManager(),
             storage = AIAgentStorage(),
             strategyName = strategy.name,
-            pipeline = pipeline
+            pipeline = pipeline,
+            executionInfo = executionInfo,
+        )
+
+        // Updated environment
+        val environmentProxy = GenericAgentEnvironmentProxy(
+            environment = preparedEnvironment,
+            context = agentContext,
+        )
+
+        // Updated prompt executor
+        val promptExecutorProxy = PromptExecutorProxy(
+            executor = promptExecutor,
+            pipeline = pipeline,
+            runId = runId,
+            context = agentContext
+        )
+
+        val updatedLLMContext = agentContext.llm.copy(
+            promptExecutor = promptExecutorProxy,
+            environment = environmentProxy
+        )
+
+        // Update the environment and llm with a created context instance
+        return agentContext.copy(
+            llm = updatedLLMContext,
+            environment = environmentProxy
         )
     }
+
+    //region Private Methods
+
+    private fun prepareEnvironment(): AIAgentEnvironment {
+        val baseEnvironment = GenericAgentEnvironment(
+            agentId = id,
+            logger = logger,
+            toolRegistry = toolRegistry,
+        )
+
+        return baseEnvironment
+    }
+
+    //endregion Private Methods
 }

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

@@ -4,11 +4,15 @@ import ai.koog.agents.core.agent.config.AIAgentConfig
 import ai.koog.agents.core.agent.context.AIAgentGraphContext
 import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
 import ai.koog.agents.core.agent.context.AIAgentLLMContext
+import ai.koog.agents.core.agent.context.AgentExecutionInfo
+import ai.koog.agents.core.agent.context.AgentExecutionPath
 import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
 import ai.koog.agents.core.agent.entity.AIAgentStateManager
 import ai.koog.agents.core.agent.entity.AIAgentStorage
 import ai.koog.agents.core.annotation.InternalAgentsApi
+import ai.koog.agents.core.environment.AIAgentEnvironment
 import ai.koog.agents.core.environment.GenericAgentEnvironment
+import ai.koog.agents.core.environment.GenericAgentEnvironmentProxy
 import ai.koog.agents.core.feature.AIAgentFeature
 import ai.koog.agents.core.feature.AIAgentGraphFeature
 import ai.koog.agents.core.feature.PromptExecutorProxy
@@ -66,14 +70,6 @@ public open class GraphAIAgent<Input, Output>(
 
     override val pipeline: AIAgentGraphPipeline = AIAgentGraphPipeline(clock)
 
-    private val environment = GenericAgentEnvironment(
-        agentId = this.id,
-        strategyId = strategy.name,
-        logger = logger,
-        toolRegistry = toolRegistry,
-        pipeline = pipeline
-    )
-
     /**
      * The context for adding and configuring features in a Kotlin AI Agent instance.
      *
@@ -104,40 +100,82 @@ public open class GraphAIAgent<Input, Output>(
         val stateManager = AIAgentStateManager()
         val storage = AIAgentStorage()
 
-        // Environment (initially equal to the current agent), transformed by some features
-        //   (ex: testing feature transforms it into a MockEnvironment with mocked tools)
-        val preparedEnvironment =
-            pipeline.onAgentEnvironmentTransforming(
-                strategy = strategy,
-                agent = this,
-                baseEnvironment = environment
-            )
+        val executionPath = AgentExecutionPath(id)
+        val executionInfo = AgentExecutionInfo(id = id, parent = null, path = executionPath)
+        val preparedEnvironment = prepareEnvironment(executionInfo)
+
+        val initialAgentLLMContext = AIAgentLLMContext(
+            tools = toolRegistry.tools.map { it.descriptor },
+            toolRegistry = toolRegistry,
+            prompt = agentConfig.prompt,
+            model = agentConfig.model,
+            promptExecutor = promptExecutor,
+            environment = preparedEnvironment,
+            config = agentConfig,
+            clock = clock
+        )
 
-        return AIAgentGraphContext(
+        // Context
+        val agentContext = AIAgentGraphContext(
             environment = preparedEnvironment,
             agentId = id,
             agentInput = agentInput,
             agentInputType = inputType,
             config = agentConfig,
-            llm = AIAgentLLMContext(
-                tools = toolRegistry.tools.map { it.descriptor },
-                toolRegistry = toolRegistry,
-                prompt = agentConfig.prompt,
-                model = agentConfig.model,
-                promptExecutor = PromptExecutorProxy(
-                    executor = promptExecutor,
-                    pipeline = pipeline,
-                    runId = runId
-                ),
-                environment = preparedEnvironment,
-                config = agentConfig,
-                clock = clock
-            ),
+            llm = initialAgentLLMContext,
             stateManager = stateManager,
             storage = storage,
             runId = runId,
             strategyName = strategy.name,
             pipeline = pipeline,
+            executionInfo = executionInfo,
+        )
+
+        // Updated environment
+        val environmentProxy = GenericAgentEnvironmentProxy(
+            environment = preparedEnvironment,
+            context = agentContext,
+        )
+
+        // Updated prompt executor
+        val promptExecutorProxy = PromptExecutorProxy(
+            executor = promptExecutor,
+            pipeline = pipeline,
+            runId = runId,
+            context = agentContext
+        )
+
+        val updatedLLMContext = agentContext.llm.copy(
+            promptExecutor = promptExecutorProxy,
+            environment = environmentProxy
+        )
+
+        // Update the environment and llm with a created context instance
+        return agentContext.copy(
+            llm = updatedLLMContext,
+            environment = environmentProxy
+        )
+    }
+
+    // Environment (initially equal to the current agent), transformed by some features
+    //   (ex: testing feature transforms it into a MockEnvironment with mocked tools)
+    private suspend fun prepareEnvironment(executionInfo: AgentExecutionInfo): AIAgentEnvironment {
+        val baseEnvironment = GenericAgentEnvironment(
+            agentId = id,
+            logger = logger,
+            toolRegistry = toolRegistry,
         )
+
+        // Environment (initially equal to the current agent), transformed by some features
+        //   (ex: testing feature transforms it into a MockEnvironment with mocked tools)
+        val preparedEnvironment =
+            pipeline.onAgentEnvironmentTransforming(
+                executionInfo = executionInfo,
+                strategy = strategy,
+                agent = this,
+                baseEnvironment = baseEnvironment
+            )
+
+        return preparedEnvironment
     }
 }