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

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.Assistant's metadata.

The training infrastructure creates one per leaf prompt call (via the session's capturingExecutorFactory) and reads accumulated consumption between retry attempts via collectAndClear. Subclasses supply the provider-specific extractConsumption.

Inheritors

LiteLLMConsumptionCapturingPromptExecutor

Constructors

ConsumptionCapturingPromptExecutor

constructor(delegate: PromptExecutor)

Functions

open override fun close()

Closes the underlying delegate executor.

collectAndClear

open override fun collectAndClear(): LLMConsumption?

Returns the consumption accumulated since the previous call (or since construction), then resets the accumulator. Returns null when nothing was captured.

execute

open suspend override fun execute(prompt: Prompt, model: ResolvedModel, tools: List<ToolDescriptor>): Message.Assistant

Pre-resolved-model overload of execute; captures consumption the same way.

open suspend override fun execute(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): Message.Assistant

Runs the call through delegate and accumulates any LLMConsumption reported by extractConsumption for the returned response. Accumulation is thread-safe.

executeBlocking

fun executeBlocking(prompt: Prompt, resolvedModel: ResolvedModel, tools: List<ToolDescriptor>): Message.Assistant

fun executeBlocking(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): Message.Assistant

executeMultipleChoices

open suspend fun executeMultipleChoices(prompt: Prompt, resolvedModel: ResolvedModel, tools: List<ToolDescriptor>): LLMChoice

open suspend fun executeMultipleChoices(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): LLMChoice

executeMultipleChoicesBlocking

fun executeMultipleChoicesBlocking(prompt: Prompt, resolvedModel: ResolvedModel, tools: List<ToolDescriptor>): LLMChoice

fun executeMultipleChoicesBlocking(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): LLMChoice

executeStreaming

open override fun executeStreaming(prompt: Prompt, resolvedModel: ResolvedModel, tools: List<ToolDescriptor>): Flow<StreamFrame>

Pre-resolved-model overload of executeStreaming; also captures nothing.

open override fun executeStreaming(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): Flow<StreamFrame>

Delegates streaming to delegate; no consumption is captured from streamed frames.

executeStreamingBlocking

fun executeStreamingBlocking(prompt: Prompt, resolvedModel: ResolvedModel, tools: List<ToolDescriptor>): Flow.Publisher<StreamFrame>

fun executeStreamingBlocking(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): Flow.Publisher<StreamFrame>

executeStructuredOrThrow

inline suspend fun <OutputStructT> PromptExecutor.executeStructuredOrThrow(prompt: Prompt, model: LLModel): OutputStructT

Executes prompt against model, requesting a structured response of type OutputStructT, and returns the parsed data.

getBasicJsonSchemaGenerator

open fun getBasicJsonSchemaGenerator(model: LLModel): BasicJsonSchemaGenerator

getStandardJsonSchemaGenerator

open fun getStandardJsonSchemaGenerator(model: LLModel): StandardJsonSchemaGenerator

models

open suspend fun models(): List<LLModel>

modelsBlocking

fun modelsBlocking(): List<LLModel>

moderate

open suspend override fun moderate(prompt: Prompt, model: ResolvedModel): ModerationResult

Pre-resolved-model overload of moderate; also captures nothing.

open suspend override fun moderate(prompt: Prompt, model: LLModel): ModerationResult

Delegates moderation to delegate; no consumption is captured.

moderateBlocking

fun moderateBlocking(prompt: Prompt, resolvedModel: ResolvedModel): ModerationResult

fun moderateBlocking(prompt: Prompt, model: LLModel): ModerationResult

resolveModel

open suspend override fun resolveModel(model: LLModel, promptExecutorOperation: PromptExecutorOperation): ResolvedModel

Delegates model resolution, so the wrapped executor keeps its fallback/routing behaviour.

resolveModelBlocking

fun resolveModelBlocking(requestedModel: LLModel, promptExecutorOperation: PromptExecutorOperation): ResolvedModel