Skip to content

Commit 00e446a

Browse files
sdubovsdubov
committed
KG-178. Introduce hierarchy into agent execution flow
1 parent 8eadf65 commit 00e446a

126 files changed

Lines changed: 3989 additions & 2616 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
package ai.koog.agents.core.agent
22

33
import ai.koog.agents.core.agent.context.AIAgentFunctionalContext
4+
import ai.koog.agents.core.agent.context.withParent
45
import ai.koog.agents.core.agent.entity.AIAgentStrategy
6+
import io.github.oshai.kotlinlogging.KotlinLogging
7+
import kotlin.reflect.typeOf
8+
import kotlin.uuid.ExperimentalUuidApi
59

610
/**
711
* A strategy for implementing AI agent behavior that operates in a loop-based manner.
812
*
913
* The [AIAgentFunctionalStrategy] class allows for the definition of a custom looping logic
10-
* that processes input and produces output by utilizing an [ai.koog.agents.core.agent.context.AIAgentFunctionalContext]. This strategy
14+
* that processes input and produces output by using an [ai.koog.agents.core.agent.context.AIAgentFunctionalContext]. This strategy
1115
* can be used to define iterative decision-making or execution processes for AI agents.
1216
*
1317
* @param TInput The type of input data processed by the strategy.
@@ -20,10 +24,24 @@ public class AIAgentFunctionalStrategy<TInput, TOutput>(
2024
override val name: String,
2125
public val func: suspend AIAgentFunctionalContext.(TInput) -> TOutput
2226
) : AIAgentStrategy<TInput, TOutput, AIAgentFunctionalContext> {
27+
28+
private companion object {
29+
private val logger = KotlinLogging.logger { }
30+
}
31+
32+
@OptIn(ExperimentalUuidApi::class)
2333
override suspend fun execute(
2434
context: AIAgentFunctionalContext,
2535
input: TInput
26-
): TOutput = context.func(input)
36+
): TOutput = withParent(context = context, partName = name) { executionInfo ->
37+
context.pipeline.onStrategyStarting(executionInfo, this@AIAgentFunctionalStrategy, context)
38+
val result = context.func(input)
39+
context.pipeline.onStrategyCompleted(executionInfo, this@AIAgentFunctionalStrategy, context, result, typeOf<Any?>())
40+
41+
logger.debug { "Finished executing strategy (name: $name) with result: $result" }
42+
result
43+
}
44+
2745
}
2846

2947
/**

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package ai.koog.agents.core.agent
22

33
import ai.koog.agents.core.agent.GraphAIAgent.FeatureContext
44
import ai.koog.agents.core.agent.config.AIAgentConfig
5+
import ai.koog.agents.core.agent.context.withParent
56
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
67
import ai.koog.agents.core.annotation.InternalAgentsApi
78
import ai.koog.agents.core.tools.Tool

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

Lines changed: 63 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,14 @@ package ai.koog.agents.core.agent
33
import ai.koog.agents.core.agent.config.AIAgentConfig
44
import ai.koog.agents.core.agent.context.AIAgentFunctionalContext
55
import ai.koog.agents.core.agent.context.AIAgentLLMContext
6+
import ai.koog.agents.core.agent.context.AgentExecutionInfo
7+
import ai.koog.agents.core.agent.context.AgentExecutionPath
68
import ai.koog.agents.core.agent.entity.AIAgentStateManager
79
import ai.koog.agents.core.agent.entity.AIAgentStorage
810
import ai.koog.agents.core.annotation.InternalAgentsApi
11+
import ai.koog.agents.core.environment.AIAgentEnvironment
912
import ai.koog.agents.core.environment.GenericAgentEnvironment
13+
import ai.koog.agents.core.environment.GenericAgentEnvironmentProxy
1014
import ai.koog.agents.core.feature.AIAgentFeature
1115
import ai.koog.agents.core.feature.AIAgentFunctionalFeature
1216
import ai.koog.agents.core.feature.PromptExecutorProxy
@@ -52,14 +56,6 @@ public class FunctionalAIAgent<Input, Output>(
5256

5357
override val pipeline: AIAgentFunctionalPipeline = AIAgentFunctionalPipeline(clock)
5458

55-
private val environment = GenericAgentEnvironment(
56-
agentId = this.id,
57-
strategyId = strategy.name,
58-
logger = logger,
59-
toolRegistry = toolRegistry,
60-
pipeline = pipeline
61-
)
62-
6359
/**
6460
* Represents a context for managing and configuring features in an AI agent.
6561
* Provides functionality to install and configure features into a specific instance of an AI agent.
@@ -84,32 +80,81 @@ public class FunctionalAIAgent<Input, Output>(
8480
}
8581

8682
override suspend fun prepareContext(agentInput: Input, runId: String): AIAgentFunctionalContext {
87-
val llm = AIAgentLLMContext(
83+
84+
val environment = GenericAgentEnvironment(
85+
agentId = id,
86+
logger = logger,
87+
toolRegistry = toolRegistry,
88+
)
89+
90+
val initialAgentLLMContext = AIAgentLLMContext(
8891
tools = toolRegistry.tools.map { it.descriptor },
8992
toolRegistry = toolRegistry,
9093
prompt = agentConfig.prompt,
9194
model = agentConfig.model,
92-
promptExecutor = PromptExecutorProxy(
93-
executor = promptExecutor,
94-
pipeline = pipeline,
95-
runId = runId
96-
),
95+
promptExecutor = promptExecutor,
9796
environment = environment,
9897
config = agentConfig,
9998
clock = clock
10099
)
101100

102-
return AIAgentFunctionalContext(
103-
environment = environment,
101+
102+
val executionPath = AgentExecutionPath(id)
103+
val executionInfo = AgentExecutionInfo(id = id, parent = null, path = executionPath)
104+
val preparedEnvironment = prepareEnvironment()
105+
106+
// Context
107+
val agentContext = AIAgentFunctionalContext(
108+
environment = preparedEnvironment,
104109
agentId = id,
105110
runId = runId,
106111
agentInput = agentInput,
107112
config = agentConfig,
108-
llm = llm,
113+
llm = initialAgentLLMContext,
109114
stateManager = AIAgentStateManager(),
110115
storage = AIAgentStorage(),
111116
strategyName = strategy.name,
112-
pipeline = pipeline
117+
pipeline = pipeline,
118+
executionInfo = executionInfo,
119+
)
120+
121+
// Updated environment
122+
val environmentProxy = GenericAgentEnvironmentProxy(
123+
environment = preparedEnvironment,
124+
context = agentContext,
125+
)
126+
127+
// Updated prompt executor
128+
val promptExecutorProxy = PromptExecutorProxy(
129+
executor = promptExecutor,
130+
pipeline = pipeline,
131+
runId = runId,
132+
context = agentContext
133+
)
134+
135+
val updatedLLMContext = agentContext.llm.copy(
136+
promptExecutor = promptExecutorProxy,
137+
environment = environmentProxy
138+
)
139+
140+
// Update the environment and llm with a created context instance
141+
return agentContext.copy(
142+
llm = updatedLLMContext,
143+
environment = environmentProxy
113144
)
114145
}
146+
147+
//region Private Methods
148+
149+
private fun prepareEnvironment(): AIAgentEnvironment {
150+
val baseEnvironment = GenericAgentEnvironment(
151+
agentId = id,
152+
logger = logger,
153+
toolRegistry = toolRegistry,
154+
)
155+
156+
return baseEnvironment
157+
}
158+
159+
//endregion Private Methods
115160
}

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

Lines changed: 69 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,15 @@ import ai.koog.agents.core.agent.config.AIAgentConfig
44
import ai.koog.agents.core.agent.context.AIAgentGraphContext
55
import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
66
import ai.koog.agents.core.agent.context.AIAgentLLMContext
7+
import ai.koog.agents.core.agent.context.AgentExecutionInfo
8+
import ai.koog.agents.core.agent.context.AgentExecutionPath
79
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
810
import ai.koog.agents.core.agent.entity.AIAgentStateManager
911
import ai.koog.agents.core.agent.entity.AIAgentStorage
1012
import ai.koog.agents.core.annotation.InternalAgentsApi
13+
import ai.koog.agents.core.environment.AIAgentEnvironment
1114
import ai.koog.agents.core.environment.GenericAgentEnvironment
15+
import ai.koog.agents.core.environment.GenericAgentEnvironmentProxy
1216
import ai.koog.agents.core.feature.AIAgentFeature
1317
import ai.koog.agents.core.feature.AIAgentGraphFeature
1418
import ai.koog.agents.core.feature.PromptExecutorProxy
@@ -66,14 +70,6 @@ public open class GraphAIAgent<Input, Output>(
6670

6771
override val pipeline: AIAgentGraphPipeline = AIAgentGraphPipeline(clock)
6872

69-
private val environment = GenericAgentEnvironment(
70-
agentId = this.id,
71-
strategyId = strategy.name,
72-
logger = logger,
73-
toolRegistry = toolRegistry,
74-
pipeline = pipeline
75-
)
76-
7773
/**
7874
* The context for adding and configuring features in a Kotlin AI Agent instance.
7975
*
@@ -104,40 +100,82 @@ public open class GraphAIAgent<Input, Output>(
104100
val stateManager = AIAgentStateManager()
105101
val storage = AIAgentStorage()
106102

107-
// Environment (initially equal to the current agent), transformed by some features
108-
// (ex: testing feature transforms it into a MockEnvironment with mocked tools)
109-
val preparedEnvironment =
110-
pipeline.onAgentEnvironmentTransforming(
111-
strategy = strategy,
112-
agent = this,
113-
baseEnvironment = environment
114-
)
103+
val executionPath = AgentExecutionPath(id)
104+
val executionInfo = AgentExecutionInfo(id = id, parent = null, path = executionPath)
105+
val preparedEnvironment = prepareEnvironment(executionInfo)
106+
107+
val initialAgentLLMContext = AIAgentLLMContext(
108+
tools = toolRegistry.tools.map { it.descriptor },
109+
toolRegistry = toolRegistry,
110+
prompt = agentConfig.prompt,
111+
model = agentConfig.model,
112+
promptExecutor = promptExecutor,
113+
environment = preparedEnvironment,
114+
config = agentConfig,
115+
clock = clock
116+
)
115117

116-
return AIAgentGraphContext(
118+
// Context
119+
val agentContext = AIAgentGraphContext(
117120
environment = preparedEnvironment,
118121
agentId = id,
119122
agentInput = agentInput,
120123
agentInputType = inputType,
121124
config = agentConfig,
122-
llm = AIAgentLLMContext(
123-
tools = toolRegistry.tools.map { it.descriptor },
124-
toolRegistry = toolRegistry,
125-
prompt = agentConfig.prompt,
126-
model = agentConfig.model,
127-
promptExecutor = PromptExecutorProxy(
128-
executor = promptExecutor,
129-
pipeline = pipeline,
130-
runId = runId
131-
),
132-
environment = preparedEnvironment,
133-
config = agentConfig,
134-
clock = clock
135-
),
125+
llm = initialAgentLLMContext,
136126
stateManager = stateManager,
137127
storage = storage,
138128
runId = runId,
139129
strategyName = strategy.name,
140130
pipeline = pipeline,
131+
executionInfo = executionInfo,
132+
)
133+
134+
// Updated environment
135+
val environmentProxy = GenericAgentEnvironmentProxy(
136+
environment = preparedEnvironment,
137+
context = agentContext,
138+
)
139+
140+
// Updated prompt executor
141+
val promptExecutorProxy = PromptExecutorProxy(
142+
executor = promptExecutor,
143+
pipeline = pipeline,
144+
runId = runId,
145+
context = agentContext
146+
)
147+
148+
val updatedLLMContext = agentContext.llm.copy(
149+
promptExecutor = promptExecutorProxy,
150+
environment = environmentProxy
151+
)
152+
153+
// Update the environment and llm with a created context instance
154+
return agentContext.copy(
155+
llm = updatedLLMContext,
156+
environment = environmentProxy
157+
)
158+
}
159+
160+
// Environment (initially equal to the current agent), transformed by some features
161+
// (ex: testing feature transforms it into a MockEnvironment with mocked tools)
162+
private suspend fun prepareEnvironment(executionInfo: AgentExecutionInfo): AIAgentEnvironment {
163+
val baseEnvironment = GenericAgentEnvironment(
164+
agentId = id,
165+
logger = logger,
166+
toolRegistry = toolRegistry,
141167
)
168+
169+
// Environment (initially equal to the current agent), transformed by some features
170+
// (ex: testing feature transforms it into a MockEnvironment with mocked tools)
171+
val preparedEnvironment =
172+
pipeline.onAgentEnvironmentTransforming(
173+
executionInfo = executionInfo,
174+
strategy = strategy,
175+
agent = this,
176+
baseEnvironment = baseEnvironment
177+
)
178+
179+
return preparedEnvironment
142180
}
143181
}

0 commit comments

Comments
 (0)