Skip to content

Commit 42e78b2

Browse files
sdubovsdubov
committed
KG-178. Add execution info in AIAgentPipeline
- Add execution info parameter to pipeline event handlers; - Propagate execution info through agent initialization in FunctionalAIAgent and GraphAIAgent; - Add AIAgentContext.with() API for execution info agent execution hierarchy; - Update AIAgentGraphContext to make environment and executionInfo mutable properties to pass reference into an agent context; - Update agent event context classes to accept execution info; - Update tests to check execution info and execution path.
1 parent 2a8f025 commit 42e78b2

File tree

26 files changed

+1022
-389
lines changed

26 files changed

+1022
-389
lines changed

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

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -85,59 +85,60 @@ public class FunctionalAIAgent<Input, Output>(
8585
toolRegistry = toolRegistry,
8686
)
8787

88-
val initialAgentLLMContext = AIAgentLLMContext(
88+
val initialLLMContext = AIAgentLLMContext(
8989
tools = toolRegistry.tools.map { it.descriptor },
9090
toolRegistry = toolRegistry,
9191
prompt = agentConfig.prompt,
9292
model = agentConfig.model,
9393
responseProcessor = agentConfig.responseProcessor,
94-
promptExecutor = ContextualPromptExecutor(
95-
executor = promptExecutor,
96-
pipeline = pipeline,
97-
runId = runId
98-
),
94+
promptExecutor = promptExecutor,
9995
environment = environment,
10096
config = agentConfig,
10197
clock = clock
10298
)
10399

104-
val agentExecutionInfo = AgentExecutionInfo(parent = null, partName = id)
105-
val runExecutionInfo = AgentExecutionInfo(parent = agentExecutionInfo, partName = runId)
106-
100+
val executionInfo = AgentExecutionInfo(parent = null, partName = id)
107101
val preparedEnvironment = prepareEnvironment()
108102

109103
// Context
110-
val agentContext = AIAgentFunctionalContext(
104+
val initialAgentContext = AIAgentFunctionalContext(
111105
environment = preparedEnvironment,
112106
agentId = id,
113107
runId = runId,
114108
agentInput = agentInput,
115109
config = agentConfig,
116-
llm = initialAgentLLMContext,
110+
llm = initialLLMContext,
117111
stateManager = AIAgentStateManager(),
118112
storage = AIAgentStorage(),
119113
strategyName = strategy.name,
120114
pipeline = pipeline,
121-
executionInfo = runExecutionInfo,
115+
executionInfo = executionInfo,
122116
parentContext = null
123117
)
124118

125119
// Updated environment
126-
val environmentProxy = ContextualAgentEnvironment(
120+
val contextualEnvironment = ContextualAgentEnvironment(
127121
environment = preparedEnvironment,
128-
context = agentContext,
122+
context = initialAgentContext,
123+
)
124+
125+
val contextualPromptExecutor = ContextualPromptExecutor(
126+
executor = promptExecutor,
127+
context = initialAgentContext,
129128
)
130129

131-
val updatedLLMContext = agentContext.llm.copy(
132-
environment = environmentProxy
130+
val updatedLLMContext = initialAgentContext.llm.copy(
131+
environment = contextualEnvironment,
132+
promptExecutor = contextualPromptExecutor,
133133
)
134134

135-
// Update the environment and llm with a created context instance
136-
return agentContext.copy(
137-
parentRootContext = agentContext.parentContext,
135+
val updatedAgentContext = initialAgentContext.copy(
138136
llm = updatedLLMContext,
139-
environment = environmentProxy
137+
environment = contextualEnvironment,
138+
parentRootContext = initialAgentContext.parentContext, // Keep the original parent context
140139
)
140+
141+
return updatedAgentContext
141142
}
142143

143144
//region Private Methods

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

Lines changed: 37 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -99,57 +99,60 @@ public open class GraphAIAgent<Input, Output>(
9999
val stateManager = AIAgentStateManager()
100100
val storage = AIAgentStorage()
101101

102-
val agentExecutionInfo = AgentExecutionInfo(parent = null, partName = id)
103-
val runExecutionInfo = AgentExecutionInfo(parent = agentExecutionInfo, partName = runId)
102+
val executionInfo = AgentExecutionInfo(parent = null, partName = id)
103+
val initialEnvironment = prepareAgentEnvironment(executionInfo = executionInfo)
104104

105-
val preparedEnvironment = prepareAgentEnvironment()
105+
// Initial
106+
val initialLLMContext = AIAgentLLMContext(
107+
tools = toolRegistry.tools.map { it.descriptor },
108+
toolRegistry = toolRegistry,
109+
prompt = agentConfig.prompt,
110+
model = agentConfig.model,
111+
responseProcessor = agentConfig.responseProcessor,
112+
promptExecutor = promptExecutor,
113+
environment = initialEnvironment,
114+
config = agentConfig,
115+
clock = clock
116+
)
106117

107-
val context = AIAgentGraphContext(
108-
environment = preparedEnvironment,
118+
val agentContext = AIAgentGraphContext(
119+
environment = initialEnvironment,
109120
agentId = id,
110121
agentInput = agentInput,
111122
agentInputType = inputType,
112123
config = agentConfig,
113-
llm = AIAgentLLMContext(
114-
tools = toolRegistry.tools.map { it.descriptor },
115-
toolRegistry = toolRegistry,
116-
prompt = agentConfig.prompt,
117-
model = agentConfig.model,
118-
responseProcessor = agentConfig.responseProcessor,
119-
promptExecutor = ContextualPromptExecutor(
120-
executor = promptExecutor,
121-
pipeline = pipeline,
122-
runId = runId
123-
),
124-
environment = preparedEnvironment,
125-
config = agentConfig,
126-
clock = clock
127-
),
124+
llm = initialLLMContext,
128125
stateManager = stateManager,
129126
storage = storage,
130127
runId = runId,
131128
strategyName = strategy.name,
132129
pipeline = pipeline,
133-
executionInfo = runExecutionInfo,
130+
executionInfo = executionInfo,
134131
parentContext = null,
135132
)
136133

137-
// Update Environment
138134
val contextualEnvironment = ContextualAgentEnvironment(
139-
environment = preparedEnvironment,
140-
context = context,
135+
environment = initialEnvironment,
136+
context = agentContext,
141137
)
142138

143-
val updatedLLMContext = context.llm.copy(
144-
environment = contextualEnvironment
139+
val contextualPromptExecutor = ContextualPromptExecutor(
140+
executor = promptExecutor,
141+
context = agentContext,
145142
)
146143

147-
// Update the environment and llm with a created context instance
148-
return context.copy(
149-
parentContext = context.parentContext,
150-
llm = updatedLLMContext,
151-
environment = contextualEnvironment
144+
val updatedLLMContext = agentContext.llm.copy(
145+
environment = contextualEnvironment,
146+
promptExecutor = contextualPromptExecutor,
152147
)
148+
149+
agentContext.replace(agentContext.copy(
150+
executionInfo = executionInfo,
151+
llm = updatedLLMContext,
152+
environment = contextualEnvironment,
153+
))
154+
155+
return agentContext
153156
}
154157

155158
/**
@@ -162,7 +165,7 @@ public open class GraphAIAgent<Input, Output>(
162165
* @return An instance of `AIAgentEnvironment` that represents the finalized environment
163166
* for the AI agent after applying all transformations.
164167
*/
165-
private suspend fun prepareAgentEnvironment(): AIAgentEnvironment {
168+
private suspend fun prepareAgentEnvironment(executionInfo: AgentExecutionInfo): AIAgentEnvironment {
166169
// Create a base environment implementation
167170
val environment = GenericAgentEnvironment(
168171
agentId = id,
@@ -172,7 +175,8 @@ public open class GraphAIAgent<Input, Output>(
172175

173176
val preparedEnvironment = pipeline.onAgentEnvironmentTransforming(
174177
agent = this,
175-
baseEnvironment = environment
178+
baseEnvironment = environment,
179+
executionInfo = executionInfo,
176180
)
177181

178182
return preparedEnvironment

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

Lines changed: 43 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@ import ai.koog.agents.core.agent.AIAgent.Companion.State
44
import ai.koog.agents.core.agent.AIAgent.Companion.State.NotStarted
55
import ai.koog.agents.core.agent.context.AIAgentContext
66
import ai.koog.agents.core.agent.context.element.AgentRunInfoContextElement
7+
import ai.koog.agents.core.agent.context.with
78
import ai.koog.agents.core.agent.entity.AIAgentStrategy
9+
import ai.koog.agents.core.annotation.InternalAgentsApi
810
import ai.koog.agents.core.feature.AIAgentFeature
911
import ai.koog.agents.core.feature.pipeline.AIAgentPipeline
1012
import io.github.oshai.kotlinlogging.KLogger
@@ -96,60 +98,39 @@ public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgent
9698
strategyName = strategy.name
9799
)
98100
) {
99-
pipeline.withPreparedPipeline {
100-
agentStateMutex.withLock {
101-
state = State.Running(context)
102-
}
103-
104-
logger.debug {
105-
formatLog(
106-
agentId = id,
107-
runId = runId,
108-
message = "Starting agent execution"
109-
)
110-
}
111-
112-
pipeline.onAgentStarting<Input, Output>(
113-
runId = runId,
114-
agent = this@StatefulSingleUseAIAgent,
115-
context = context
116-
)
101+
context.withPreparedPipeline {
102+
context.with(partName = runId) { executionInfo ->
103+
agentStateMutex.withLock {
104+
@OptIn(InternalAgentsApi::class)
105+
state = State.Running(context.parentContext ?: context)
106+
}
117107

118-
val result = try {
119-
strategy.execute(context = context, input = agentInput)
120-
} catch (e: Throwable) {
121-
logger.error(e) { "Execution exception reported by server!" }
122-
pipeline.onAgentExecutionFailed(
123-
agentId = id,
124-
runId = runId,
125-
throwable = e
126-
)
127-
agentStateMutex.withLock { state = State.Failed(e) }
128-
throw e
129-
}
108+
logger.debug { formatLog(id, runId, "Starting agent execution") }
109+
pipeline.onAgentStarting<Input, Output>(executionInfo, runId, this@StatefulSingleUseAIAgent, context)
110+
111+
val result = try {
112+
@Suppress("UNCHECKED_CAST")
113+
strategy.execute(context = context, input = agentInput)
114+
} catch (e: Throwable) {
115+
logger.error(e) { "Execution exception reported by server!" }
116+
pipeline.onAgentExecutionFailed(executionInfo, id, runId, e)
117+
agentStateMutex.withLock { state = State.Failed(e) }
118+
throw e
119+
}
130120

131-
logger.debug {
132-
formatLog(
133-
agentId = id,
134-
runId = runId,
135-
message = "Finished agent execution"
136-
)
137-
}
138-
pipeline.onAgentCompleted(
139-
agentId = id,
140-
runId = runId,
141-
result = result
142-
)
121+
logger.debug { formatLog(id, runId, "Finished agent execution") }
122+
pipeline.onAgentCompleted(executionInfo, id, runId, result)
143123

144-
agentStateMutex.withLock {
145-
state = if (result != null) {
146-
State.Finished(result)
147-
} else {
148-
State.Failed(Exception("result is null"))
124+
agentStateMutex.withLock {
125+
state = if (result != null) {
126+
State.Finished(result)
127+
} else {
128+
State.Failed(Exception("result is null"))
129+
}
149130
}
150-
}
151131

152-
result ?: error("result is null")
132+
result ?: error("result is null")
133+
}
153134
}
154135
}
155136
}
@@ -208,14 +189,22 @@ public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgent
208189
*
209190
* @return The result of the block's execution.
210191
*/
211-
private suspend fun <T> AIAgentPipeline.withPreparedPipeline(block: suspend () -> T): T =
212-
try {
213-
prepareAllFeatures()
192+
private suspend fun <T> AIAgentContext.withPreparedPipeline(block: suspend () -> T): T {
193+
require(executionInfo.parent == null) {
194+
"withPreparedPipeline() should be called from a top level agent context."
195+
}
196+
197+
return try {
198+
pipeline.prepareAllFeatures()
214199
block.invoke()
215200
} finally {
216-
onAgentClosing(agentId = id)
217-
closeAllFeaturesMessageProcessors()
201+
pipeline.onAgentClosing(
202+
executionInfo = executionInfo.parent ?: executionInfo,
203+
agentId = id
204+
)
205+
pipeline.closeAllFeaturesMessageProcessors()
218206
}
207+
}
219208
}
220209

221210
/**

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

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ public interface AIAgentContext {
7878
* This variable provides synchronized access to the agent's state to ensure thread safety
7979
* and consistent state transitions during concurrent operations. It acts as a central
8080
* mechanism for managing state updates and validations across different
81-
* nodes and subgraphes of the AI agent's execution flow.
81+
* nodes and subgraphs of the AI agent's execution flow.
8282
*
8383
* The [stateManager] is utilized extensively in coordinating state changes, such as
8484
* tracking the number of iterations made by the agent and enforcing execution limits
@@ -192,3 +192,35 @@ public inline fun <reified TFeature : Any> AIAgentContext.feature(feature: AIAge
192192
*/
193193
public inline fun <reified TFeature : Any> AIAgentContext.featureOrThrow(feature: AIAgentFeature<*, TFeature>): TFeature =
194194
feature(feature) ?: throw NoSuchElementException("Feature ${feature.key} is not found.")
195+
196+
/**
197+
* Executes a block of code with a modified execution context.
198+
*
199+
* @param T The return type of the block being executed.
200+
* @param executionInfo The execution info to be set for the context.
201+
* @param block The suspend function to execute with the modified execution context.
202+
* @return The result of executing the provided block.
203+
*/
204+
public inline fun <T> AIAgentContext.with(executionInfo: AgentExecutionInfo, block: (executionInfo: AgentExecutionInfo) -> T): T {
205+
val originalExecutionInfo = this.executionInfo
206+
return try {
207+
this.executionInfo = executionInfo
208+
block(executionInfo)
209+
} finally {
210+
this.executionInfo = originalExecutionInfo
211+
}
212+
}
213+
214+
/**
215+
* Executes a block of code with a modified execution context, creating a parent-child relationship
216+
* between execution contexts for tracing purposes.
217+
*
218+
* @param T The return type of the block being executed.
219+
* @param partName The name of the execution part to append to the execution path.
220+
* @param block The suspend function to execute with the modified execution context.
221+
* @return The result of executing the provided block.
222+
*/
223+
public inline fun <T> AIAgentContext.with(partName: String, block: (executionInfo: AgentExecutionInfo) -> T): T {
224+
val executionInfo = AgentExecutionInfo(parent = this.executionInfo, partName = partName)
225+
return with(executionInfo = executionInfo, block = block)
226+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ import ai.koog.prompt.message.Message
3131
* during execution.
3232
*/
3333
@OptIn(InternalAgentsApi::class)
34-
@Suppress("UNCHECKED_CAST")
3534
public class AIAgentFunctionalContext(
3635
override val environment: AIAgentEnvironment,
3736
override val agentId: String,
@@ -53,6 +52,7 @@ public class AIAgentFunctionalContext(
5352
storeMap[key] = value
5453
}
5554

55+
@Suppress("UNCHECKED_CAST")
5656
override fun <T> get(key: AIAgentStorageKey<*>): T? = storeMap[key] as T?
5757

5858
override fun remove(key: AIAgentStorageKey<*>): Boolean = storeMap.remove(key) != null

0 commit comments

Comments
 (0)