Skip to content

Commit e0211e4

Browse files
sebstosebsto
committed
add basic support for agentruntime
1 parent e07b33d commit e0211e4

File tree

11 files changed

+439
-0
lines changed

11 files changed

+439
-0
lines changed

‎Examples/retrieve/.gitignore‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
.DS_Store
2+
/.build
3+
/Packages
4+
xcuserdata/
5+
DerivedData/
6+
.swiftpm/configuration/registries.json
7+
.swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata
8+
.netrc

‎Examples/retrieve/Package.swift‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
// swift-tools-version: 6.0
2+
// The swift-tools-version declares the minimum version of Swift required to build this package.
3+
4+
import PackageDescription
5+
6+
let package = Package(
7+
name: "Retrieve",
8+
platforms: [.macOS(.v15), .iOS(.v18), .tvOS(.v18)],
9+
products: [
10+
.executable(name: "Retrieve", targets: ["Retrieve"])
11+
],
12+
dependencies: [
13+
// for production use, uncomment the following line
14+
// .package(url: "https://github.com/build-on-aws/swift-bedrock-library.git", branch: "main"),
15+
16+
// for local development, use the following line
17+
.package(name: "swift-bedrock-library", path: "../.."),
18+
19+
.package(url: "https://github.com/apple/swift-log.git", from: "1.5.0"),
20+
],
21+
targets: [
22+
.executableTarget(
23+
name: "Retrieve",
24+
dependencies: [
25+
.product(name: "BedrockService", package: "swift-bedrock-library"),
26+
.product(name: "Logging", package: "swift-log"),
27+
]
28+
)
29+
]
30+
)

‎Examples/retrieve/Sources/Retrieve.swift‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the Swift Bedrock Library open source project
4+
//
5+
// Copyright (c) 2025 Amazon.com, Inc. or its affiliates
6+
// and the Swift Bedrock Library project authors
7+
// Licensed under Apache License v2.0
8+
//
9+
// See LICENSE.txt for license information
10+
// See CONTRIBUTORS.txt for the list of Swift Bedrock Library project authors
11+
//
12+
// SPDX-License-Identifier: Apache-2.0
13+
//
14+
//===----------------------------------------------------------------------===//
15+
16+
import BedrockService
17+
import Logging
18+
19+
@main
20+
struct Main {
21+
static func main() async throws {
22+
do {
23+
try await Main.retrieve()
24+
} catch {
25+
print("Error:\n\(error)")
26+
}
27+
}
28+
29+
static func retrieve() async throws {
30+
var logger = Logger(label: "Retrieve")
31+
logger.logLevel = .debug
32+
33+
let bedrock = try await BedrockService(
34+
region: .uswest2,
35+
logger: logger
36+
// uncomment if you use SSO with AWS Identity Center
37+
// authentication: .sso
38+
)
39+
40+
let knowledgeBaseId = "EQ13XRVPLE"
41+
let query = "should I write open source or open-source"
42+
let numberOfResults = 3
43+
44+
print("Retrieving from knowledge base...")
45+
print("Knowledge Base ID: \(knowledgeBaseId)")
46+
print("Query: \(query)")
47+
print("Number of results: \(numberOfResults)")
48+
print()
49+
50+
let response = try await bedrock.retrieve(
51+
knowledgeBaseId: knowledgeBaseId,
52+
retrievalQuery: query,
53+
numberOfResults: numberOfResults
54+
)
55+
56+
print("Retrieved \(response.results?.count ?? 0) results:")
57+
58+
// Show best match using convenience function
59+
if let bestMatch = response.bestMatch() {
60+
print("\n--- Best Match (Score: \(bestMatch.score ?? 0)) ---")
61+
if let content = bestMatch.content?.text {
62+
print("Content: \(content)")
63+
}
64+
}
65+
66+
// Show all results using convenience property
67+
// if let results = response.results {
68+
// for (index, result) in results.enumerated() {
69+
// print("\n--- Result \(index + 1) ---")
70+
// if let content = result.content?.text {
71+
// print("Content: \(content)")
72+
// }
73+
// if let score = result.score {
74+
// print("Score: \(score)")
75+
// }
76+
// if let location = result.location?.s3Location {
77+
// print("Source: s3://\(location.uri ?? "unknown")")
78+
// }
79+
// }
80+
// }
81+
}
82+
}

