feat(prompt): Include examples in the prompt with StructuredRequest.Native to help LLMs better understand desired data format (#1396)

In certain cases, it's not enough to just constrain an LLM to a given
data format, since it might "get lost". It helps to provide an LLM with
examples regarding what type of data we want in the output, not only in
which format.

Fixes #1328

Author: EugeneTheDev

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/PromptExecutorExtensions.kt

@@ -77,7 +77,14 @@ public data class StructuredRequestConfig<T>(
 
             // Rely on built-in model capabilities to provide structured response.
             is StructuredRequest.Native -> {
-                prompt.withUpdatedParams { schema = mode.structure.schema }
+                prompt(prompt) {
+                    // If examples are supplied, append them
+                    if (mode.structure.examples.isNotEmpty()) {
+                        user {
+                            mode.structure.examples(this)
+                        }
+                    }
+                }.withUpdatedParams { schema = mode.structure.schema }
             }
         }
     }

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/Structure.kt

@@ -1,6 +1,7 @@
 package ai.koog.prompt.structure
 
 import ai.koog.prompt.params.LLMParams
+import ai.koog.prompt.text.TextContentBuilderBase
 
 /**
  * Represents a generic structure for handling data with a specific schema.
@@ -34,4 +35,12 @@ public abstract class Structure<TStruct, TSchema : LLMParams.Schema>(
      * @return A string representing the pretty-printed version of the input structured data.
      */
     public abstract fun pretty(value: TStruct): String
+
+    /**
+     * Defines formatted examples using the provided [TextContentBuilderBase].
+     *
+     * @param builder The [TextContentBuilderBase] instance for constructing textual content.
+     * @return The modified [TextContentBuilderBase] containing formatted examples.
+     */
+    public abstract fun examples(builder: TextContentBuilderBase<*>): TextContentBuilderBase<*>
 }

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredPrompts.kt

@@ -8,9 +8,7 @@ import ai.koog.prompt.text.TextContentBuilderBase
  */
 public object StructuredOutputPrompts {
     /**
-     * Formats and appends the structured data output to the provided MarkdownContentBuilder.
-     *
-     * @param structure The StructuredData instance containing the format ID and definition for the output.
+     * Formats and appends the structured data output to the provided text builder.
      */
     public fun outputInstructionPrompt(builder: TextContentBuilderBase<*>, structure: Structure<*, *>): TextContentBuilderBase<*> = builder.apply {
         markdown {
@@ -21,4 +19,30 @@ public object StructuredOutputPrompts {
             structure.definition(this)
         }
     }
+
+    /**
+     * Formats and appends structure examples, if they are present in the provided [structure], to the provided text builder,
+     * to show an LLM expected output format. If [Structure.examples] is empty, nothing is appended.
+     */
+    public fun <T> examplesPrompt(builder: TextContentBuilderBase<*>, structure: Structure<T, *>): TextContentBuilderBase<*> = builder.apply {
+        markdown {
+            if (structure.examples.isNotEmpty()) {
+                h4("EXAMPLES")
+
+                if (structure.examples.size == 1) {
+                    +"Here is an example of a valid response:"
+                } else {
+                    +"Here are some examples of valid responses:"
+                }
+
+                structure.examples.forEach { example ->
+                    codeblock(
+                        code = ai.koog.prompt.text.text {
+                            structure(structure, example)
+                        },
+                    )
+                }
+            }
+        }
+    }
 }

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructure.kt

@@ -3,9 +3,9 @@ package ai.koog.prompt.structure.json
 import ai.koog.prompt.markdown.markdown
 import ai.koog.prompt.params.LLMParams
 import ai.koog.prompt.structure.Structure
+import ai.koog.prompt.structure.StructuredOutputPrompts
 import ai.koog.prompt.structure.json.generator.JsonSchemaGenerator
 import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator
-import ai.koog.prompt.structure.structure
 import ai.koog.prompt.text.TextContentBuilderBase
 import kotlinx.serialization.KSerializer
 import kotlinx.serialization.json.ClassDiscriminatorMode
@@ -21,8 +21,10 @@ import kotlinx.serialization.serializer
  * @property examples A list of example data items that conform to the structure.
  * @property serializer The serializer used to convert the data to and from JSON.
  * @property json [kotlinx.serialization.json.Json] instance to perform de/serialization.
- * @property definitionPrompt Prompt with definition, explaining the structure to the LLM.
+ * @param definitionPrompt Prompt with definition, explaining the structure to the LLM.
  * Default is [JsonStructure.defaultDefinitionPrompt]
+ * @param examplesPrompt Prompt with examples of valid formats for the structured data.
+ * Default is [StructuredOutputPrompts.examplesPrompt]
 */
 public class JsonStructure<TStruct>(
     id: String,
@@ -32,8 +34,12 @@ public class JsonStructure<TStruct>(
     public val json: Json,
     private val definitionPrompt: (
         builder: TextContentBuilderBase<*>,
-        structuredData: JsonStructure<TStruct>
-    ) -> TextContentBuilderBase<*> = ::defaultDefinitionPrompt
+        structure: JsonStructure<TStruct>
+    ) -> TextContentBuilderBase<*> = ::defaultDefinitionPrompt,
+    private val examplesPrompt: (
+        builder: TextContentBuilderBase<*>,
+        structure: JsonStructure<TStruct>,
+    ) -> TextContentBuilderBase<*> = StructuredOutputPrompts::examplesPrompt,
 ) : Structure<TStruct, LLMParams.Schema.JSON>(id, schema, examples) {
 
     override fun parse(text: String): TStruct = json.decodeFromString(serializer, text.trim().stripMarkdown())
@@ -42,6 +48,8 @@ public class JsonStructure<TStruct>(
 
     override fun definition(builder: TextContentBuilderBase<*>): TextContentBuilderBase<*> = definitionPrompt(builder, this)
 
+    override fun examples(builder: TextContentBuilderBase<*>): TextContentBuilderBase<*> = examplesPrompt(builder, this)
+
     // LLMs often enclose JSON output in Markdown code blocks.
     // Stripping them saves extra LLM call which would be required to fix the string.
     private fun String.stripMarkdown() = when {
@@ -88,24 +96,7 @@ public class JsonStructure<TStruct>(
                     +json.encodeToString(schema.schema)
                     br()
 
-                    if (examples.isNotEmpty()) {
-                        h4("EXAMPLES")
-
-                        if (examples.size == 1) {
-                            +"Here is an example of a valid response:"
-                        } else {
-                            +"Here are some examples of valid responses:"
-                        }
-
-                        examples.forEach { example ->
-                            codeblock(
-                                code = ai.koog.prompt.text.text {
-                                    structure(this@with, example)
-                                },
-                                language = "json"
-                            )
-                        }
-                    }
+                    StructuredOutputPrompts.examplesPrompt(this, structuredData)
 
                     h2("RESULT")
                     +"Provide ONLY the resulting JSON, WITHOUT ANY free text comments, backticks, or other symbols."