Add LLM as a Judge component

agents/agents-ext/build.gradle.kts

@@ -32,6 +32,7 @@ kotlin {
         jvmTest {
             dependencies {
                 implementation(kotlin("test-junit5"))
+                implementation(libs.mockk)
             }
         }
     }

agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/AIAgentSubgraphExt.kt

@@ -3,6 +3,7 @@ package ai.koog.agents.ext.agent
 import ai.koog.agents.core.agent.context.AIAgentGraphContextBase
 import ai.koog.agents.core.agent.entity.ToolSelectionStrategy
 import ai.koog.agents.core.agent.entity.createStorageKey
+import ai.koog.agents.core.annotation.InternalAgentsApi
 import ai.koog.agents.core.dsl.builder.AIAgentBuilderDslMarker
 import ai.koog.agents.core.dsl.builder.AIAgentSubgraphBuilderBase
 import ai.koog.agents.core.dsl.builder.AIAgentSubgraphDelegate
@@ -25,22 +26,9 @@ import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.params.LLMParams
 import kotlinx.serialization.KSerializer
-import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.serializer
 
-/**
- * Represents the result of a verification process for a subgraph.
- *
- * @property correct Indicates whether the subgraph verification was successful.
- * @property message A message providing details about the verification outcome.
- */
-@Serializable
-public data class VerifiedSubgraphResult(
-    val correct: Boolean,
-    val message: String,
-)
-
 /**
  * Utility object providing tools and methods for working with subgraphs and tasks in a controlled
  * and structured way. These utilities are designed to help finalize subgraph-related tasks and
@@ -240,22 +228,43 @@ public inline fun <reified Input, reified Output, reified OutputTransformed> AIA
 }
 
 /**
- * [subgraphWithTask] with [VerifiedSubgraphResult] result.
+ * [subgraphWithTask] with [CriticResult] result.
  * It verifies if the task was performed correctly or not, and describes the problems if any.
  */
+@OptIn(InternalAgentsApi::class)
 @Suppress("unused")
 @AIAgentBuilderDslMarker
-public inline fun <reified Input> AIAgentSubgraphBuilderBase<*, *>.subgraphWithVerification(
+public inline fun <reified Input : Any> AIAgentSubgraphBuilderBase<*, *>.subgraphWithVerification(
     toolSelectionStrategy: ToolSelectionStrategy,
     llmModel: LLModel? = null,
     llmParams: LLMParams? = null,
     noinline defineTask: suspend AIAgentGraphContextBase.(input: Input) -> String
-): AIAgentSubgraphDelegate<Input, VerifiedSubgraphResult> = subgraphWithTask(
-    toolSelectionStrategy = toolSelectionStrategy,
-    llmModel = llmModel,
-    llmParams = llmParams,
-    defineTask = defineTask
-)
+): AIAgentSubgraphDelegate<Input, CriticResult<Input>> = subgraph {
+    val inputKey = createStorageKey<Input>("subgraphWithVerification-input-key")
+
+    val saveInput by node<Input, Input> { input ->
+        storage.set(inputKey, input)
+
+        input
+    }
+
+    val verifyTask by subgraphWithTask<Input, CriticResultFromLLM>(
+        toolSelectionStrategy = toolSelectionStrategy,
+        llmModel = llmModel,
+        llmParams = llmParams,
+        defineTask = defineTask
+    )
+
+    val provideResult by node<CriticResultFromLLM, CriticResult<Input>> { result ->
+        CriticResult(
+            successful = result.isCorrect,
+            feedback = result.feedback,
+            input = storage.get(inputKey)!!
+        )
+    }
+
+    nodeStart then saveInput then verifyTask then provideResult then nodeFinish
+}
 
 /**
  * Constructs a subgraph within an AI agent's strategy graph with additional verification capabilities.
@@ -271,16 +280,16 @@ public inline fun <reified Input> AIAgentSubgraphBuilderBase<*, *>.subgraphWithV
  * @param defineTask A suspendable function defining the task that the subgraph will execute,
  *                   which takes an input and produces a string-based task description.
  * @return A delegate representing the constructed subgraph with input type `Input` and output type
- *         as a verified subgraph result `VerifiedSubgraphResult`.
+ *         as a verified subgraph result `CriticResult`.
  */
 @Suppress("unused")
 @AIAgentBuilderDslMarker
