Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/pull_request.yml
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ jobs:
# We pass the list of examples here, but we can't pass an array as argument
# Instead, we pass a String with a valid JSON array.
# The workaround is mentioned here https://github.com/orgs/community/discussions/11692
examples: "[ 'converse', 'converse-stream', 'text_chat' ]"
examples: "[ 'api-key', 'converse', 'converse-stream', 'text_chat' ]"

swift-6-language-mode:
name: Swift 6 Language Mode
Expand Down
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ Package.resolved
.vscode
.env
Makefile

**/temp
node_modules

# **/backend
Expand Down
4 changes: 3 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,9 @@ let bedrock = try await BedrockService(
)
```

### API Key Authentication
### API Key Authentication (temporary unavailable)

This capability will be available when issue [#1979](https://github.com/awslabs/aws-sdk-swift/issues/1979) from the AWS SDK for Swift will be fixed or a workaround provided.

Use an API key for authentication. API keys are generated in the AWS console and provide a simpler authentication method for specific use cases.

Expand Down
20 changes: 16 additions & 4 deletions Sources/BedrockAuthentication.swift
Original file line number Diff line number Diff line change
Expand Up @@ -56,16 +56,15 @@ public enum BedrockAuthentication: Sendable, CustomStringConvertible {
/// Creates an AWS credential identity resolver depending on the authentication parameter.
/// - Parameters:
/// - authentication: The authentication type to use
/// - Returns: An optional AWS credential identity resolver. A nil return value means that the default AWS credential provider chain will be used.
/// - Returns: An optional AWS credential identity resolver. A nil return value means that the default AWS credential provider chain will be used or that the authentication type does not require a specific resolver (like `apiKey`).
///
func getAWSCredentialIdentityResolver(
logger: Logger
) async throws -> (any SmithyIdentity.AWSCredentialIdentityResolver)? {

switch self {
case .default,
.apiKey(_):
return nil
case .default, .apiKey(_):
return nil //TODO should we throw an error when apiKey is used ?
case .profile(let profileName):
return try? ProfileAWSCredentialIdentityResolver(profileName: profileName)
case .sso(let profileName):
Expand All @@ -84,4 +83,17 @@ public enum BedrockAuthentication: Sendable, CustomStringConvertible {
return StaticAWSCredentialIdentityResolver(creds)
}
}

/// Creates a BearerTokenIdentityResolver depending on the authentication parameter.
/// - Returns: An optional BearerTokenIdentityResolver. A nil return value means that the authentication type requires an AWSCredentialsProvider instead (like `default`, `profile`, `sso`, `webIdentity`, or `static`).
/// - Note: Only `apiKey` authentication uses BearerTokenIdentityResolver.
func getBearerTokenIdentityResolver(logger: Logger) -> (any SmithyIdentity.BearerTokenIdentityResolver)? {
guard case .apiKey(let key) = self else {
return nil // Only apiKey authentication uses BearerTokenIdentityResolver
}

// Create a StaticBearerTokenIdentityResolver with the provided API key
let identity = BearerTokenIdentity(token: key)
return StaticBearerTokenIdentityResolver(token: identity)
}
}
43 changes: 32 additions & 11 deletions Sources/BedrockService.swift
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ public struct BedrockService: Sendable {
) async throws {
self.logger = logger ?? BedrockService.createLogger("bedrock.service")
self.logger.trace(
"Initializing SwiftBedrock",
"Initializing BedrockService",
metadata: ["region": .string(region.rawValue)]
)
self.region = region
Expand Down Expand Up @@ -121,7 +121,7 @@ public struct BedrockService: Sendable {
-> BedrockClient
{
let config: BedrockClient.BedrockClientConfiguration = try await prepareConfig(
region: region,
initialConfig: BedrockClient.BedrockClientConfiguration(region: region.rawValue),
authentication: authentication,
logger: logger
)
Expand All @@ -143,7 +143,9 @@ public struct BedrockService: Sendable {
-> BedrockRuntimeClient
{
let config: BedrockRuntimeClient.BedrockRuntimeClientConfiguration = try await prepareConfig(
region: region,
initialConfig: BedrockRuntimeClient.BedrockRuntimeClientConfiguration(
region: region.rawValue
),
authentication: authentication,
logger: logger
)
Expand All @@ -152,13 +154,18 @@ public struct BedrockService: Sendable {

/// Generic function to create client configuration and avoid duplication code.
internal static func prepareConfig<C: BedrockConfigProtocol>(
region: Region,
initialConfig: C,
authentication: BedrockAuthentication,
logger: Logging.Logger
) async throws -> C {
var config: C = try await .init()

config.region = region.rawValue
var config = initialConfig

if logger.logLevel == .trace {
// enable trace HTTP requests and responses for the SDK
// see https://github.com/smithy-lang/smithy-swift/blob/main/Sources/ClientRuntime/Telemetry/Logging/ClientLogMode.swift
config.clientLogMode = .requestAndResponse
}

// support profile, SSO, web identity and static authentication
if let awsCredentialIdentityResolver = try? await authentication.getAWSCredentialIdentityResolver(
Expand All @@ -168,11 +175,25 @@ public struct BedrockService: Sendable {
}

// support API keys
if case .apiKey(let key) = authentication {
config.httpClientConfiguration.defaultHeaders.add(
name: "Authorization",
value: "Bearer \(key)"
)
if case .apiKey(_) = authentication {
// config.httpClientConfiguration.defaultHeaders.add(
// name: "Authorization",
// value: "Bearer \(key)"
// )
if let bearerTokenIdentityresolver = authentication.getBearerTokenIdentityResolver(logger: logger) {
config.bearerTokenIdentityResolver = bearerTokenIdentityresolver

// force utilisation of a bearer token instead of AWS credentials + Signv4
// see https://github.com/awslabs/aws-sdk-swift/blob/15b8951d108968f767f4199a3c011e27ac519d61/Sources/Services/AWSBedrockRuntime/Sources/AWSBedrockRuntime/AuthSchemeResolver.swift#L58
config.authSchemeResolver = DefaultBedrockRuntimeAuthSchemeResolver(authSchemePreference: [
"httpBearerAuth"
])
} else {
// TODO: should we throw an error here ?
logger.error(
"API Key authentication is used but no BearerTokenIdentityResolver is provided. This will lead to issues."
)
}
logger.trace("Using API Key for authentication")
} else {
logger.trace("Using AWS credentials for authentication")
Expand Down
16 changes: 13 additions & 3 deletions Sources/Protocols/BedrockConfigProtocol.swift
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,23 @@
import AWSBedrock
import AWSBedrockRuntime
import ClientRuntime
import SmithyHTTPAuthAPI
import SmithyIdentity

protocol BedrockConfigProtocol {
init() async throws
// support regular AWS Credentials + Sigv4 authentication
var awsCredentialIdentityResolver: any SmithyIdentity.AWSCredentialIdentityResolver { get set }
var httpClientConfiguration: ClientRuntime.HttpClientConfiguration { get set }
var region: String? { get set }

// support bearer token authentication (for API Keys)
var bearerTokenIdentityResolver: any SmithyIdentity.BearerTokenIdentityResolver { get set }
var authSchemeResolver: SmithyHTTPAuthAPI.AuthSchemeResolver { get set }

// not used at the moment, we use the bearer token instead
//var httpClientConfiguration: ClientRuntime.HttpClientConfiguration { get set }

// for debugging
var clientLogMode: ClientRuntime.ClientLogMode { get set }

}
extension BedrockClient.BedrockClientConfiguration: @retroactive @unchecked Sendable, BedrockConfigProtocol {}
extension BedrockRuntimeClient.BedrockRuntimeClientConfiguration: @retroactive @unchecked Sendable,
Expand Down
22 changes: 17 additions & 5 deletions Tests/AuthenticationTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,10 @@
//===----------------------------------------------------------------------===//

import AWSBedrock
import AWSBedrockRuntime
import AwsCommonRuntimeKit
import Logging
import SmithyIdentity
import Testing

@testable import BedrockService
Expand Down Expand Up @@ -88,16 +90,24 @@ extension BedrockServiceTests {
// when
// create bedrock configuration with API Key authentication
let config: BedrockClient.BedrockClientConfiguration = try await BedrockService.prepareConfig(
region: .useast1,
initialConfig: BedrockClient.BedrockClientConfiguration(
region: "us-east-1" // default region
),
authentication: auth,
logger: Logger(label: "test.logger"),
)

// then
#expect(config.region == Region.useast1.rawValue) // default region
#expect(
config.httpClientConfiguration.defaultHeaders.value(for: "Authorization") == "Bearer test-api-key-12345"
)

// check token
let resolver = config.bearerTokenIdentityResolver as? StaticBearerTokenIdentityResolver
let token = try await resolver?.getIdentity(identityProperties: nil).token
#expect(token == testApiKey, "Expected token to match the API key")

// check bearer auth scheme
let authScheme = (config.authSchemeResolver as? DefaultBedrockRuntimeAuthSchemeResolver)?.authSchemePreference
#expect(authScheme?.contains("httpBearerAuth") == true, "Expected auth scheme to be HTTP Bearer")

}

Expand Down Expand Up @@ -147,7 +157,9 @@ extension BedrockServiceTests {

// when
let _: BedrockClient.BedrockClientConfiguration = try await BedrockService.prepareConfig(
region: .useast1,
initialConfig: BedrockClient.BedrockClientConfiguration(
region: "us-east-1" // default region
),
authentication: auth,
logger: logger
)
Expand Down