agents-core/ai.koog.agents.core.feature/AIAgentPipeline

AIAgentPipeline

Pipeline for AI agent features that provides interception points for various agent lifecycle events.

The pipeline allows features to:

Be installed into agent contexts
Intercept agent creation and environment transformation
Intercept strategy execution before and after it happens
Intercept node execution before and after it happens
Intercept LLM (Language Learning Model) calls before and after they happen
Intercept tool calls before and after they happen

This pipeline serves as the central mechanism for extending and customizing agent behavior through a flexible interception system. Features can be installed with custom configurations and can hook into different stages of the agent's execution lifecycle.

Constructors

AIAgentPipeline

constructor()

Functions

getAgentFeatures

fun getAgentFeatures(context: AIAgentContextBase): Map<AIAgentStorageKey<*>, Any>

Retrieves all features associated with the given agent context.

install

fun <Config : FeatureConfig, Feature : Any> install(feature: AIAgentFeature<Config, Feature>, configure: Config.() -> Unit)

Installs a feature into the pipeline with the provided configuration.

interceptAfterLLMCall

fun <TFeature : Any> interceptAfterLLMCall(interceptContext: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: AfterLLMCallContext) -> Unit)

Intercepts LLM calls after they are made to process or log the response.

interceptAfterNode

fun <TFeature : Any> interceptAfterNode(interceptContext: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: NodeAfterExecuteContext) -> Unit)

Intercepts node execution after it completes.

interceptAgentBeforeClosed

fun <TFeature : Any> interceptAgentBeforeClosed(context: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: AgentBeforeCloseContext) -> Unit)

Intercepts and sets a handler to be invoked before an agent is closed.

interceptAgentFinished

fun <TFeature : Any> interceptAgentFinished(context: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: AgentFinishedContext) -> Unit)

Intercepts the completion of an agent's operation and assigns a custom handler to process the result.

interceptAgentRunError

fun <TFeature : Any> interceptAgentRunError(context: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: AgentRunErrorContext) -> Unit)

Intercepts and handles errors occurring during the execution of an AI agent's strategy.

interceptBeforeAgentStarted

fun <TFeature : Any> interceptBeforeAgentStarted(context: InterceptContext<TFeature>, handle: suspend (AgentStartContext<TFeature>) -> Unit)

Intercepts on before an agent started to modify or enhance the agent.

interceptBeforeLLMCall

fun <TFeature : Any> interceptBeforeLLMCall(interceptContext: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: BeforeLLMCallContext) -> Unit)

Intercepts LLM calls before they are made to modify or log the prompt.

interceptBeforeNode

fun <TFeature : Any> interceptBeforeNode(interceptContext: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: NodeBeforeExecuteContext) -> Unit)

Intercepts node execution before it starts.

interceptContextAgentFeature

fun <TFeature : Any> interceptContextAgentFeature(feature: AIAgentFeature<*, TFeature>, handler: AgentContextHandler<TFeature>)

Sets a feature handler for agent context events.

interceptEnvironmentCreated

fun <TFeature : Any> interceptEnvironmentCreated(context: InterceptContext<TFeature>, transform: AgentTransformEnvironmentContext<TFeature>.(AIAgentEnvironment) -> AIAgentEnvironment)

Intercepts environment creation to allow features to modify or enhance the agent environment.

interceptStrategyFinished

fun <TFeature : Any> interceptStrategyFinished(context: InterceptContext<TFeature>, handle: suspend (StrategyFinishContext<TFeature>) -> Unit)

Sets up an interceptor to handle the completion of a strategy for the given feature.

interceptStrategyStarted

fun <TFeature : Any> interceptStrategyStarted(context: InterceptContext<TFeature>, handle: suspend (StrategyStartContext<TFeature>) -> Unit)

Intercepts strategy started event to perform actions when an agent strategy begins execution.

interceptToolCall

fun <TFeature : Any> interceptToolCall(interceptContext: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: ToolCallContext) -> Unit)

