Skip to content

Commit 6d04ba2

Browse files
authored
Add Bedrock support in Ktor for configuring and initializing Bedrock LLM clients. (#1141)
Fixes #803 ## Motivation and Context Adds missing support for Amazon Bedrock in Ktor integration ## Breaking Changes No --- #### Type of the changes - [x] New feature (non-breaking change which adds functionality) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Tests improvement - [ ] Refactoring #### Checklist - [x] The pull request has a description of the proposed change - [x] I read the [Contributing Guidelines](https://github.com/JetBrains/koog/blob/main/CONTRIBUTING.md) before opening the pull request - [x] The pull request uses **`develop`** as the base branch - [ ] Tests for the changes have been added - [x] All new and existing tests passed ##### Additional steps for pull requests adding a new feature - [x] An issue describing the proposed change exists - [x] The pull request includes a link to the issue - [ ] The change was discussed and approved in the issue - [x] Docs have been added / updated
1 parent b4967d8 commit 6d04ba2

File tree

1 file changed

+75
-0
lines changed

1 file changed

+75
-0
lines changed
Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
package ai.koog.ktor
2+
3+
import ai.koog.prompt.executor.clients.LLMClient
4+
import ai.koog.prompt.executor.clients.bedrock.BedrockGuardrailsSettings
5+
import ai.koog.prompt.executor.clients.bedrock.BedrockLLMClient
6+
import ai.koog.prompt.executor.clients.bedrock.StaticBearerTokenProvider
7+
import ai.koog.prompt.llm.LLMProvider
8+
import aws.sdk.kotlin.services.bedrockruntime.BedrockRuntimeClient
9+
import kotlinx.datetime.Clock
10+
11+
/**
12+
* Configuration to create a new Bedrock LLM client configured with the specified identity provider and settings.
13+
*/
14+
public class BedrockClientConfig {
15+
16+
/**
17+
* Configure the underlying BedrockRuntimeClient from the AWS SDK.
18+
*/
19+
public var configure: BedrockRuntimeClient.Config.Builder.() -> Unit = {}
20+
21+
/**
22+
* Set moderation guardrails settings for Bedrock.
23+
*/
24+
public var moderationGuardrailsSettings: BedrockGuardrailsSettings? = null
25+
26+
/**
27+
* Override the clock used for time-based operations.
28+
*/
29+
public var clock: Clock = Clock.System
30+
}
31+
32+
/**
33+
* Configures and initializes a Bedrock LLM client with optional configuration.
34+
*
35+
* @param apiKey The API key used for authenticating with the Bedrock API.
36+
* @param clock A clock used for time-based operations
37+
* @param moderationGuardrailsSettings Optional settings of the AWS bedrock Guardrails (see [AWS documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-use-independent-api.html) ) that would be used for the [LLMClient.moderate] request
38+
* @param configure A lambda receiver to customize the OpenAI configuration such as base URL, timeout settings, and paths.
39+
*/
40+
public fun KoogAgentsConfig.bedrock(
41+
apiKey: String,
42+
clock: Clock = Clock.System,
43+
moderationGuardrailsSettings: BedrockGuardrailsSettings? = null,
44+
configure: BedrockRuntimeClient.Config.Builder.() -> Unit = {}
45+
) {
46+
val client = BedrockRuntimeClient {
47+
configure()
48+
bearerTokenProvider = StaticBearerTokenProvider(apiKey)
49+
}
50+
addLLMClient(
51+
LLMProvider.Bedrock,
52+
BedrockLLMClient(client, moderationGuardrailsSettings, clock)
53+
)
54+
}
55+
56+
/**
57+
* Configures and initializes a Bedrock LLM client with AWS SDK builder
58+
*
59+
* @param clock A clock used for time-based operations
60+
* @param moderationGuardrailsSettings Optional settings of the AWS bedrock Guardrails (see [AWS documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-use-independent-api.html) ) that would be used for the [LLMClient.moderate] request
61+
* @param configure A lambda receiver to customize the OpenAI configuration such as base URL, timeout settings, and paths.
62+
*/
63+
public fun KoogAgentsConfig.bedrock(
64+
clock: Clock = Clock.System,
65+
moderationGuardrailsSettings: BedrockGuardrailsSettings? = null,
66+
configure: BedrockRuntimeClient.Config.Builder.() -> Unit
67+
) {
68+
val client = BedrockRuntimeClient {
69+
configure()
70+
}
71+
addLLMClient(
72+
LLMProvider.Bedrock,
73+
BedrockLLMClient(client, moderationGuardrailsSettings, clock)
74+
)
75+
}

0 commit comments

Comments
 (0)