Skip to content

Commit 7f5f406

Browse files
sdubovsdubov
committed
KG-178. Drop context elements
1 parent f9a66bf commit 7f5f406

File tree

7 files changed

+87
-400
lines changed

7 files changed

+87
-400
lines changed

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

Lines changed: 26 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,13 @@ package ai.koog.agents.core.agent
33
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
6-
import ai.koog.agents.core.agent.context.element.AgentRunInfoContextElement
76
import ai.koog.agents.core.agent.entity.AIAgentStrategy
87
import ai.koog.agents.core.annotation.InternalAgentsApi
98
import ai.koog.agents.core.feature.AIAgentFeature
109
import ai.koog.agents.core.feature.pipeline.AIAgentPipeline
1110
import io.github.oshai.kotlinlogging.KLogger
1211
import kotlinx.coroutines.sync.Mutex
1312
import kotlinx.coroutines.sync.withLock
14-
import kotlinx.coroutines.withContext
1513
import kotlin.reflect.KClass
1614
import kotlin.uuid.ExperimentalUuidApi
1715
import kotlin.uuid.Uuid
@@ -90,46 +88,37 @@ public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgent
9088
val eventId = Uuid.random().toString()
9189
val context = prepareContext(agentInput, runId, eventId)
9290

93-
return withContext(
94-
AgentRunInfoContextElement(
95-
agentId = id,
96-
runId = runId,
97-
agentConfig = agentConfig,
98-
strategyName = strategy.name
99-
)
100-
) {
101-
context.withPreparedPipeline {
102-
agentStateMutex.withLock {
103-
@OptIn(InternalAgentsApi::class)
104-
state = State.Running(context.parentContext ?: context)
105-
}
91+
return context.withPreparedPipeline {
92+
agentStateMutex.withLock {
93+
@OptIn(InternalAgentsApi::class)
94+
state = State.Running(context.parentContext ?: context)
95+
}
10696

107-
logger.debug { formatLog(id, runId, "Starting agent execution") }
108-
pipeline.onAgentStarting<Input, Output>(eventId, context.executionInfo, runId, this@StatefulSingleUseAIAgent, context)
109-
110-
val result = try {
111-
@Suppress("UNCHECKED_CAST")
112-
strategy.execute(context = context, input = agentInput)
113-
} catch (e: Throwable) {
114-
logger.error(e) { "Execution exception reported by server!" }
115-
pipeline.onAgentExecutionFailed(eventId, context.executionInfo, id, runId, e)
116-
agentStateMutex.withLock { state = State.Failed(e) }
117-
throw e
118-
}
97+
logger.debug { formatLog(id, runId, "Starting agent execution") }
98+
pipeline.onAgentStarting<Input, Output>(eventId, context.executionInfo, runId, this@StatefulSingleUseAIAgent, context)
99+
100+
val result = try {
101+
@Suppress("UNCHECKED_CAST")
102+
strategy.execute(context = context, input = agentInput)
103+
} catch (e: Throwable) {
104+
logger.error(e) { "Execution exception reported by server!" }
105+
pipeline.onAgentExecutionFailed(eventId, context.executionInfo, id, runId, e)
106+
agentStateMutex.withLock { state = State.Failed(e) }
107+
throw e
108+
}
119109

120-
logger.debug { formatLog(id, runId, "Finished agent execution") }
121-
pipeline.onAgentCompleted(eventId, context.executionInfo, id, runId, result)
110+
logger.debug { formatLog(id, runId, "Finished agent execution") }
111+
pipeline.onAgentCompleted(eventId, context.executionInfo, id, runId, result)
122112

123-
agentStateMutex.withLock {
124-
state = if (result != null) {
125-
State.Finished(result)
126-
} else {
127-
State.Failed(Exception("result is null"))
128-
}
113+
agentStateMutex.withLock {
114+
state = if (result != null) {
115+
State.Finished(result)
116+
} else {
117+
State.Failed(Exception("result is null"))
129118
}
130-
131-
result ?: error("result is null")
132119
}
120+
121+
result ?: error("result is null")
133122
}
134123
}
135124

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

Lines changed: 0 additions & 56 deletions
This file was deleted.

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

Lines changed: 0 additions & 43 deletions
This file was deleted.

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

Lines changed: 18 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,12 @@
11
package ai.koog.agents.core.agent.entity
22

33
import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
4-
import ai.koog.agents.core.agent.context.element.NodeInfoContextElement
5-
import ai.koog.agents.core.agent.context.element.getNodeInfoElement
64
import ai.koog.agents.core.agent.context.with
75
import ai.koog.agents.core.annotation.InternalAgentsApi
86
import io.github.oshai.kotlinlogging.KotlinLogging
97
import kotlinx.coroutines.CancellationException
10-
import kotlinx.coroutines.withContext
118
import kotlin.reflect.KType
129
import kotlin.uuid.ExperimentalUuidApi
13-
import kotlin.uuid.Uuid
1410

1511
/**
1612
* Represents an abstract node in an AI agent strategy graph, responsible for executing a specific
@@ -158,26 +154,24 @@ public open class AIAgentNode<TInput, TOutput> internal constructor(
158154
@OptIn(ExperimentalUuidApi::class)
159155
override suspend fun execute(context: AIAgentGraphContextBase, input: TInput): TOutput =
160156
context.with(id) { executionInfo, eventId ->
161-
withContext(NodeInfoContextElement(Uuid.random().toString(), getNodeInfoElement()?.id, name, input, inputType)) {
162-
logger.debug { "Start executing node (name: $name)" }
163-
context.pipeline.onNodeExecutionStarting(eventId, executionInfo, this@AIAgentNode, context, input, inputType)
164-
165-
val output =
166-
try {
167-
val executeResult = context.execute(input)
168-
logger.trace { "Finished executing node (name: $name) with output: $executeResult" }
169-
executeResult
170-
} catch (e: CancellationException) {
171-
throw e
172-
} catch (e: Exception) {
173-
logger.error(e) { "Error executing node (name: $name): ${e.message}" }
174-
context.pipeline.onNodeExecutionFailed(eventId, executionInfo, this@AIAgentNode, context, input, inputType, e)
175-
throw e
176-
}
177-
178-
context.pipeline.onNodeExecutionCompleted(eventId, executionInfo, this@AIAgentNode, context, input, inputType, output, outputType)
179-
output
180-
}
157+
logger.debug { "Start executing node (name: $name)" }
158+
context.pipeline.onNodeExecutionStarting(eventId, executionInfo, this@AIAgentNode, context, input, inputType)
159+
160+
val output =
161+
try {
162+
val executeResult = context.execute(input)
163+
logger.trace { "Finished executing node (name: $name) with output: $executeResult" }
164+
executeResult
165+
} catch (e: CancellationException) {
166+
throw e
167+
} catch (e: Exception) {
168+
logger.error(e) { "Error executing node (name: $name): ${e.message}" }
169+
context.pipeline.onNodeExecutionFailed(eventId, executionInfo, this@AIAgentNode, context, input, inputType, e)
170+
throw e
171+
}
172+
173+
context.pipeline.onNodeExecutionCompleted(eventId, executionInfo, this@AIAgentNode, context, input, inputType, output, outputType)
174+
output
181175
}
182176
}
183177

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

Lines changed: 43 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@ package ai.koog.agents.core.agent.entity
33
import ai.koog.agents.core.agent.context.AIAgentContext
44
import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
55
import ai.koog.agents.core.agent.context.DetachedPromptExecutorAPI
6-
import ai.koog.agents.core.agent.context.element.NodeInfoContextElement
7-
import ai.koog.agents.core.agent.context.element.getNodeInfoElement
86
import ai.koog.agents.core.agent.context.getAgentContextData
97
import ai.koog.agents.core.agent.context.store
108
import ai.koog.agents.core.agent.context.with
@@ -26,11 +24,9 @@ import ai.koog.prompt.structure.json.JsonStructure
2624
import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator
2725
import io.github.oshai.kotlinlogging.KotlinLogging
2826
import kotlinx.coroutines.CancellationException
29-
import kotlinx.coroutines.withContext
3027
import kotlinx.serialization.Serializable
3128
import kotlin.reflect.KType
3229
import kotlin.uuid.ExperimentalUuidApi
33-
import kotlin.uuid.Uuid
3430

3531
/**
3632
* [AIAgentSubgraph] represents a structured subgraph within an AI agent workflow. It serves as a logical
@@ -163,61 +159,59 @@ public open class AIAgentSubgraph<TInput, TOutput>(
163159
@OptIn(InternalAgentsApi::class, DetachedPromptExecutorAPI::class, ExperimentalUuidApi::class)
164160
override suspend fun execute(context: AIAgentGraphContextBase, input: TInput): TOutput? =
165161
context.with { executionInfo, eventId ->
166-
withContext(NodeInfoContextElement(Uuid.random().toString(), getNodeInfoElement()?.id, name, input, inputType)) {
167-
val newTools = selectTools(context)
168-
169-
// Copy inner context with new tools, model and LLM params.
170-
val initialLLMContext = context.llm
171-
172-
context.replace(
173-
context.copy(
174-
llm = context.llm.copy(
175-
tools = newTools,
176-
model = llmModel ?: context.llm.model,
177-
prompt = context.llm.prompt.copy(params = llmParams ?: context.llm.prompt.params),
178-
responseProcessor = responseProcessor
179-
),
162+
val newTools = selectTools(context)
163+
164+
// Copy inner context with new tools, model and LLM params.
165+
val initialLLMContext = context.llm
166+
167+
context.replace(
168+
context.copy(
169+
llm = context.llm.copy(
170+
tools = newTools,
171+
model = llmModel ?: context.llm.model,
172+
prompt = context.llm.prompt.copy(params = llmParams ?: context.llm.prompt.params),
173+
responseProcessor = responseProcessor
180174
),
181-
)
175+
),
176+
)
182177

183-
runIfNotStrategy(context) {
184-
pipeline.onSubgraphExecutionStarting(eventId, executionInfo, this@AIAgentSubgraph, context, input, inputType)
185-
}
178+
runIfNotStrategy(context) {
179+
pipeline.onSubgraphExecutionStarting(eventId, executionInfo, this@AIAgentSubgraph, context, input, inputType)
180+
}
186181

187-
// Execute the subgraph with an inner context and get the result and updated prompt.
188-
val result = try {
189-
executeWithInnerContext(context, input)
190-
} catch (e: CancellationException) {
191-
throw e
192-
} catch (e: Exception) {
193-
logger.error(e) { "Exception during executing subgraph '$name': ${e.message}" }
194-
runIfNotStrategy(context) {
195-
pipeline.onSubgraphExecutionFailed(eventId, executionInfo, this@AIAgentSubgraph, context, input, inputType, e)
196-
}
197-
throw e
182+
// Execute the subgraph with an inner context and get the result and updated prompt.
183+
val result = try {
184+
executeWithInnerContext(context, input)
185+
} catch (e: CancellationException) {
186+
throw e
187+
} catch (e: Exception) {
188+
logger.error(e) { "Exception during executing subgraph '$name': ${e.message}" }
189+
runIfNotStrategy(context) {
190+
pipeline.onSubgraphExecutionFailed(eventId, executionInfo, this@AIAgentSubgraph, context, input, inputType, e)
198191
}
192+
throw e
193+
}
199194

200-
// Restore original LLM context with updated message history.
201-
context.replace(
202-
context.copy(
203-
llm = initialLLMContext.copy(
204-
prompt = context.llm.prompt.copy(params = initialLLMContext.prompt.params)
205-
),
195+
// Restore original LLM context with updated message history.
196+
context.replace(
197+
context.copy(
198+
llm = initialLLMContext.copy(
199+
prompt = context.llm.prompt.copy(params = initialLLMContext.prompt.params)
206200
),
207-
)
201+
),
202+
)
208203

209-
val innerForcedData = context.getAgentContextData()
210-
211-
if (innerForcedData != null) {
212-
context.store(innerForcedData)
213-
}
204+
val innerForcedData = context.getAgentContextData()
214205

215-
runIfNotStrategy(context) {
216-
pipeline.onSubgraphExecutionCompleted(eventId, executionInfo, this@AIAgentSubgraph, context, input, inputType, result, outputType)
217-
}
206+
if (innerForcedData != null) {
207+
context.store(innerForcedData)
208+
}
218209

219-
result
210+
runIfNotStrategy(context) {
211+
pipeline.onSubgraphExecutionCompleted(eventId, executionInfo, this@AIAgentSubgraph, context, input, inputType, result, outputType)
220212
}
213+
214+
result
221215
}
222216

223217
@OptIn(InternalAgentsApi::class)

0 commit comments

Comments
 (0)