Intercepts and handles tool calls for the specified feature and its implementation. Updates the tool call handler for the given feature key with a custom handler.

interceptToolCallFailure

fun <TFeature : Any> interceptToolCallFailure(interceptContext: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: ToolCallFailureContext) -> Unit)

Sets up an interception mechanism to handle tool call failures for a specific feature.

interceptToolCallResult

fun <TFeature : Any> interceptToolCallResult(interceptContext: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: ToolCallResultContext) -> Unit)

Intercepts the result of a tool call with a custom handler for a specific feature.

interceptToolValidationError

fun <TFeature : Any> interceptToolValidationError(interceptContext: InterceptContext<TFeature>, handle: suspend TFeature.(eventContext: ToolValidationErrorContext) -> Unit)

Intercepts validation errors encountered during the execution of tools associated with the specified feature.

onAfterLLMCall

suspend fun onAfterLLMCall(runId: String, prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>, responses: List<Message.Response>, moderationResponse: ModerationResult? = null)

Notifies all registered LLM handlers after a language model call has completed.

onAfterNode

suspend fun onAfterNode(node: AIAgentNodeBase<*, *>, context: AIAgentContextBase, input: Any?, output: Any?, inputType: KType, outputType: KType)

Notifies all registered node handlers after a node has been executed.

onAgentBeforeClosed

suspend fun onAgentBeforeClosed(agentId: String)

Invoked before an agent is closed to perform necessary pre-closure operations.

onAgentFinished

suspend fun onAgentFinished(agentId: String, runId: String, result: Any?, resultType: KType)

Notifies all registered handlers that an agent has finished execution.

onAgentRunError

suspend fun onAgentRunError(agentId: String, runId: String, throwable: Throwable)

Notifies all registered handlers about an error that occurred during agent execution.

onBeforeAgentStarted

suspend fun onBeforeAgentStarted(runId: String, agent: AIAgent<*, *>, strategy: AIAgentStrategy<*, *>, context: AIAgentContextBase)

Notifies all registered handlers that an agent has started execution.

onBeforeLLMCall

suspend fun onBeforeLLMCall(runId: String, prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>)

Notifies all registered LLM handlers before a language model call is made.

onBeforeNode

suspend fun onBeforeNode(node: AIAgentNodeBase<*, *>, context: AIAgentContextBase, input: Any?, inputType: KType)

Notifies all registered node handlers before a node is executed.

onStrategyFinished

suspend fun onStrategyFinished(strategy: AIAgentStrategy<*, *>, context: AIAgentContextBase, result: Any?, resultType: KType)

Notifies all registered strategy handlers that a strategy has finished execution.

onStrategyStarted

suspend fun onStrategyStarted(strategy: AIAgentStrategy<*, *>, context: AIAgentContextBase)

Notifies all registered strategy handlers that a strategy has started execution.

onToolCall

suspend fun onToolCall(runId: String, toolCallId: String?, tool: Tool<*, *>, toolArgs: ToolArgs)

Notifies all registered tool handlers when a tool is called.

onToolCallFailure

suspend fun onToolCallFailure(runId: String, toolCallId: String?, tool: Tool<*, *>, toolArgs: ToolArgs, throwable: Throwable)

Notifies all registered tool handlers when a tool call fails with an exception.

onToolCallResult

suspend fun onToolCallResult(runId: String, toolCallId: String?, tool: Tool<*, *>, toolArgs: ToolArgs, result: ToolResult?)

Notifies all registered tool handlers about the result of a tool call.

onToolValidationError

suspend fun onToolValidationError(runId: String, toolCallId: String?, tool: Tool<*, *>, toolArgs: ToolArgs, error: String)

Notifies all registered tool handlers when a validation error occurs during a tool call.

transformEnvironment

fun transformEnvironment(strategy: AIAgentStrategy<*, *>, agent: AIAgent<*, *>, baseEnvironment: AIAgentEnvironment): AIAgentEnvironment

Transforms the agent environment by applying all registered environment transformers.