support media types

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

@@ -6,6 +6,7 @@ import ai.koog.agents.core.environment.SafeTool
 import ai.koog.agents.core.environment.toSafeResult
 import ai.koog.agents.core.tools.Tool
 import ai.koog.agents.core.tools.ToolResult
+import ai.koog.prompt.message.MediaContent
 import ai.koog.prompt.message.Message
 import kotlin.reflect.KClass
 
@@ -174,3 +175,20 @@ public infix fun <IncomingOutput, IntermediateOutput, OutgoingInput>
         .onCondition { signature -> block(signature) }
         .transformed { it.content }
 }
+
+/**
+ * Creates an edge that filters assistant messages based on a custom condition and provides access to media content.
+ *
+ * @param block A function that evaluates whether to accept an assistant message with media
+ */
+public infix fun <IncomingOutput, IntermediateOutput, OutgoingInput>
+        AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput, OutgoingInput>.onAssistantMessageWithMedia(
+    block: suspend (Message.Assistant) -> Boolean
+): AIAgentEdgeBuilderIntermediate<IncomingOutput, MediaContent, OutgoingInput> {
+    return onIsInstance(Message.Assistant::class)
+        .onCondition {
+            it.mediaContent != null
+        }
+        .onCondition { signature -> block(signature) }
+        .transformed { it.mediaContent!! }
+}

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

@@ -77,9 +77,9 @@ class EventHandlerTest {
             "OnBeforeNode (node: __start__, input: ${agentInput})",
             "OnAfterNode (node: __start__, input: ${agentInput}, output: ${agentInput})",
             "OnBeforeNode (node: test LLM call, input: Test LLM call prompt)",
-            "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), finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts))], tools: [])",
+            "OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=[]), 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=[])], tools: [])",
             "OnAfterLLMCall (responses: [Assistant: Default test response])",
-            "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), finishReason=null))",
+            "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))",
             "OnStrategyFinished (strategy: $strategyName, result: Done)",
             "OnAgentFinished (strategy: $strategyName, result: Done)",
         )
@@ -123,9 +123,9 @@ class EventHandlerTest {
             "OnBeforeNode (node: __start__, input: ${agentInput})",
             "OnAfterNode (node: __start__, input: ${agentInput}, output: ${agentInput})",
             "OnBeforeNode (node: test LLM call, input: Test LLM call prompt)",
-            "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), finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts))], tools: [dummy])",
+            "OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=[]), 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=[])], tools: [dummy])",
             "OnAfterLLMCall (responses: [Assistant: Default test response])",
-            "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), finishReason=null))",
+            "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))",
             "OnStrategyFinished (strategy: $strategyName, result: Done)",
             "OnAgentFinished (strategy: $strategyName, result: Done)",
         )
@@ -171,13 +171,13 @@ class EventHandlerTest {
             "OnBeforeNode (node: __start__, input: ${agentInput})",
             "OnAfterNode (node: __start__, input: ${agentInput}, output: ${agentInput})",
             "OnBeforeNode (node: test LLM call, input: Test LLM call prompt)",
-            "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), finishReason=null), User(content=Test LLM call prompt, metaInfo=RequestMetaInfo(timestamp=$ts))], tools: [dummy])",
+            "OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=[]), 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=[])], tools: [dummy])",
             "OnAfterLLMCall (responses: [Assistant: Default test response])",
-            "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), finishReason=null))",
+            "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))",
             "OnBeforeNode (node: test LLM call with tools, input: Test LLM call with tools prompt)",
-            "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), 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), finishReason=null), User(content=Test LLM call with tools prompt, metaInfo=RequestMetaInfo(timestamp=$ts))], tools: [dummy])",
+            "OnBeforeLLMCall (prompt: [System(content=Test system message, metaInfo=RequestMetaInfo(timestamp=$ts)), User(content=Test user message, metaInfo=RequestMetaInfo(timestamp=$ts), mediaContent=[]), 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=[]), 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=[])], tools: [dummy])",
             "OnAfterLLMCall (responses: [Assistant: Default test response])",
-            "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), finishReason=null))",
+            "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))",
             "OnStrategyFinished (strategy: $strategyName, result: Done)",
             "OnAgentFinished (strategy: $strategyName, result: Done)",
         )

