AWS Bedrock support (#285)

Author: kaeawc

examples/build.gradle.kts

@@ -29,6 +29,7 @@ dependencies {
     api(project(":prompt:prompt-structure"))
     api(project(":prompt:prompt-executor:prompt-executor-clients:prompt-executor-openai-client"))
     api(project(":prompt:prompt-executor:prompt-executor-clients:prompt-executor-anthropic-client"))
+    api(project(":prompt:prompt-executor:prompt-executor-clients:prompt-executor-bedrock-client"))
     api(project(":prompt:prompt-executor:prompt-executor-llms"))
     api(project(":prompt:prompt-executor:prompt-executor-llms-all"))
 
@@ -80,6 +81,7 @@ registerRunExampleTask("runExampleInstagramPostDescriber", "ai.koog.agents.examp
 registerRunExampleTask("runExampleRoutingViaGraph", "ai.koog.agents.example.banking.routing.RoutingViaGraphKt")
 registerRunExampleTask("runExampleRoutingViaAgentsAsTools", "ai.koog.agents.example.banking.routing.RoutingViaAgentsAsToolsKt")
 registerRunExampleTask("runExampleFeatureOpenTelemetry", "ai.koog.agents.example.feature.OpenTelemetryKt")
+registerRunExampleTask("runExampleBedrockAgent", "ai.koog.agents.example.client.BedrockAgentKt")
 
 dokka {
     dokkaSourceSets.named("main") {

examples/env.template.properties

@@ -6,4 +6,7 @@
 
 OPENAI_API_KEY=
 
-ANTHROPIC_API_KEY=
+ANTHROPIC_API_KEY=
+
+AWS_ACCESS_KEY_ID=
+AWS_SECRET_ACCESS_KEY=

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

@@ -12,4 +12,10 @@ internal object ApiKeyService {
 
     val openRouterApiKey: String
         get() = System.getenv("OPENROUTER_API_KEY") ?: throw IllegalArgumentException("OPENROUTER_API_KEY env is not set")
+
+    val awsAccessKey: String
+        get() = System.getenv("AWS_ACCESS_KEY_ID") ?: throw IllegalArgumentException("AWS_ACCESS_KEY_ID env is not set")
+
+    val awsSecretAccessKey: String
+        get() = System.getenv("AWS_SECRET_ACCESS_KEY") ?: throw IllegalArgumentException("AWS_SECRET_ACCESS_KEY env is not set")
 }

examples/src/main/kotlin/ai/koog/agents/example/client/BedrockAgent.kt

@@ -0,0 +1,48 @@
+package ai.koog.agents.example.client
+
+import ai.koog.agents.core.agent.AIAgent
+import ai.koog.agents.core.tools.ToolRegistry
+import ai.koog.agents.core.tools.reflect.asTools
+import ai.koog.agents.example.ApiKeyService
+import ai.koog.agents.example.simpleapi.Switch
+import ai.koog.agents.example.simpleapi.SwitchTools
+import ai.koog.prompt.executor.clients.bedrock.BedrockClientSettings
+import ai.koog.prompt.executor.clients.bedrock.BedrockModels
+import ai.koog.prompt.executor.llms.all.simpleBedrockExecutor
+import kotlinx.coroutines.runBlocking
+
+fun main(): Unit = runBlocking {
+    val switch = Switch()
+
+    val toolRegistry = ToolRegistry {
+        tools(SwitchTools(switch).asTools())
+    }
+
+    // Create Bedrock client settings
+    val bedrockSettings = BedrockClientSettings(
+        region = "us-east-1", // Change this to your preferred region
+        maxRetries = 3
+    )
+
+    // Create Bedrock LLM client
+    val executor = simpleBedrockExecutor(
+        awsAccessKeyId = ApiKeyService.awsAccessKey,
+        awsSecretAccessKey = ApiKeyService.awsSecretAccessKey,
+        settings = bedrockSettings
+    )
+
+    val agent = AIAgent(
+        executor = executor,
+        llmModel = BedrockModels.AnthropicClaude35SonnetV2, // Use Claude 3.5 Sonnet
+        systemPrompt = "You're responsible for running a Switch and perform operations on it by request",
+        temperature = 0.0,
+        toolRegistry = toolRegistry
+    )
+
+    println("Bedrock Agent with Switch Tools - Chat started")
+    println("You can ask me to turn the switch on/off or check its current state.")
+    println("Type your request:")
+
+    val input = readln()
+    agent.run(input)
+}

gradle/libs.versions.toml

@@ -21,6 +21,7 @@ kover = "0.9.1"
 spring-boot = "3.5.3"
 spring-management = "1.1.7"
 opentelemetry = "1.51.0"
+aws-sdk-kotlin = "1.4.119"
 
 [libraries]
 jetbrains-annotations = { module = "org.jetbrains:annotations", version.ref = "annotations" }
@@ -55,6 +56,7 @@ opentelemetry-bom = { module = "io.opentelemetry:opentelemetry-bom", version.ref
 opentelemetry-sdk = { module = "io.opentelemetry:opentelemetry-sdk" }
 opentelemetry-exporter-otlp = { module = "io.opentelemetry:opentelemetry-exporter-otlp" }
 opentelemetry-exporter-logging = { module = "io.opentelemetry:opentelemetry-exporter-logging" }
+aws-sdk-kotlin-bedrockruntime = { module = "aws.sdk.kotlin:bedrockruntime", version.ref = "aws-sdk-kotlin" }
 
 # Spring
 spring-boot-bom = { module = "org.springframework.boot:spring-boot-dependencies", version.ref = "spring-boot" }

koog-agents/build.gradle.kts

@@ -39,6 +39,7 @@ val included = setOf(
     ":prompt:prompt-executor:prompt-executor-clients:prompt-executor-openai-client",
     ":prompt:prompt-executor:prompt-executor-clients:prompt-executor-openrouter-client",
     ":prompt:prompt-executor:prompt-executor-clients:prompt-executor-ollama-client",
+    ":prompt:prompt-executor:prompt-executor-clients:prompt-executor-bedrock-client",
     ":prompt:prompt-executor:prompt-executor-llms",
     ":prompt:prompt-executor:prompt-executor-llms-all",
     ":prompt:prompt-executor:prompt-executor-model",

prompt/prompt-executor/prompt-executor-clients/prompt-executor-anthropic-client/build.gradle.kts

@@ -54,6 +54,9 @@ kotlin {
     }
 
     explicitApi()
+    compilerOptions {
+        freeCompilerArgs.add("-opt-in=ai.koog.prompt.executor.clients.InternalLLMClientApi")
+    }
 }
 
 publishToMaven()

prompt/prompt-executor/prompt-executor-clients/prompt-executor-anthropic-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/anthropic/DataModel.kt

@@ -1,13 +1,15 @@
 package ai.koog.prompt.executor.clients.anthropic
 
+import ai.koog.prompt.executor.clients.InternalLLMClientApi
 import kotlinx.serialization.EncodeDefault
 import kotlinx.serialization.EncodeDefault.Mode.ALWAYS
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.JsonObject
 
+@InternalLLMClientApi
 @Serializable
-internal data class AnthropicMessageRequest(
+public data class AnthropicMessageRequest(
     val model: String,
     val messages: List<AnthropicMessage>,
     val maxTokens: Int = 2048,
@@ -18,90 +20,98 @@ internal data class AnthropicMessageRequest(
     val toolChoice: AnthropicToolChoice? = null,
 )
 
+@InternalLLMClientApi
 @Serializable
-internal data class AnthropicMessage(
+public data class AnthropicMessage(
     val role: String,
     val content: List<AnthropicContent>
 )
 
+@InternalLLMClientApi
 @Serializable
-internal data class SystemAnthropicMessage(
+public data class SystemAnthropicMessage(
     val text: String,
     @EncodeDefault(ALWAYS) val type: String = "text"
 )
 
+@InternalLLMClientApi
 @Serializable
-internal sealed class AnthropicContent {
+public sealed class AnthropicContent {
     @Serializable
     @SerialName("text")
-    data class Text(val text: String) : AnthropicContent()
+    public data class Text(val text: String) : AnthropicContent()
 
     @Serializable
     @SerialName("image")
-    data class Image(val source: ImageSource) : AnthropicContent()
+    public data class Image(val source: ImageSource) : AnthropicContent()
 
     @Serializable
     @SerialName("document")
-    data class Document(val source: DocumentSource) : AnthropicContent()
+    public data class Document(val source: DocumentSource) : AnthropicContent()
 
     @Serializable
     @SerialName("tool_use")
-    data class ToolUse(
+    public data class ToolUse(
         val id: String,
         val name: String,
         val input: JsonObject
     ) : AnthropicContent()
 
     @Serializable
     @SerialName("tool_result")
-    data class ToolResult(
+    public data class ToolResult(
         val toolUseId: String,
         val content: String
     ) : AnthropicContent()
 }
 
+@InternalLLMClientApi
 @Serializable
-internal sealed interface ImageSource {
+public sealed interface ImageSource {
     @Serializable
     @SerialName("url")
-    data class Url(val url: String) : ImageSource
+    public data class Url(val url: String) : ImageSource
 
     @Serializable
     @SerialName("base64")
-    data class Base64(val data: String, val mediaType: String) : ImageSource
+    public data class Base64(val data: String, val mediaType: String) : ImageSource
 }
 
+@InternalLLMClientApi
 @Serializable
-internal sealed interface DocumentSource {
+public sealed interface DocumentSource {
     @Serializable
     @SerialName("url")
-    data class Url(val url: String) : DocumentSource
+    public data class Url(val url: String) : DocumentSource
 
     @Serializable
     @SerialName("base64")
-    data class Base64(val data: String, val mediaType: String) : DocumentSource
+    public data class Base64(val data: String, val mediaType: String) : DocumentSource
 
     @Serializable
     @SerialName("text")
-    data class PlainText(val data: String, val mediaType: String) : DocumentSource
+    public data class PlainText(val data: String, val mediaType: String) : DocumentSource
 }
 
+@InternalLLMClientApi
 @Serializable
-internal data class AnthropicTool(
+public data class AnthropicTool(
     val name: String,
     val description: String,
     val inputSchema: AnthropicToolSchema
 )
 
+@InternalLLMClientApi
 @Serializable
-internal data class AnthropicToolSchema(
+public data class AnthropicToolSchema(
     val type: String = "object",
     val properties: JsonObject,
     val required: List<String>
 )
 
+@InternalLLMClientApi
 @Serializable
-internal data class AnthropicResponse(
+public data class AnthropicResponse(
     val id: String,
     val type: String,
     val role: String,
@@ -111,57 +121,62 @@ internal data class AnthropicResponse(
     val usage: AnthropicUsage? = null
 )
 
+@InternalLLMClientApi
 @Serializable
-internal sealed class AnthropicResponseContent {
+public sealed class AnthropicResponseContent {
     @Serializable
     @SerialName("text")
-    data class Text(val text: String) : AnthropicResponseContent()
+    public data class Text(val text: String) : AnthropicResponseContent()
 
     @Serializable
     @SerialName("tool_use")
-    data class ToolUse(
+    public data class ToolUse(
         val id: String,
         val name: String,
         val input: JsonObject
     ) : AnthropicResponseContent()
 }
 
+@InternalLLMClientApi
 @Serializable
-internal data class AnthropicUsage(
+public data class AnthropicUsage(
     val inputTokens: Int,
     val outputTokens: Int
 )
 
+@InternalLLMClientApi
 @Serializable
-internal data class AnthropicStreamResponse(
+public data class AnthropicStreamResponse(
     val type: String,
     val delta: AnthropicStreamDelta? = null,
     val message: AnthropicResponse? = null
 )
 
+@InternalLLMClientApi
 @Serializable
-internal data class AnthropicStreamDelta(
+public data class AnthropicStreamDelta(
     val type: String,
     val text: String? = null,
     val toolUse: AnthropicResponseContent.ToolUse? = null
 )
 
 
+@InternalLLMClientApi
 @Serializable
-internal sealed interface AnthropicToolChoice {
+public sealed interface AnthropicToolChoice {
     @Serializable
     @SerialName("auto")
-    data object Auto : AnthropicToolChoice
+    public data object Auto : AnthropicToolChoice
 
     @Serializable
     @SerialName("any")
-    data object Any : AnthropicToolChoice
+    public data object Any : AnthropicToolChoice
 
     @Serializable
     @SerialName("none")
-    data object None : AnthropicToolChoice
+    public data object None : AnthropicToolChoice
 
     @Serializable
     @SerialName("tool")
-    data class Tool(val name: String) : AnthropicToolChoice
+    public data class Tool(val name: String) : AnthropicToolChoice
 }