‎Package.swift‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ let package = Package(
2323
.product(name: "AWSClientRuntime", package: "aws-sdk-swift"),
2424
.product(name: "AWSBedrock", package: "aws-sdk-swift"),
2525
.product(name: "AWSBedrockRuntime", package: "aws-sdk-swift"),
26+
.product(name: "AWSBedrockAgentRuntime", package: "aws-sdk-swift"),
2627
.product(name: "AWSSSOOIDC", package: "aws-sdk-swift"),
2728
.product(name: "Smithy", package: "smithy-swift"),
2829
.product(name: "Logging", package: "swift-log"),

‎Sources/BedrockService/BedrockService+AgentRuntime.swift‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the Swift Bedrock Library open source project
4+
//
5+
// Copyright (c) 2025 Amazon.com, Inc. or its affiliates
6+
// and the Swift Bedrock Library project authors
7+
// Licensed under Apache License v2.0
8+
//
9+
// See LICENSE.txt for license information
10+
// See CONTRIBUTORS.txt for the list of Swift Bedrock Library project authors
11+
//
12+
// SPDX-License-Identifier: Apache-2.0
13+
//
14+
//===----------------------------------------------------------------------===//
15+
16+
@preconcurrency import AWSBedrockAgentRuntime
17+
import Logging
18+
19+
extension BedrockService {
20+
/// Creates a BedrockAgentRuntimeClient
21+
/// - Parameters:
22+
/// - region: The AWS region to configure the client for
23+
/// - authentication: The authentication type to use
24+
/// - logger: Logger instance
25+
/// - Returns: Configured BedrockAgentRuntimeProtocol instance
26+
/// - Throws: Error if client creation fails
27+
internal static func createBedrockAgentRuntimeClient(
28+
region: Region,
29+
authentication: BedrockAuthentication,
30+
logger: Logging.Logger
31+
) async throws -> BedrockAgentRuntimeClient {
32+
let config: BedrockAgentRuntimeClient.BedrockAgentRuntimeClientConfiguration = try await prepareConfig(
33+
initialConfig: BedrockAgentRuntimeClient.BedrockAgentRuntimeClientConfiguration(region: region.rawValue),
34+
authentication: authentication,
35+
logger: logger
36+
)
37+
return BedrockAgentRuntimeClient(config: config)
38+
}
39+
/// Retrieves information from a knowledge base
40+
/// - Parameters:
41+
/// - knowledgeBaseId: The unique identifier of the knowledge base to query
42+
/// - retrievalQuery: The query to search for in the knowledge base
43+
/// - numberOfResults: The number of results to return (optional, defaults to 3)
44+
/// - Returns: RetrieveResult containing the retrieved results
45+
/// - Throws: BedrockLibraryError or other errors from the underlying service
46+
public func retrieve(
47+
knowledgeBaseId: String,
48+
retrievalQuery: String,
49+
numberOfResults: Int = 3
50+
) async throws -> RetrieveResult {
51+
logger.trace("Retrieving from knowledge base", metadata: [
52+
"knowledgeBaseId": .string(knowledgeBaseId),
53+
"numberOfResults": .stringConvertible(numberOfResults)
54+
])
55+
56+
let input = RetrieveInput(
57+
knowledgeBaseId: knowledgeBaseId,
58+
retrievalQuery: BedrockAgentRuntimeClientTypes.KnowledgeBaseQuery(text: retrievalQuery)
59+
)
60+
61+
do {
62+
let response = try await bedrockAgentRuntimeClient.retrieve(input: input)
63+
logger.trace("Successfully retrieved from knowledge base")
64+
return RetrieveResult(response)
65+
} catch {
66+
try handleCommonError(error, context: "retrieving from knowledge base")
67+
}
68+
}
69+
}

‎Sources/BedrockService/BedrockService.swift‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
@preconcurrency import AWSBedrock
1717
@preconcurrency import AWSBedrockRuntime
18+
@preconcurrency import AWSBedrockAgentRuntime
1819
import AWSClientRuntime
1920
import AWSSSOOIDC
2021
import AwsCommonRuntimeKit
@@ -38,6 +39,7 @@ public struct BedrockService: Sendable {
3839
package let logger: Logging.Logger
3940
package let bedrockClient: BedrockClientProtocol
4041
package let bedrockRuntimeClient: BedrockRuntimeClientProtocol
42+
package let bedrockAgentRuntimeClient: BedrockAgentRuntimeProtocol
4143

4244
// MARK: - Initialization
4345

@@ -47,13 +49,15 @@ public struct BedrockService: Sendable {
4749
/// - logger: Optional custom logger instance
4850
/// - bedrockClient: Optional custom Bedrock client
4951
/// - bedrockRuntimeClient: Optional custom Bedrock Runtime client
52+
/// - bedrockAgentRuntimeClient: Optional custom Bedrock Agent Runtime client
5053
/// - authentication: The authentication type to use (defaults to .default)
5154
/// - Throws: Error if client initialization fails
5255
public init(
5356
region: Region = .useast1,
5457
logger: Logging.Logger? = nil,
5558
bedrockClient: BedrockClientProtocol? = nil,
5659
bedrockRuntimeClient: BedrockRuntimeClientProtocol? = nil,
60+
bedrockAgentRuntimeClient: BedrockAgentRuntimeProtocol? = nil,
5761
authentication: BedrockAuthentication = .default
5862
) async throws {
5963
self.logger = logger ?? BedrockService.createLogger("bedrock.service")
@@ -93,6 +97,22 @@ public struct BedrockService: Sendable {
9397
metadata: ["authentication type": "\(authentication)"]
9498
)
9599
}
100+
if bedrockAgentRuntimeClient != nil {
101+
self.logger.trace("Using supplied bedrockAgentRuntimeClient")
102+
self.bedrockAgentRuntimeClient = bedrockAgentRuntimeClient!
103+
} else {
104+
self.logger.trace("Creating bedrockAgentRuntimeClient")
105+
self.bedrockAgentRuntimeClient = try await BedrockService.createBedrockAgentRuntimeClient(
106+
region: region,
107+
authentication: authentication,
108+
logger: self.logger
109+
)
110+
self.logger.trace(
111+
"Created bedrockAgentRuntimeClient",
112+
metadata: ["authentication type": "\(authentication)"]
113+
)
114+
}
115+
96116
self.logger.trace(
97117
"Initialized SwiftBedrock",
98118
metadata: ["region": .string(region.rawValue)]

‎Sources/BedrockService/Docs.docc/BedrockService.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ BedrockService is a lightweight layer on top of the AWS SDK for Swift that provi
3434
- <doc:Documents>
3535
- <doc:Reasoning>
3636
- <doc:Streaming>
37+
- <doc:Retrieval>
3738

3839
### Extending the Library
3940

0 commit comments

Comments
 (0)