Skip to content

Commit 2823045

Browse files
authored
Add support for ServiceTier (#61)
* Add support for ServiceTier * revert log level to debug in examples * rationalize the default value * fix test
1 parent 4e51e09 commit 2823045

23 files changed

+491
-51
lines changed

Examples/openai/Sources/Converse/main.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ let bedrock = try await BedrockService(
2626

2727
var builder = try ConverseRequestBuilder(with: .openai_gpt_oss_20b)
2828
.withPrompt("Who are you?")
29+
.withServiceTier(.priority)
2930

3031
var reply = try await bedrock.converse(with: builder)
3132

Examples/openai/Sources/Invoke/main.swift

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ let bedrock = try await BedrockService(
2626

2727
let textCompletion = try await bedrock.completeText(
2828
"Who are you?",
29-
with: .openai_gpt_oss_20b
29+
with: .openai_gpt_oss_20b,
30+
// serviceTier: .default
3031
)
3132

3233
if let reasoning = textCompletion.reasoning {

Package.swift

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,11 @@ let package = Package(
1010
.library(name: "BedrockService", targets: ["BedrockService"])
1111
],
1212
dependencies: [
13-
.package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.6.1"),
14-
.package(url: "https://github.com/awslabs/aws-sdk-swift", from: "1.5.51"),
15-
.package(url: "https://github.com/smithy-lang/smithy-swift", from: "0.158.0"),
13+
.package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.6.2"),
14+
.package(url: "https://github.com/awslabs/aws-sdk-swift", from: "1.6.3"),
15+
.package(url: "https://github.com/smithy-lang/smithy-swift", from: "0.173.0"),
1616
.package(url: "https://github.com/apple/swift-log.git", from: "1.6.4"),
17-
.package(url: "https://github.com/awslabs/aws-crt-swift", from: "0.53.0"),
17+
.package(url: "https://github.com/awslabs/aws-crt-swift", from: "0.54.2"),
1818
],
1919
targets: [
2020
.target(

Sources/BedrockService/BedrockRuntimeClient/Converse/BedrockService+Converse.swift

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,8 @@ extension BedrockService {
7171
systemPrompts: systemPrompts,
7272
tools: tools,
7373
enableReasoning: enableReasoning,
74-
maxReasoningTokens: maxReasoningTokens
74+
maxReasoningTokens: maxReasoningTokens,
75+
serviceTier: .default
7576
)
7677
}
7778

@@ -87,6 +88,7 @@ extension BedrockService {
8788
/// - tools: Optional array of tools the model can use
8889
/// - enableReasoning: Optional flag to enable reasoning output
8990
/// - maxReasoningTokens: Optional maximum number of reasoning tokens to generate
91+
/// - serviceTier: Optional. The service tier to serve this request (.default | .priority | .flex)
9092
/// - Throws: BedrockLibraryError.notSupported for parameters or functionalities that are not supported
9193
/// BedrockLibraryError.invalidParameter for invalid parameters
9294
/// BedrockLibraryError.invalidPrompt if the prompt is empty or too long
@@ -103,7 +105,8 @@ extension BedrockService {
103105
systemPrompts: [String]? = nil,
104106
tools: [Tool]? = nil,
105107
enableReasoning: Bool? = false,
106-
maxReasoningTokens: Int? = nil
108+
maxReasoningTokens: Int? = nil,
109+
serviceTier: ServiceTier
107110
) async throws -> Message {
108111
do {
109112
let modality = try model.getConverseModality()
@@ -127,6 +130,7 @@ extension BedrockService {
127130
"stopSequences": "\(String(describing: stopSequences))",
128131
"systemPrompts": "\(String(describing: systemPrompts))",
129132
"tools": "\(String(describing: tools))",
133+
"serviceTier": "\(serviceTier.rawValue)",
130134
]
131135
)
132136
let converseRequest = ConverseRequest(
@@ -138,7 +142,8 @@ extension BedrockService {
138142
stopSequences: stopSequences,
139143
systemPrompts: systemPrompts,
140144
tools: tools,
141-
maxReasoningTokens: maxReasoningTokens
145+
maxReasoningTokens: maxReasoningTokens,
146+
serviceTier: serviceTier
142147
)
143148

144149
logger.trace("Creating ConverseInput")
@@ -180,7 +185,8 @@ extension BedrockService {
180185
stopSequences: builder.stopSequences,
181186
systemPrompts: builder.systemPrompts,
182187
tools: builder.tools,
183-
maxReasoningTokens: builder.maxReasoningTokens
188+
maxReasoningTokens: builder.maxReasoningTokens,
189+
serviceTier: builder.serviceTier
184190
)
185191
history.append(assistantMessage)
186192
logger.trace(

Sources/BedrockService/BedrockRuntimeClient/Converse/BedrockService+ConverseStreaming.swift

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,8 @@ extension BedrockService {
7474
systemPrompts: systemPrompts,
7575
tools: tools,
7676
enableReasoning: enableReasoning,
77-
maxReasoningTokens: maxReasoningTokens
77+
maxReasoningTokens: maxReasoningTokens,
78+
serviceTier: .default
7879
)
7980
}
8081

@@ -90,6 +91,7 @@ extension BedrockService {
9091
/// - tools: Optional array of tools the model can use
9192
/// - enableReasoning: Optional flag to enable reasoning capabilities
9293
/// - maxReasoningTokens: Optional maximum number of tokens for reasoning
94+
/// - serviceTier: Optional. The service tier to serve this request (.default | .priority | .flex)
9395
/// - Throws: BedrockLibraryError.notSupported for parameters or functionalities that are not supported
9496
/// BedrockLibraryError.invalidParameter for invalid parameters
9597
/// BedrockLibraryError.invalidPrompt if the prompt is empty or too long
@@ -107,7 +109,8 @@ extension BedrockService {
107109
systemPrompts: [String]? = nil,
108110
tools: [Tool]? = nil,
109111
enableReasoning: Bool? = false,
110-
maxReasoningTokens: Int? = nil
112+
maxReasoningTokens: Int? = nil,
113+
serviceTier: ServiceTier
111114
) async throws -> ConverseReplyStream {
112115
do {
113116
guard model.hasConverseStreamingModality() else {
@@ -138,6 +141,7 @@ extension BedrockService {
138141
"stopSequences": "\(String(describing: stopSequences))",
139142
"systemPrompts": "\(String(describing: systemPrompts))",
140143
"tools": "\(String(describing: tools))",
144+
"serviceTier": "\(serviceTier.rawValue)",
141145
]
142146
)
143147
let converseRequest = ConverseStreamingRequest(
@@ -149,7 +153,8 @@ extension BedrockService {
149153
stopSequences: stopSequences,
150154
systemPrompts: systemPrompts,
151155
tools: tools,
152-
maxReasoningTokens: maxReasoningTokens
156+
maxReasoningTokens: maxReasoningTokens,
157+
serviceTier: serviceTier
153158
)
154159

155160
logger.trace("Creating ConverseStreamingInput")
@@ -220,7 +225,8 @@ extension BedrockService {
220225
stopSequences: builder.stopSequences,
221226
systemPrompts: builder.systemPrompts,
222227
tools: builder.tools,
223-
maxReasoningTokens: builder.maxReasoningTokens
228+
maxReasoningTokens: builder.maxReasoningTokens,
229+
serviceTier: builder.serviceTier
224230
)
225231
return streamingResponse
226232
} catch {

Sources/BedrockService/BedrockRuntimeClient/Converse/ConverseRequest.swift

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ public struct ConverseRequest {
2929
let toolConfig: ToolConfig?
3030
let systemPrompts: [String]?
3131
let maxReasoningTokens: Int?
32+
let serviceTier: ServiceTier
3233

3334
@available(
3435
*,
@@ -56,7 +57,8 @@ public struct ConverseRequest {
5657
stopSequences: stopSequences,
5758
systemPrompts: systemPrompts,
5859
tools: tools,
59-
maxReasoningTokens: maxReasoningTokens
60+
maxReasoningTokens: maxReasoningTokens,
61+
serviceTier: .default
6062
)
6163
}
6264

@@ -69,7 +71,8 @@ public struct ConverseRequest {
6971
stopSequences: [String]?,
7072
systemPrompts: [String]?,
7173
tools: [Tool]?,
72-
maxReasoningTokens: Int?
74+
maxReasoningTokens: Int?,
75+
serviceTier: ServiceTier
7376
) {
7477
self.messages = messages
7578
self.model = model
@@ -86,6 +89,7 @@ public struct ConverseRequest {
8689
} else {
8790
self.toolConfig = nil
8891
}
92+
self.serviceTier = serviceTier
8993
}
9094

9195
func getConverseInput(forRegion region: Region) throws -> ConverseInput {
@@ -94,6 +98,7 @@ public struct ConverseRequest {
9498
inferenceConfig: inferenceConfig?.getSDKInferenceConfig(),
9599
messages: try getSDKMessages(),
96100
modelId: model.getModelIdWithCrossRegionInferencePrefix(region: region),
101+
serviceTier: BedrockRuntimeClientTypes.ServiceTier(type: .init(rawValue: serviceTier.rawValue)),
97102
system: getSDKSystemPrompts(),
98103
toolConfig: try toolConfig?.getSDKToolConfig()
99104
)

0 commit comments

Comments
 (0)