agents/agents-features/agents-features-trace/src/jvmTest/kotlin/ai/koog/agents/features/tracing/writer/TraceTestAPI.kt

@@ -4,21 +4,68 @@ import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.agents.core.agent.AIAgent
 import ai.koog.agents.core.agent.config.AIAgentConfig
 import ai.koog.agents.core.agent.entity.AIAgentStrategy
+import ai.koog.prompt.dsl.AttachmentBuilder
 import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.executor.clients.openai.OpenAIModels
+import ai.koog.prompt.message.MediaContent
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.message.RequestMetaInfo
 import ai.koog.prompt.message.ResponseMetaInfo
 import kotlinx.datetime.Clock
 
-val testClock: Clock = object : Clock {
+internal val testClock: Clock = object : Clock {
     override fun now(): kotlinx.datetime.Instant = kotlinx.datetime.Instant.parse("2023-01-01T00:00:00Z")
 }
 
-fun userMessage(content: String): Message.User = Message.User(content, metaInfo = RequestMetaInfo.create(testClock))
-fun assistantMessage(content: String): Message.Assistant = Message.Assistant(content, metaInfo = ResponseMetaInfo.create(testClock))
-fun systemMessage(content: String): Message.System = Message.System(content, metaInfo = RequestMetaInfo.create(testClock))
+/**
+ * Creates a user message with optional media attachments.
+ *
+ * The method constructs a user message using the provided text content and any additional
+ * media content defined via the `attachmentsBlock`. This allows the user to include
+ * various types of media attachments, such as images, audio files, or documents, alongside
+ * the message content.
+ *
+ * @param content The text content of the user message.
+ * @param attachmentsBlock A lambda function used to configure the media attachments
+ *                          for the message, using the `AttachmentBuilder` DSL.
+ * @return A `Message.User` object containing the message content, metadata,
+ *         and any associated media attachments.
+ */
+fun userMessage(content: String, attachmentsBlock: AttachmentBuilder.() -> Unit = {}): Message.User = Message.User(
+    content,
+    metaInfo = RequestMetaInfo.create(testClock),
+    mediaContent = AttachmentBuilder().apply(attachmentsBlock).build()
+)
 
+/**
+ * Creates an instance of [Message.Assistant] with the provided content and generated metadata.
+ *
+ * @param content The textual content of the assistant's message.
+ * @return A new instance of [Message.Assistant] containing the given content and metadata generated using the test clock.
+ */
+fun assistantMessage(content: String): Message.Assistant =
+    Message.Assistant(content, metaInfo = ResponseMetaInfo.create(testClock))
+
+/**
+ * Creates a system-generated message encapsulated in a [Message.System] instance.
+ *
+ * @param content The textual content of the system message.
+ * @return A [Message.System] object containing the provided content and autogenerated metadata.
+ */
+fun systemMessage(content: String): Message.System =
+    Message.System(content, metaInfo = RequestMetaInfo.create(testClock))
+
+/**
+ * Creates an AI agent with the specified configuration, strategy, and optional prompts.
+ *
+ * @param strategy The strategy used to define the workflow and execution pattern for the AI agent.
+ * @param promptId The identifier for the prompt configuration. If null, a default prompt ID will be used.
+ * @param systemPrompt Optional system-level message to include in the prompt. If null, a default message will be used.
+ * @param userPrompt Optional user-level message to include in the prompt. If null, a default message will be used.
+ * @param assistantPrompt Optional assistant response to include in the prompt. If null, a default response will be used.
+ * @param installFeatures Lambda function allowing additional features to be installed on the agent.
+ * @return A configured instance of the AIAgent class ready for execution.
+ */
 fun createAgent(
     strategy: AIAgentStrategy,
     promptId: String? = null,

examples/build.gradle.kts

@@ -71,6 +71,7 @@ registerRunExampleTask("runExampleMarkdownStreamingWithTool", "ai.koog.agents.ex
 registerRunExampleTask("runExampleRiderProjectTemplate", "ai.koog.agents.example.rider.project.template.RiderProjectTemplateKt")
 registerRunExampleTask("runExampleExecSandbox", "ai.koog.agents.example.execsandbox.ExecSandboxKt")
 registerRunExampleTask("runExampleLoopComponent", "ai.koog.agents.example.components.loop.ProjectGeneratorKt")
+registerRunExampleTask("runExampleInstagramPostDescriber", "ai.koog.agents.example.media.InstagramPostDescriberKt")
 
 dokka {
     dokkaSourceSets.named("main") {

examples/src/main/kotlin/ai/koog/agents/example/ApiKeyService.kt

@@ -6,4 +6,7 @@ internal object ApiKeyService {
 
     val anthropicApiKey: String
         get() = System.getenv("ANTHROPIC_API_KEY") ?: throw IllegalArgumentException("ANTHROPIC_API_KEY env is not set")
+
+    val googleApiKey: String
+        get() = System.getenv("GOOGLE_API_KEY") ?: throw IllegalArgumentException("GOOGLE_API_KEY env is not set")
 }

examples/src/main/kotlin/ai/koog/agents/example/media/InstagramPostDescriber.kt

@@ -0,0 +1,54 @@
+package ai.koog.agents.example.media
+
+import ai.koog.agents.example.ApiKeyService
+import ai.koog.prompt.dsl.prompt
+import ai.koog.prompt.executor.clients.anthropic.AnthropicModels
+import ai.koog.prompt.executor.clients.google.GoogleModels
+import ai.koog.prompt.executor.model.PromptExecutorExt.execute
+import ai.koog.prompt.executor.clients.openai.OpenAIModels
+import ai.koog.prompt.executor.llms.all.simpleAnthropicExecutor
+import ai.koog.prompt.executor.llms.all.simpleGoogleAIExecutor
+import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor
+import ai.koog.prompt.markdown.markdown
+import kotlinx.coroutines.runBlocking
+
+fun main() {
+    val openaiExecutor = simpleOpenAIExecutor(ApiKeyService.openAIApiKey)
+//    val anthropicExecutor = simpleAnthropicExecutor(ApiKeyService.anthropicApiKey)
+//    val googleExecutor = simpleGoogleAIExecutor(ApiKeyService.googleApiKey)
+
+    val resourcePath =
+        object {}.javaClass.classLoader.getResource("images")?.path ?: error("images directory not found")
+
+    val prompt = prompt("example-prompt") {
+        system("You are professional assistant that can write cool and funny descriptions for Instagram posts.")
+
+        user {
+            markdown {
+                +"I want to create a new post on Instagram."
+                br()
+                +"Can you write something creative under my instagram post with the following photos?"
+                br()
+                h2("Requirements")
+                bulleted {
+                    item("It must be very funny and creative")
+                    item("It must increase my chance of becoming an ultra-famous blogger!!!!")
+                    item("It not contain explicit content, harassment or bullying")
+                    item("It must be a short catching phrase")
+                    item("You must include relevant hashtags that would increase the visibility of my post")
+                }
+            }
+
+            attachments {
+                image("$resourcePath/photo1.png")
+                image("$resourcePath/photo2.png")
+            }
+        }
+    }
+
+    runBlocking {
+        openaiExecutor.execute(prompt, OpenAIModels.Chat.GPT4_1).content.also(::println)
+//        anthropicExecutor.execute(prompt, AnthropicModels.Sonnet_4).content.also(::println)
+//        googleExecutor.execute(prompt, GoogleModels.Gemini2_0Flash).content.also(::println)
+    }
+}

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/Models.kt

@@ -23,19 +23,24 @@ object Models {
             OpenAIModels.CostOptimized.O4Mini,
             OpenAIModels.CostOptimized.GPT4_1Nano,
             OpenAIModels.CostOptimized.GPT4_1Mini,
+
+            OpenAIModels.Audio.GPT4oMiniAudio,
+            OpenAIModels.Audio.GPT4oAudio,
         )
     }
 
     @JvmStatic
     fun anthropicModels(): Stream<LLModel> {
         return Stream.of(
-            AnthropicModels.Opus,
+            AnthropicModels.Opus_3,
+            AnthropicModels.Opus_4,
 
             AnthropicModels.Haiku_3,
             AnthropicModels.Haiku_3_5,
 
             AnthropicModels.Sonnet_3_5,
             AnthropicModels.Sonnet_3_7,
+            AnthropicModels.Sonnet_4,
         )
     }