Skip to content

Commit 6375ce0

Browse files
devcrocoddevcrocod
committed
update user dsl for prompting
1 parent b4919a6 commit 6375ce0

File tree

12 files changed

+647
-339
lines changed

12 files changed

+647
-339
lines changed

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentSubgraph.kt

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,9 +122,7 @@ public open class AIAgentSubgraph<Input, Output>(
122122
replaceHistoryWithTLDR()
123123
updatePrompt {
124124
user {
125-
text {
126-
selectRelevantTools(tools, toolSelectionStrategy.subtaskDescription)
127-
}
125+
selectRelevantTools(tools, toolSelectionStrategy.subtaskDescription)
128126
}
129127
}
130128

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMWriteSession.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -444,7 +444,7 @@ public class AIAgentLLMWriteSession internal constructor(
444444
if (definition != null) {
445445
val prompt = prompt(prompt, clock) {
446446
user {
447-
text { definition.definition(this) }
447+
definition.definition(this)
448448
}
449449
}
450450
this.prompt = prompt

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/HistoryCompressionStrategies.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ public abstract class HistoryCompressionStrategy {
4444
prompt = prompt.withMessages { messages -> messages.dropLastWhile { it is Message.Tool.Call } }
4545
updatePrompt {
4646
user {
47-
text { summarizeInTLDR() }
47+
summarizeInTLDR()
4848
}
4949
}
5050
listOf(llmSession.requestLLMWithoutTools())

‎agents/agents-features/agents-features-event-handler/src/jvmTest/kotlin/ai/koog/agents/features/eventHandler/feature/EventHandlerFeatureTest.kt

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ class EventHandlerTest {
7777
"OnBeforeNode (node: __start__, input: ${agentInput})",
7878
"OnAfterNode (node: __start__, input: ${agentInput}, output: ${agentInput})",
7979
"OnBeforeNode (node: test LLM call, input: Test LLM call prompt)",
80-
"OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts)), Assistant(content=Test assistant response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts))], tools: [])",
80+
"OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null), Assistant(content=Test assistant response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null)], tools: [])",
8181
"OnAfterLLMCall (responses: [Assistant: Default test response])",
8282
"OnAfterNode (node: test LLM call, input: Test LLM call prompt, output: Assistant(content=Default test response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null))",
8383
"OnStrategyFinished (strategy: $strategyName, result: Done)",
@@ -123,7 +123,7 @@ class EventHandlerTest {
123123
"OnBeforeNode (node: __start__, input: ${agentInput})",
124124
"OnAfterNode (node: __start__, input: ${agentInput}, output: ${agentInput})",
125125
"OnBeforeNode (node: test LLM call, input: Test LLM call prompt)",
126-
"OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts)), Assistant(content=Test assistant response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts))], tools: [dummy])",
126+
"OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null), Assistant(content=Test assistant response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null)], tools: [dummy])",
127127
"OnAfterLLMCall (responses: [Assistant: Default test response])",
128128
"OnAfterNode (node: test LLM call, input: Test LLM call prompt, output: Assistant(content=Default test response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null))",
129129
"OnStrategyFinished (strategy: $strategyName, result: Done)",
@@ -171,11 +171,11 @@ class EventHandlerTest {
171171
"OnBeforeNode (node: __start__, input: ${agentInput})",
172172
"OnAfterNode (node: __start__, input: ${agentInput}, output: ${agentInput})",
173173
"OnBeforeNode (node: test LLM call, input: Test LLM call prompt)",
174-
"OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts)), Assistant(content=Test assistant response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts))], tools: [dummy])",
174+
"OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null), Assistant(content=Test assistant response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null)], tools: [dummy])",
175175
"OnAfterLLMCall (responses: [Assistant: Default test response])",
176176
"OnAfterNode (node: test LLM call, input: Test LLM call prompt, output: Assistant(content=Default test response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null))",
177177
"OnBeforeNode (node: test LLM call with tools, input: Test LLM call with tools prompt)",
178-
"OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts)), Assistant(content=Test assistant response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts)), Assistant(content=Default test response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, mediaContent=null, finishReason=null), User(content=Test LLM call with tools prompt, metaInfo=RequestMetaInfo(timestamp=$ts))], tools: [dummy])",
178+
"OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null), Assistant(content=Test assistant response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null), Assistant(content=Default test response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null), User(content=Test LLM call with tools prompt, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=null)], tools: [dummy])",
179179
"OnAfterLLMCall (responses: [Assistant: Default test response])",
180180
"OnAfterNode (node: test LLM call with tools, input: Test LLM call with tools prompt, output: Assistant(content=Default test response, metaInfo=ResponseMetaInfo(timestamp=$ts, totalTokensCount=null, inputTokensCount=null, outputTokensCount=null), mediaContent=null, finishReason=null))",
181181
"OnStrategyFinished (strategy: $strategyName, result: Done)",

‎prompt/prompt-model/src/commonMain/kotlin/ai/koog/prompt/dsl/AttachmentBuilder.kt

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
package ai.koog.prompt.dsl
2+
3+
import ai.koog.prompt.message.MediaContent
4+
5+
/**
6+
* A builder for constructing media attachments for prompt messages.
7+
*
8+
* This builder provides a fluent DSL API for creating collections of media content
9+
* that can be attached to user messages. It supports images, audio files, and documents,
10+
* enabling rich multimedia content in prompt construction.
11+
*
12+
* Example usage:
13+
* ```kotlin
14+
* val attachments = AttachmentBuilder().apply {
15+
* image("screenshot.png")
16+
* audio(audioData, "mp3")
17+
* document("report.pdf")
18+
* }.build()
19+
* ```
20+
*
21+
* This class is part of the new DSL structure for prompt construction. It focuses specifically
22+
* on building media attachments, while text content is handled by TextContentBuilder.
23+
* For combining both text and attachments, use ContentBuilderWithAttachment.
24+
*
25+
* @see MediaContent for the types of media content supported
26+
* @see ContentBuilderWithAttachment for using this builder in prompt construction
27+
*/
28+
@PromptDSL
29+
public class AttachmentBuilder() {
30+
/**
31+
* Internal collection to accumulate media content during the building process.
32+
*
33+
* This mutable list stores all the media content items created through the various builder methods
34+
* and is used to construct the final media content list when [build] is called.
35+
*/
36+
private val mediaContents = mutableListOf<MediaContent>()
37+
38+
/**
39+
* Adds an image attachment to the media content collection.
40+
*
41+
* Creates an image media content item from the specified source.
42+
* The source can be either a local file path or a URL.
43+
*
44+
* Model type support:
45+
* - **Anthropic** — supports local and URL images. Formats: `png`, `jpeg`, `webp`, `gif`
46+
* - **Gemini** — supports local images only. Formats: `png`, `jpeg`, `webp`, `heic`, `heif`, `gif`
47+
* - **Ollama** — does not support images yet
48+
* - **OpenAI** — supports local and URL images. Formats: `png`, `jpeg`, `webp`, `gif`
49+
* - **OpenRouter** — supports local and URL images. Formats: `png`, `jpeg`, `webp`, `gif`
50+
*
51+
* Example:
52+
* ```kotlin
53+
* image("screenshot.png") // Local file
54+
* image("https://example.com/pic.jpg") // URL
55+
* ```
56+
*
57+
* @param source The path to the local image file or URL of the image
58+
*/
59+
public fun image(source: String) {
60+
mediaContents.add(MediaContent.Image(source))
61+
}
62+
63+
/**
64+
* Adds an audio attachment to the media content collection.
65+
*
66+
* Creates an audio media content item with the specified data and format.
67+
* The audio data should be provided as a byte array.
68+
*
69+
* - **Anthropic** — does not support audio
70+
* - **Gemini** — formats: `wav`, `mp3`, `aiff`, `aac`, `ogg`, `flac`
71+
* - **Ollama** — does not support audio
72+
* - **OpenAI** — formats: `wav`, `mp3`
73+
* - **OpenRouter** — formats: `wav`, `mp3`
74+
*
75+
* Example:
76+
* ```kotlin
77+
* audio(audioByteArray, "mp3")
78+
* audio(recordedSpeech, "wav")
79+
* ```
80+
*
81+
* @param data The audio data as a byte array
82+
* @param format The audio file format (e.g., "mp3", "wav", "ogg")
83+
*/
84+
public fun audio(data: ByteArray, format: String) {
85+
mediaContents.add(MediaContent.Audio(data, format))
86+
}
87+
88+
/**
89+
* Adds a document attachment to the media content collection.
90+
*
91+
* Creates a document media content item from the specified local file path.
92+
* URLs are not supported for security reasons.
93+
*
94+
* - Anthropic — supports local and URL-based PDF files. Local support also for TXT and MD files.
95+
* - Gemini — supports only local files of the following formats: `pdf`, `js`, `py`, `txt`, `html`, `css`, `md`, `csv`, `xml`, `rtf`
96+
* - Ollama — does not support documents yet
97+
* - OpenAI — supports only local PDF files
98+
* - OpenRouter — supports only local PDF files
99+
*
100+
* Example:
101+
* ```kotlin
102+
* document("report.pdf")
103+
* document("/path/to/document.docx")
104+
* ```
105+
*
106+
* @param source The local file path to the document
107+
*/
108+
public fun document(source: String) {
109+
mediaContents.add(MediaContent.File(source))
110+
}
111+
112+
/**
113+
* Constructs and returns the accumulated list of media content items.
114+
*
115+
* This method finalizes the building process and returns all the media content
116+
* items that were added through the various builder methods.
117+
*
118+
* @return A list containing all the media content items created through the builder methods
119+
*/
120+
public fun build(): List<MediaContent> = mediaContents
121+
}

‎prompt/prompt-model/src/commonMain/kotlin/ai/koog/prompt/dsl/ContentBuilderWithAttachment.kt

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
package ai.koog.prompt.dsl
2+
3+
import ai.koog.prompt.message.MediaContent
4+
import ai.koog.prompt.text.TextContentBuilder
5+
6+
/**
7+
* A builder class that extends TextContentBuilder to support both text and media attachments.
8+
*
9+
* This class combines text content building capabilities with media attachment support,
10+
* allowing users to create rich content that includes both textual information and
11+
* media elements like images, audio files, and documents.
12+
*
13+
* Example usage:
14+
* ```kotlin
15+
* val contentBuilder = ContentBuilderWithAttachment()
16+
* contentBuilder.text("Here's my analysis:")
17+
* contentBuilder.attachments {
18+
* image("chart.png")
19+
* document("report.pdf")
20+
* }
21+
* val (text, attachments) = contentBuilder.buildWithAttachments()
22+
* ```
23+
*
24+
* This class is part of the new DSL structure for prompt construction, replacing the previous
25+
* UserContentBuilder approach. It provides a more structured way to combine text and media
26+
* attachments in a single builder.
27+
*
28+
* @see TextContentBuilder for text-only content building
29+
* @see AttachmentBuilder for media attachment building
30+
*/
31+
@PromptDSL
32+
public class ContentBuilderWithAttachment : TextContentBuilder() {
33+
private var attachments: List<MediaContent> = emptyList()
34+
35+
/**
36+
* Configures media attachments for this content builder.
37+
*
38+
* This method allows you to specify multiple media attachments using the
39+
* AttachmentBuilder DSL. The attachments can include images, audio files,
40+
* and documents.
41+
*
42+
* Example:
43+
* ```kotlin
44+
* attachments {
45+
* image("photo.jpg")
46+
* audio(audioData, "mp3")
47+
* document("report.pdf")
48+
* }
49+
* ```
50+
*
51+
* @param body The configuration block for building attachments using AttachmentBuilder
52+
*/
53+
public fun attachments(body: AttachmentBuilder.() -> Unit) {
54+
this.attachments = AttachmentBuilder().apply(body).build()
55+
}
56+
57+
/**
58+
* Builds and returns both the text content and media attachments.
59+
*
60+
* This method combines the text content built through the inherited TextContentBuilder
61+
* methods with any media attachments configured through the [attachments] method.
62+
*
63+
* @return A Pair containing the built text content as the first element and
64+
* the list of media attachments as the second element
65+
*/
66+
public fun buildWithAttachments(): Pair<String, List<MediaContent>> = build() to attachments
67+
}

‎prompt/prompt-model/src/commonMain/kotlin/ai/koog/prompt/dsl/PromptBuilder.kt

Lines changed: 46 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package ai.koog.prompt.dsl
22

3+
import ai.koog.prompt.message.MediaContent
34
import ai.koog.prompt.message.Message
45
import ai.koog.prompt.message.RequestMetaInfo
56
import ai.koog.prompt.message.ResponseMetaInfo
@@ -80,39 +81,73 @@ public class PromptBuilder internal constructor(
8081
}
8182

8283
/**
83-
* Adds a user message to the prompt.
84+
* Adds a user message to the prompt with optional media attachments.
8485
*
8586
* User messages represent input from the user to the language model.
87+
* This method supports adding text content along with a list of media attachments such as images, audio, or documents.
88+
*
89+
* @param content The content of the user message.
90+
* @param attachments The list of media attachments associated with the user message. Defaults to an empty list if no attachments are provided.
91+
*/
92+
public fun user(content: String, attachments: List<MediaContent> = emptyList()) {
93+
when (attachments.size) {
94+
0 -> messages.add(Message.User(content, RequestMetaInfo.create(clock)))
95+
1 -> messages.add(Message.User(content, RequestMetaInfo.create(clock), attachments.first()))
96+
else -> with(messages) {
97+
add(Message.User(content, RequestMetaInfo.create(clock), attachments.first()))
98+
addAll(
99+
attachments
100+
.subList(1, attachments.lastIndex + 1)
101+
.map { Message.User("", RequestMetaInfo.create(clock), it) })
102+
}
103+
}
104+
}
105+
106+
/**
107+
* Adds a user message to the prompt with media attachments.
108+
*
109+
* User messages represent input from the user to the language model.
110+
* This method allows attaching media content like images, audio, or documents.
86111
*
87112
* Example:
88113
* ```kotlin
114+
* // Simple text message
89115
* user("What is the capital of France?")
116+
*
117+
* // Message with attachments using a lambda
118+
* user("Please analyze this image") {
119+
* image("photo.jpg")
120+
* }
90121
* ```
91122
*
92123
* @param content The content of the user message
124+
* @param block Optional lambda to configure attachments using AttachmentBuilder
93125
*/
94-
public fun user(content: String) {
95-
messages.add(Message.User(content, RequestMetaInfo.create(clock)))
126+
public fun user(content: String, block: AttachmentBuilder.() -> Unit) {
127+
user(content, AttachmentBuilder().apply(block).build())
96128
}
97129

98130
/**
99-
* Adds a user message to the prompt using a UserContentBuilder.
131+
* Adds a user message to the prompt using a ContentBuilderWithAttachment.
100132
*
101-
* This allows for complex user message construction, supporting features like textual
102-
* and media content.
133+
* This allows for more complex message construction with both text and attachments.
103134
*
104135
* Example:
105136
* ```kotlin
106137
* user {
107-
* text("What is in this image?")
108-
* image("https://example.com/image.jpg")
138+
* text("I have a question about programming.")
139+
* text("How do I implement a binary search in Kotlin?")
140+
* attachments {
141+
* image("screenshot.png")
142+
* }
109143
* }
110144
* ```
111145
*
112-
* @param body The initialization block for the UserContentBuilder.
146+
* @param body The initialization block for the ContentBuilderWithAttachment
113147
*/
114-
public fun user(body: UserContentBuilder.() -> Unit) {
115-
messages.addAll(UserContentBuilder(clock).apply(body).build())
148+
public fun user(body: ContentBuilderWithAttachment.() -> Unit) {
149+
val (content, media) = ContentBuilderWithAttachment().apply(body).buildWithAttachments()
150+
user(content, media)
116151
}
117152

118153
/**

0 commit comments

Comments
 (0)