Skip to content

Commit 45228f4

Browse files
RizzenRizzen
committed
moving graph-specific pipeline methods to graph pipeline
1 parent 411b84b commit 45228f4

File tree

33 files changed

+348
-171
lines changed

33 files changed

+348
-171
lines changed

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

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

3-
import ai.koog.agents.core.agent.entity.AIAgentStrategy
3+
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
44
import ai.koog.agents.core.dsl.builder.forwardTo
55
import ai.koog.agents.core.dsl.builder.strategy
66
import ai.koog.agents.core.dsl.extension.nodeExecuteMultipleTools
@@ -30,7 +30,7 @@ import ai.koog.agents.core.dsl.extension.onToolCall
3030
* - SingleRunMode.PARALLEL: Executes multiple tool calls in parallel.
3131
* @return An instance of AIAgentStrategy configured according to the specified single-run mode.
3232
*/
33-
public fun singleRunStrategy(runMode: ToolCalls = ToolCalls.SINGLE_RUN_SEQUENTIAL): AIAgentStrategy<String, String> =
33+
public fun singleRunStrategy(runMode: ToolCalls = ToolCalls.SINGLE_RUN_SEQUENTIAL): AIAgentGraphStrategy<String, String> =
3434
when (runMode) {
3535
ToolCalls.SEQUENTIAL -> singleRunWithParallelAbility(false)
3636
ToolCalls.PARALLEL -> singleRunWithParallelAbility(true)

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ package ai.koog.agents.core.agent
33
import ai.koog.agents.core.agent.GraphAIAgent.FeatureContext
44
import ai.koog.agents.core.agent.config.AIAgentConfig
55
import ai.koog.agents.core.agent.config.AIAgentConfigBase
6-
import ai.koog.agents.core.agent.entity.AIAgentStrategy
6+
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
77
import ai.koog.agents.core.tools.ToolRegistry
88
import ai.koog.prompt.dsl.prompt
99
import ai.koog.prompt.executor.model.PromptExecutor
@@ -30,7 +30,7 @@ import kotlin.uuid.Uuid
3030
@OptIn(ExperimentalUuidApi::class)
3131
public inline fun <reified Input, reified Output> AIAgent(
3232
promptExecutor: PromptExecutor,
33-
strategy: AIAgentStrategy<Input, Output>,
33+
strategy: AIAgentGraphStrategy<Input, Output>,
3434
agentConfig: AIAgentConfigBase,
3535
id: String = Uuid.random().toString(),
3636
toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
@@ -52,7 +52,7 @@ public inline fun <reified Input, reified Output> AIAgent(
5252
* Convenience builder that creates an instance of an [AIAgent] with string input and output and the specified parameters.
5353
*
5454
* @param executor The [PromptExecutor] responsible for executing prompts.
55-
* @param strategy The [AIAgentStrategy] defining the agent's behavior. Default is a single-run strategy.
55+
* @param strategy The [AIAgentGraphStrategy] defining the agent's behavior. Default is a single-run strategy.
5656
* @param systemPrompt The system-level prompt context for the agent. Default is an empty string.
5757
* @param llmModel The language model to be used by the agent.
5858
* @param temperature The sampling temperature for the language model, controlling randomness. Default is 1.0.
@@ -67,7 +67,7 @@ public fun AIAgent(
6767
executor: PromptExecutor,
6868
llmModel: LLModel,
6969
id: String = Uuid.random().toString(),
70-
strategy: AIAgentStrategy<String, String> = singleRunStrategy(),
70+
strategy: AIAgentGraphStrategy<String, String> = singleRunStrategy(),
7171
systemPrompt: String = "",
7272
temperature: Double = 1.0,
7373
numberOfChoices: Int = 1,

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

Lines changed: 61 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,79 @@
11
package ai.koog.agents.core.agent
22

33
import ai.koog.agents.core.agent.config.AIAgentConfigBase
4+
import ai.koog.agents.core.feature.AIAgentFeature
5+
import ai.koog.agents.core.feature.AIAgentNonGraphFeature
6+
import ai.koog.agents.core.feature.AIAgentNonGraphPipeline
7+
import ai.koog.agents.core.feature.config.FeatureConfig
48
import ai.koog.agents.core.tools.ToolRegistry
59
import ai.koog.prompt.executor.model.PromptExecutor
10+
import kotlinx.coroutines.sync.Mutex
11+
import kotlinx.coroutines.sync.withLock
612

7-
internal class EssentialAIAgent<Input, Output>(
13+
/**
14+
* Represents the core AI agent for processing input and generating output using
15+
* a defined configuration, toolset, and prompt execution pipeline.
16+
*
17+
* @param Input The type of input data expected by the agent.
18+
* @param Output The type of output data produced by the agent.
19+
* @param id The unique identifier for the agent instance.
20+
* @param promptExecutor The executor responsible for processing prompts and interacting with language models.
21+
* @param agentConfig The configuration for the agent, including the prompt structure and execution parameters.
22+
* @param toolRegistry The registry of tools available for the agent. Defaults to an empty registry if not specified.
23+
*/
24+
public class EssentialAIAgent<Input, Output>(
825
override val id: String,
926
public val promptExecutor: PromptExecutor,
1027
public val agentConfig: AIAgentConfigBase,
1128
public val toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
1229
) : AIAgent<Input, Output> {
1330

31+
/**
32+
* Represents a context for managing and configuring features in an AI agent.
33+
* Provides functionality to install and configure features into a specific instance of an AI agent.
34+
*/
35+
public class FeatureContext internal constructor(private val agent: EssentialAIAgent<*, *>) {
36+
/**
37+
* Installs and configures a feature into the current AI agent context.
38+
*
39+
* @param feature the feature to be added, defined by an implementation of [AIAgentFeature], which provides specific functionality
40+
* @param configure an optional lambda to customize the configuration of the feature, where the provided [Config] can be modified
41+
*/
42+
public fun <Config : FeatureConfig, Feature : Any> install(
43+
feature: AIAgentNonGraphFeature<Config, Feature>,
44+
configure: Config.() -> Unit = {}
45+
) {
46+
agent.install(feature, configure)
47+
}
48+
}
49+
50+
private var isRunning = false
51+
52+
private val runningMutex = Mutex()
53+
private val pipeline = AIAgentNonGraphPipeline()
54+
55+
private fun <Config : FeatureConfig, Feature : Any> install(feature: AIAgentNonGraphFeature<Config, Feature>, configure: Config.() -> Unit) {
56+
pipeline.install(feature, configure)
57+
}
58+
1459
override suspend fun run(agentInput: Input): Output {
15-
TODO("Not yet implemented")
60+
runningMutex.withLock {
61+
if (isRunning) {
62+
throw IllegalStateException("Agent is already running")
63+
}
64+
isRunning = true
65+
}
66+
67+
68+
runningMutex.withLock {
69+
isRunning = false
70+
}
71+
72+
TODO()
1673
}
1774

1875
override suspend fun close() {
19-
TODO("Not yet implemented")
76+
pipeline.onAgentBeforeClosed(agentId = id)
77+
pipeline.closeFeaturesStreamProviders()
2078
}
2179
}

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package ai.koog.agents.core.agent
2+
3+
import ai.koog.agents.core.agent.config.AIAgentConfigBase
4+
import ai.koog.agents.core.agent.context.AIAgentContextBase
5+
import ai.koog.agents.core.agent.context.AIAgentLLMContext
6+
import ai.koog.agents.core.agent.entity.AIAgentStateManager
7+
import ai.koog.agents.core.agent.entity.AIAgentStorage
8+
import ai.koog.agents.core.agent.entity.AIAgentStorageKey
9+
import ai.koog.agents.core.environment.AIAgentEnvironment
10+
import ai.koog.agents.core.feature.AIAgentFeature
11+
import ai.koog.prompt.message.Message
12+
13+
/**
14+
*/
15+
public class EssentialAIAgentContext(
16+
override val environment: AIAgentEnvironment,
17+
override val agentId: String,
18+
override val runId: String,
19+
override val agentInput: Any?,
20+
override val config: AIAgentConfigBase,
21+
override val llm: AIAgentLLMContext,
22+
override val stateManager: AIAgentStateManager,
23+
override val storage: AIAgentStorage,
24+
override val strategyName: String
25+
) : AIAgentContextBase {
26+
27+
private val storeMap: MutableMap<AIAgentStorageKey<*>, Any> = mutableMapOf()
28+
29+
override fun store(key: AIAgentStorageKey<*>, value: Any) {
30+
storeMap[key] = value
31+
}
32+
33+
override fun <T> get(key: AIAgentStorageKey<*>): T? {
34+
@Suppress("UNCHECKED_CAST")
35+
return storeMap[key] as T?
36+
}
37+
38+
override fun remove(key: AIAgentStorageKey<*>): Boolean {
39+
return storeMap.remove(key) != null
40+
}
41+
42+
override fun <Feature : Any> feature(key: AIAgentStorageKey<Feature>): Feature? {
43+
// Essential context has no pipeline-installed features by default.
44+
// Users may store arbitrary values using store/get if needed.
45+
return null
46+
}
47+
48+
override fun <Feature : Any> feature(feature: AIAgentFeature<*, Feature>): Feature? = feature(feature.key)
49+
50+
override suspend fun getHistory(): List<Message> {
51+
return llm.readSession { prompt.messages }
52+
}
53+
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@ import ai.koog.agents.core.agent.context.AIAgentLLMContext
88
import ai.koog.agents.core.agent.context.element.AgentRunInfoContextElement
99
import ai.koog.agents.core.agent.context.getAgentContextData
1010
import ai.koog.agents.core.agent.context.removeAgentContextData
11+
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
1112
import ai.koog.agents.core.agent.entity.AIAgentStateManager
1213
import ai.koog.agents.core.agent.entity.AIAgentStorage
13-
import ai.koog.agents.core.agent.entity.AIAgentStrategy
1414
import ai.koog.agents.core.annotation.InternalAgentsApi
1515
import ai.koog.agents.core.environment.GenericAgentEnvironment
1616
import ai.koog.agents.core.feature.AIAgentFeature
@@ -56,7 +56,7 @@ public open class GraphAIAgent<Input, Output>(
5656
public val inputType: KType,
5757
public val outputType: KType,
5858
public val promptExecutor: PromptExecutor,
59-
private val strategy: AIAgentStrategy<Input, Output>,
59+
private val strategy: AIAgentGraphStrategy<Input, Output>,
6060
public val agentConfig: AIAgentConfigBase,
6161
override val id: String = Uuid.random().toString(),
6262
public val toolRegistry: ToolRegistry = ToolRegistry.EMPTY,

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

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import ai.koog.agents.core.agent.entity.AIAgentStorageKey
77
import ai.koog.agents.core.environment.AIAgentEnvironment
88
import ai.koog.agents.core.feature.AIAgentFeature
99
import ai.koog.prompt.message.Message
10-
import kotlin.reflect.KType
1110

1211
/**
1312
* The [AIAgentContextBase] interface represents the context of an AI agent in the lifecycle.
@@ -48,11 +47,6 @@ public interface AIAgentContextBase {
4847
*/
4948
public val agentInput: Any?
5049

51-
/**
52-
* [KType] representing the type of the [agentInput]
53-
*/
54-
public val agentInputType: KType
55-
5650
/**
5751
* Represents the configuration for an AI agent.
5852
*

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,11 @@ public interface AIAgentGraphContextBase : AIAgentContextBase {
4141
*/
4242
public val pipeline: AIAgentGraphPipeline
4343

44+
/**
45+
* [KType] representing the type of the [agentInput]
46+
*/
47+
public val agentInputType: KType
48+
4449
/**
4550
* Creates a copy of the current [AIAgentGraphContext], allowing for selective overriding of its properties.
4651
*

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentStrategy.kt renamed to ‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentGraphStrategy.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import kotlinx.serialization.serializer
1616
* @property nodeFinish The finishing node of the strategy, marking the subgraph's endpoint.
1717
* @property toolSelectionStrategy The strategy responsible for determining the toolset available during subgraph execution.
1818
*/
19-
public class AIAgentStrategy<Input, Output>(
19+
public class AIAgentGraphStrategy<Input, Output>(
2020
override val name: String,
2121
public val nodeStart: StartNode<Input>,
2222
public val nodeFinish: FinishNode<Output>,

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/builder/AIAgentStrategyBuilder.kt

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package ai.koog.agents.core.dsl.builder
22

3-
import ai.koog.agents.core.agent.entity.AIAgentStrategy
3+
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
44
import ai.koog.agents.core.agent.entity.FinishNode
55
import ai.koog.agents.core.agent.entity.StartNode
66
import ai.koog.agents.core.agent.entity.ToolSelectionStrategy
@@ -20,12 +20,12 @@ public class AIAgentStrategyBuilder<Input, Output>(
2020
inputType: KType,
2121
outputType: KType,
2222
private val toolSelectionStrategy: ToolSelectionStrategy,
23-
) : AIAgentSubgraphBuilderBase<Input, Output>(), BaseBuilder<AIAgentStrategy<Input, Output>> {
23+
) : AIAgentSubgraphBuilderBase<Input, Output>(), BaseBuilder<AIAgentGraphStrategy<Input, Output>> {
2424
public override val nodeStart: StartNode<Input> = StartNode(type = inputType)
2525
public override val nodeFinish: FinishNode<Output> = FinishNode(type = outputType)
2626

27-
override fun build(): AIAgentStrategy<Input, Output> {
28-
val strategy = AIAgentStrategy(
27+
override fun build(): AIAgentGraphStrategy<Input, Output> {
28+
val strategy = AIAgentGraphStrategy(
2929
name = name,
3030
nodeStart = nodeStart,
3131
nodeFinish = nodeFinish,
@@ -49,7 +49,7 @@ public inline fun <reified Input, reified Output> strategy(
4949
name: String,
5050
toolSelectionStrategy: ToolSelectionStrategy = ToolSelectionStrategy.ALL,
5151
init: AIAgentStrategyBuilder<Input, Output>.() -> Unit,
52-
): AIAgentStrategy<Input, Output> {
52+
): AIAgentGraphStrategy<Input, Output> {
5353
return AIAgentStrategyBuilder<Input, Output>(
5454
name = name,
5555
inputType = typeOf<Input>(),

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/builder/AIAgentSubgraphBuilder.kt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ package ai.koog.agents.core.dsl.builder
55
import ai.koog.agents.core.agent.context.AIAgentContextBase
66
import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
77
import ai.koog.agents.core.agent.context.getAgentContextData
8+
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
89
import ai.koog.agents.core.agent.entity.AIAgentNodeBase
9-
import ai.koog.agents.core.agent.entity.AIAgentStrategy
1010
import ai.koog.agents.core.agent.entity.AIAgentSubgraph
1111
import ai.koog.agents.core.agent.entity.FinishNode
1212
import ai.koog.agents.core.agent.entity.StartNode
@@ -184,7 +184,7 @@ public abstract class AIAgentSubgraphBuilderBase<Input, Output> {
184184
return "$parentPath:${node.id}"
185185
}
186186

187-
internal fun buildSubgraphMetadata(start: StartNode<Input>, parentName: String, strategy: AIAgentStrategy<Input, Output>): SubgraphMetadata {
187+
internal fun buildSubgraphMetadata(start: StartNode<Input>, parentName: String, strategy: AIAgentGraphStrategy<Input, Output>): SubgraphMetadata {
188188
val subgraphNodes = buildSubGraphNodesMap(start, parentName)
189189
subgraphNodes[parentName] = strategy
190190

0 commit comments

Comments
 (0)