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

Package-level declarations

Types

class ActionLogBuilder(val truncation: ActionLogTruncation, val json: Json)

DSL for building a stage action log. Obtained via StageScope.logAction.

data class CompletedAgentRun<Output>(val output: Output, val score: Double, val isSolved: Boolean, val usedAgent: GraphAIAgent<*, *>? = null)

Result of a successful agent run produced by StageScope.runAgent.

MatchedAgentRun

class MatchedAgentRun<Output, RunData>(val agentRun: CompletedAgentRun<Output>, val runData: RunData)

A successful CompletedAgentRun together with the PreparedAgentRun.runData that produced it, returned by the generic StageScope.runAgentWithRetries.

PreparedAgentRun

class PreparedAgentRun<Input, Output, RunData>(val agent: GraphAIAgent<Input, Output>, val runData: RunData)

An agent paired with caller-supplied per-attempt data, returned by the agentProvider of the generic StageScope.runAgentWithRetries so the data can be matched back to a successful run.

RunUntil

enum RunUntil : Enum<RunUntil>

Stop condition for the re-sampling StageScope.runAgentWithRetries overloads.

StageScope

interface StageScope<Input, Output, InputLabel>

The single unified scope for the training DSL.

TrainingDsl

@DslMarker

annotation class TrainingDsl

DslMarker for the training DSL; scopes receiver resolution within builders like ActionLogBuilder.

Functions

executePrompt

suspend fun StageScope<*, *, *>.executePrompt(prompt: Prompt, model: LLModel, retryPolicy: RetryPolicy? = null): Result<List<Message.Response>>

Executes an LLM prompt with consumption tracking. Creates a PromptExecutionRecord as a leaf record, recording elapsed time and consumption. On success, returns the raw LLM response messages.

executePromptOrThrow

suspend fun StageScope<*, *, *>.executePromptOrThrow(prompt: Prompt, model: LLModel, retryPolicy: RetryPolicy? = null): List<Message.Response>

.getOrThrow() shortcut over executePrompt; throws on a recorded execution failure.

executePromptStructured

inline suspend fun <T> StageScope<*, *, *>.executePromptStructured(prompt: Prompt, model: LLModel, retryPolicy: RetryPolicy? = null): Result<T>

Executes a structured LLM prompt with consumption tracking. Creates a PromptExecutionRecord and returns the deserialized structure.

executePromptStructuredOrThrow

inline suspend fun <T> StageScope<*, *, *>.executePromptStructuredOrThrow(prompt: Prompt, model: LLModel, retryPolicy: RetryPolicy? = null): T

.getOrThrow() shortcut over executePromptStructured; throws on a recorded execution failure.

mapIterableStageOrThrow

suspend fun <Input, Output, InputLabel, TFrom, TInto> StageScope<Input, Output, InputLabel>.mapIterableStageOrThrow(collection: Collection<TFrom>, name: String, itemBlock: suspend StageScope<Input, Output, InputLabel>.(TFrom) -> TInto): List<TInto>

Convenience wrapper that maps a collection, enclosing it in a dedicated stage. Pros: automatically sets substagesTotal to enable ETC calculation.

runAgentOrThrow

suspend fun <Input, Output, InputLabel> StageScope<Input, Output, InputLabel>.runAgentOrThrow(item: TrainSetItem<Input, InputLabel>, agentToRun: GraphAIAgent<Input, Output> = trackedAgent, retryPolicy: RetryPolicy? = null): CompletedAgentRun<Output>

.getOrThrow() shortcut over StageScope.runAgent; throws on a recorded run failure.

runIterableStage

suspend fun <Input, Output, InputLabel, T> StageScope<Input, Output, InputLabel>.runIterableStage(collection: Collection<T>, name: String, itemBlock: suspend StageScope<Input, Output, InputLabel>.(T) -> Unit): Result<Unit>

Convenience wrapper that iterates a collection, enclosing it in a dedicated stage. Pros: automatically sets substagesTotal to enable ETC calculation.

runIterableStageOrThrow

suspend fun <Input, Output, InputLabel, T> StageScope<Input, Output, InputLabel>.runIterableStageOrThrow(collection: Collection<T>, name: String, itemBlock: suspend StageScope<Input, Output, InputLabel>.(T) -> Unit)

runStageOrThrow

suspend fun <Input, Output, InputLabel, T> StageScope<Input, Output, InputLabel>.runStageOrThrow(name: String, substagesTotal: Int? = null, metrics: MetricsMap = standardStageMetrics(), block: suspend StageScope<Input, Output, InputLabel>.() -> T): T

.getOrThrow() shortcut over StageScope.runStage. See its KDoc for the metrics preset conventions; the default and parameter semantics are identical.

setAdditionalData

inline fun <T> StageScope<*, *, *>.setAdditionalData(data: T)

Serializes data with defaultExperimentsJson and stores it as this stage's additional data via StageScope.setAdditionalDataJson. Overwrites any previously set value.

truncateStrings

fun JsonElement.truncateStrings(truncation: ActionLogTruncation): JsonElement

Recursive string truncation over the entire JSON element tree