-public inline fun <reified Input> AIAgentSubgraphBuilderBase<*, *>.subgraphWithVerification(
+public inline fun <reified Input : Any> AIAgentSubgraphBuilderBase<*, *>.subgraphWithVerification(
     tools: List<Tool<*, *>>,
     llmModel: LLModel? = null,
     llmParams: LLMParams? = null,
     noinline defineTask: suspend AIAgentGraphContextBase.(input: Input) -> String
-): AIAgentSubgraphDelegate<Input, VerifiedSubgraphResult> = subgraphWithVerification(
+): AIAgentSubgraphDelegate<Input, CriticResult<Input>> = subgraphWithVerification(
     toolSelectionStrategy = ToolSelectionStrategy.Tools(tools.map { it.descriptor }),
     llmModel = llmModel,
     llmParams = llmParams,

agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/LLMAsAJudge.kt

@@ -0,0 +1,125 @@
+package ai.koog.agents.ext.agent
+
+import ai.koog.agents.core.annotation.InternalAgentsApi
+import ai.koog.agents.core.dsl.builder.AIAgentNodeDelegate
+import ai.koog.agents.core.dsl.builder.AIAgentSubgraphBuilderBase
+import ai.koog.agents.core.tools.annotations.LLMDescription
+import ai.koog.prompt.dsl.prompt
+import ai.koog.prompt.executor.clients.openai.OpenAIModels
+import ai.koog.prompt.llm.LLModel
+import ai.koog.prompt.message.Message
+import ai.koog.prompt.structure.StructureFixingParser
+import kotlinx.serialization.Serializable
+
+/**
+ * Represents the result of a plan evaluation performed by an LLM (Large Language Model).
+ *
+ * This class is primarily used within internal agent-related implementations where an LLM
+ * evaluates the correctness of a plan and optionally provides feedback for improvements.
+ *
+ * @property isCorrect Indicates whether the evaluated plan is correct.
+ * @property feedback Optional feedback provided by the LLM about the evaluated plan. This property
+ *        is populated only when the plan is deemed incorrect (`isCorrect == false`) and adjustments
+ *        are suggested.
+ */
+@InternalAgentsApi
+@Serializable
+@LLMDescription("Result of the evaluation")
+public data class CriticResultFromLLM(
+    @property:LLMDescription("Was the task solved correctly?")
+    val isCorrect: Boolean,
+    @property:LLMDescription(
+        "Optional feedback about the provided solution. " +
+            "Only needed if `isCorrect == false` and if solution needs adjustments."
+    )
+    val feedback: String
+)
+
+/**
+ * Represents the result of a critique or feedback process.
+ *
+ * @property successful Indicates whether the critique operation was successful.
+ * @property feedback A textual message providing details about the*/
+public data class CriticResult<T>(
+    val successful: Boolean,
+    val feedback: String,
+    val input: T
+)
+
+/**
+ * A method to utilize a language model (LLM) as a critic or judge for evaluating tasks with context-aware feedback.
+ * This method processes a given task and the interaction history to provide structured feedback on the task's correctness.
+ *
+ * @param llmModel The optional language model to override the default model during the session. If `null`, the default model will be used.
+ * @param task The task or instruction to be presented to the language model for critical evaluation.
+ */
+@OptIn(InternalAgentsApi::class)
+public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.llmAsAJudge(
+    llmModel: LLModel? = null,
+    task: String
+): AIAgentNodeDelegate<T, CriticResult<T>> = node<T, CriticResult<T>> { nodeInput ->
+    llm.writeSession {
+        val initialPrompt = prompt.copy()
+        val initialModel = model
+
+        prompt = prompt("critic") {
+            // Combine all history into one message with XML tags
+            // to prevent LLM from continuing answering in a tool_call -> tool_result pattern
+            val combinedMessage = buildString {
+                append("<previous_conversation>\n")
+                initialPrompt.messages.forEach { message ->
+                    when (message) {
+                        is Message.System -> append("<user>\n${message.content}\n</user>\n")
+                        is Message.User -> append("<user>\n${message.content}\n</user>\n")
+                        is Message.Assistant -> append("<assistant>\n${message.content}\n</assistant>\n")
+                        is Message.Tool.Call -> append(
+                            "<tool_call tool=${message.tool}>\n${message.content}\n</tool_call>\n"
+                        )
+
+                        is Message.Tool.Result -> append(
+                            "<tool_result tool=${message.tool}>\n${message.content}\n</tool_result>\n"
+                        )
+                    }
+                }
+                append("</previous_conversation>\n")
+            }
+
+            // Put Critic Task as a System instruction
+            system(task)
+            // And rest of the history -- in a combined XML message
+            user(combinedMessage)
+        }
+
+        if (llmModel != null) {
+            model = llmModel
+        }
+
+        val result = requestLLMStructured<CriticResultFromLLM>(
+            // optional field -- recommented for LLM awareness and reliability of the output
+            examples = listOf(
+                CriticResultFromLLM(
+                    isCorrect = true,
+                    feedback = "All good"
+                ),
+                CriticResultFromLLM(
+                    isCorrect = false,
+                    feedback = "Following parts of the plan have problems: *, *, *. Please consider changing ..."
+                )
+            ),
+            // optional field -- recommented for reliability of the format
+            fixingParser = StructureFixingParser(
+                fixingModel = OpenAIModels.CostOptimized.GPT4oMini,
+                retries = 3,
+            )
+        ).getOrThrow().structure
+
+        prompt = initialPrompt
+        model = initialModel
+
+        CriticResult(
+            successful = result.isCorrect,
+            feedback = result.feedback,
+            input = nodeInput
+        )
+    }
+}