koog-agents-optimization/ai.koog.agents.optimization.training

Package-level declarations

Types

Controls truncation of string values in substep action logs.

abstract class CapturingPromptExecutor : PromptExecutor

A PromptExecutor that accumulates the LLMConsumption of the calls it intercepts and hands it back via collectAndClear. The training infrastructure creates one per leaf prompt call (via the session's capturingExecutorFactory) and reads its consumption between retry attempts; the concrete provider-aware extraction lives in the app.

ConsumptionCapturingPromptExecutor

abstract class ConsumptionCapturingPromptExecutor(delegate: PromptExecutor) : CapturingPromptExecutor

A CapturingPromptExecutor that wraps a delegate executor, runs each call through it, and accumulates the LLMConsumption extracted from the returned Message.Response metadata.

FailureAnalyzer

interface FailureAnalyzer

Converts a caught exception into an AnalyzedFailure. This is the injection point for provider-specific failure recognition: the library calls it from the retry / stage-failure paths, the app supplies the recognizers (provider-specific and network heuristics).

LiteLLMConsumptionCapturingPromptExecutor

class LiteLLMConsumptionCapturingPromptExecutor(delegate: PromptExecutor) : ConsumptionCapturingPromptExecutor

LiteLLM variant: reads inputTokensCount / outputTokensCount / totalTokensCount from the first response that carries token metadata.

NoOpFailureAnalyzer

object NoOpFailureAnalyzer : FailureAnalyzer

A provider-agnostic FailureAnalyzer that classifies every exception as FailureKind.UNKNOWN / TransiencyLevel.UNKNOWN — i.e. it makes no retryable-vs-terminal judgement.

PrematureExecutionStopDecision

data class PrematureExecutionStopDecision(val conditionMet: Boolean, val conditionMetReason: () -> String)

Outcome of an early-stop check, e.g. the earlyStop callback of ai.koog.agents.optimization.training.dsl.StageScope.iterateDataset.

StageScopeImpl

class StageScopeImpl<Input, Output, InputLabel>(resources: TrainingResources<Input, Output, InputLabel>, logger: KLogger, val record: StageRecord, parentScope: StageScopeImpl<Input, Output, InputLabel>?, onUpdate: (isStageUpdate: Boolean) -> Unit) : StageScope<Input, Output, InputLabel>

Core implementation of StageScope.

TrainingLogging

object TrainingLogging

Log formatting utilities for the training infrastructure.

TrainingResources

class TrainingResources<Input, Output, InputLabel>(val trackedAgent: GraphAIAgent<Input, Output>, val dataset: TrainSet<Input, InputLabel>, val substepPromptExecutor: PromptExecutor, val metric: (TrainSetItem<Input, InputLabel>, Output) -> Double, val threshold: Double, val serializers: DatasetExecutionSerializers<Input, Output, InputLabel>, val actionLogTruncation: ActionLogTruncation, val retriesPolicy: RetryPolicy, val failureAnalyzer: FailureAnalyzer, val capturingExecutorFactory: (PromptExecutor) -> CapturingPromptExecutor, val consumptionCollector: () -> LLMConsumption? = { null }, val runAgentAttempt: suspend (agent: GraphAIAgent<Input, Output>, input: Input) -> Pair<GraphAIAgent<Input, Output>, Output> = { agent, input -> agent to agent.run(input) }, val withItemContext: suspend (itemId: String, block: suspend () -> Unit) -> Unit = { _, block -> block() })

Shared execution resources created once per training session.

TrainingSession

class TrainingSession<Input, Output, InputLabel>(resources: TrainingResources<Input, Output, InputLabel>, logger: KLogger, trainingName: ExperimentName, recordsFilePath: ResilientPath?, throwIfRecordsFileExists: Boolean, podName: String?, spendLimit: JsonElement?, retries: JsonElement?, abortPolicies: List<AbortPolicy> = emptyList(), progressListener: (StageRecord) -> Unit = {})

Lifecycle wrapper for a training session.

Functions

trainingSession

fun <Input, Output, InputLabel> trainingSession(experimentName: ExperimentName, trackedAgent: GraphAIAgent<Input, Output>, dataset: TrainSet<Input, InputLabel>, substepPromptExecutor: PromptExecutor, metric: (TrainSetItem<Input, InputLabel>, Output) -> Double, serializers: DatasetExecutionSerializers<Input, Output, InputLabel>, threshold: Double = 0.9, consumptionCollector: () -> LLMConsumption? = { null }, logger: KLogger = KotlinLogging.logger {}): TrainingSession<Input, Output, InputLabel>

Builds a TrainingSession for the common case in a single call, filling in library defaults: no smart failure recognition (NoOpFailureAnalyzer), LiteLLM token capture (LiteLLMConsumptionCapturingPromptExecutor, which reads koog's standard token counts), single-attempt runs (RetryPolicy.None), ActionLogTruncation.DEFAULT, no records file, and no cluster/spend-limit metadata.