Add basic support for agent runtime

Examples/retrieve/.gitignore

@@ -0,0 +1,8 @@
+.DS_Store
+/.build
+/Packages
+xcuserdata/
+DerivedData/
+.swiftpm/configuration/registries.json
+.swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata
+.netrc

Examples/retrieve/Package.swift

@@ -0,0 +1,30 @@
+// swift-tools-version: 6.0
+// The swift-tools-version declares the minimum version of Swift required to build this package.
+
+import PackageDescription
+
+let package = Package(
+    name: "Retrieve",
+    platforms: [.macOS(.v15), .iOS(.v18), .tvOS(.v18)],
+    products: [
+        .executable(name: "Retrieve", targets: ["Retrieve"])
+    ],
+    dependencies: [
+        // for production use, uncomment the following line
+        // .package(url: "https://github.com/build-on-aws/swift-bedrock-library.git", branch: "main"),
+
+        // for local development, use the following line
+        .package(name: "swift-bedrock-library", path: "../.."),
+
+        .package(url: "https://github.com/apple/swift-log.git", from: "1.5.0"),
+    ],
+    targets: [
+        .executableTarget(
+            name: "Retrieve",
+            dependencies: [
+                .product(name: "BedrockService", package: "swift-bedrock-library"),
+                .product(name: "Logging", package: "swift-log"),
+            ]
+        )
+    ]
+)

Examples/retrieve/Sources/Retrieve.swift

@@ -0,0 +1,82 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Bedrock Library open source project
+//
+// Copyright (c) 2025 Amazon.com, Inc. or its affiliates
+//                    and the Swift Bedrock Library project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift Bedrock Library project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import BedrockService
+import Logging
+
+@main
+struct Main {
+    static func main() async throws {
+        do {
+            try await Main.retrieve()
+        } catch {
+            print("Error:\n\(error)")
+        }
+    }
+
+    static func retrieve() async throws {
+        var logger = Logger(label: "Retrieve")
+        logger.logLevel = .debug
+
+        let bedrock = try await BedrockService(
+            region: .uswest2,
+            logger: logger
+            // uncomment if you use SSO with AWS Identity Center
+            // authentication: .sso
+        )
+
+        let knowledgeBaseId = "EQ13XRVPLE"
+        let query = "should I write open source or open-source"
+        let numberOfResults = 3
+
+        print("Retrieving from knowledge base...")
+        print("Knowledge Base ID: \(knowledgeBaseId)")
+        print("Query: \(query)")
+        print("Number of results: \(numberOfResults)")
+        print()
+
+        let response = try await bedrock.retrieve(
+            knowledgeBaseId: knowledgeBaseId,
+            retrievalQuery: query,
+            numberOfResults: numberOfResults
+        )
+
+        print("Retrieved \(response.results?.count ?? 0) results:")
+
+        // Show best match using convenience function
+        if let bestMatch = response.bestMatch() {
+            print("\n--- Best Match (Score: \(bestMatch.score ?? 0)) ---")
+            if let content = bestMatch.content?.text {
+                print("Content: \(content)")
+            }
+        }
+
+        // Show all results using convenience property
+        // if let results = response.results {
+        //     for (index, result) in results.enumerated() {
+        //         print("\n--- Result \(index + 1) ---")
+        //         if let content = result.content?.text {
+        //             print("Content: \(content)")
+        //         }
+        //         if let score = result.score {
+        //             print("Score: \(score)")
+        //         }
+        //         if let location = result.location?.s3Location {
+        //             print("Source: s3://\(location.uri ?? "unknown")")
+        //         }
+        //     }
+        // }
+    }
+}

Package.swift

@@ -23,6 +23,7 @@ let package = Package(
                 .product(name: "AWSClientRuntime", package: "aws-sdk-swift"),
                 .product(name: "AWSBedrock", package: "aws-sdk-swift"),
                 .product(name: "AWSBedrockRuntime", package: "aws-sdk-swift"),
+                .product(name: "AWSBedrockAgentRuntime", package: "aws-sdk-swift"),
                 .product(name: "AWSSSOOIDC", package: "aws-sdk-swift"),
                 .product(name: "Smithy", package: "smithy-swift"),
                 .product(name: "Logging", package: "swift-log"),

Sources/BedrockService/BedrockService+AgentRuntime.swift

@@ -0,0 +1,69 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Bedrock Library open source project
+//
+// Copyright (c) 2025 Amazon.com, Inc. or its affiliates
+//                    and the Swift Bedrock Library project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift Bedrock Library project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+@preconcurrency import AWSBedrockAgentRuntime
+import Logging
+
+extension BedrockService {
+    /// Creates a BedrockAgentRuntimeClient
+    /// - Parameters:
+    ///   - region: The AWS region to configure the client for
+    ///   - authentication: The authentication type to use
+    ///   - logger: Logger instance
+    /// - Returns: Configured BedrockAgentRuntimeProtocol instance
+    /// - Throws: Error if client creation fails
+    internal static func createBedrockAgentRuntimeClient(
+        region: Region,
+        authentication: BedrockAuthentication,
+        logger: Logging.Logger
+    ) async throws -> BedrockAgentRuntimeClient {
+        let config: BedrockAgentRuntimeClient.BedrockAgentRuntimeClientConfiguration = try await prepareConfig(
+            initialConfig: BedrockAgentRuntimeClient.BedrockAgentRuntimeClientConfiguration(region: region.rawValue),
+            authentication: authentication,
+            logger: logger
+        )
+        return BedrockAgentRuntimeClient(config: config)
+    }
+    /// Retrieves information from a knowledge base
+    /// - Parameters:
+    ///   - knowledgeBaseId: The unique identifier of the knowledge base to query
+    ///   - retrievalQuery: The query to search for in the knowledge base
+    ///   - numberOfResults: The number of results to return (optional, defaults to 3)
+    /// - Returns: RetrieveResult containing the retrieved results
+    /// - Throws: BedrockLibraryError or other errors from the underlying service
+    public func retrieve(
+        knowledgeBaseId: String,
+        retrievalQuery: String,
+        numberOfResults: Int = 3
+    ) async throws -> RetrieveResult {
+        logger.trace("Retrieving from knowledge base", metadata: [
+            "knowledgeBaseId": .string(knowledgeBaseId),
+            "numberOfResults": .stringConvertible(numberOfResults)
+        ])
+
+        let input = RetrieveInput(
+            knowledgeBaseId: knowledgeBaseId,
+            retrievalQuery: BedrockAgentRuntimeClientTypes.KnowledgeBaseQuery(text: retrievalQuery)
+        )
+
+        do {
+            let response = try await bedrockAgentRuntimeClient.retrieve(input: input)
+            logger.trace("Successfully retrieved from knowledge base")
+            return RetrieveResult(response)
+        } catch {
+            try handleCommonError(error, context: "retrieving from knowledge base")
+        }
+    }
+}

Sources/BedrockService/BedrockService.swift

@@ -15,6 +15,7 @@
 
 @preconcurrency import AWSBedrock
 @preconcurrency import AWSBedrockRuntime
+@preconcurrency import AWSBedrockAgentRuntime
 import AWSClientRuntime
 import AWSSSOOIDC
 import AwsCommonRuntimeKit
@@ -38,6 +39,7 @@ public struct BedrockService: Sendable {
     package let logger: Logging.Logger
     package let bedrockClient: BedrockClientProtocol
     package let bedrockRuntimeClient: BedrockRuntimeClientProtocol
+    package let bedrockAgentRuntimeClient: BedrockAgentRuntimeProtocol
 
     // MARK: - Initialization
 
@@ -47,13 +49,15 @@ public struct BedrockService: Sendable {
     ///   - logger: Optional custom logger instance
     ///   - bedrockClient: Optional custom Bedrock client
     ///   - bedrockRuntimeClient: Optional custom Bedrock Runtime client
+    ///   - bedrockAgentRuntimeClient: Optional custom Bedrock Agent Runtime client
     ///   - authentication: The authentication type to use (defaults to .default)
     /// - Throws: Error if client initialization fails
     public init(
         region: Region = .useast1,
         logger: Logging.Logger? = nil,
         bedrockClient: BedrockClientProtocol? = nil,
         bedrockRuntimeClient: BedrockRuntimeClientProtocol? = nil,
+        bedrockAgentRuntimeClient: BedrockAgentRuntimeProtocol? = nil,
         authentication: BedrockAuthentication = .default
     ) async throws {
         self.logger = logger ?? BedrockService.createLogger("bedrock.service")
@@ -93,6 +97,22 @@ public struct BedrockService: Sendable {
                 metadata: ["authentication type": "\(authentication)"]
             )
         }
+        if bedrockAgentRuntimeClient != nil {
+            self.logger.trace("Using supplied bedrockAgentRuntimeClient")
+            self.bedrockAgentRuntimeClient = bedrockAgentRuntimeClient!
+        } else {
+            self.logger.trace("Creating bedrockAgentRuntimeClient")
+            self.bedrockAgentRuntimeClient = try await BedrockService.createBedrockAgentRuntimeClient(
+                region: region,
+                authentication: authentication,
+                logger: self.logger
+            )
+            self.logger.trace(
+                "Created bedrockAgentRuntimeClient",
+                metadata: ["authentication type": "\(authentication)"]
+            )
+        }
+
         self.logger.trace(
             "Initialized SwiftBedrock",
             metadata: ["region": .string(region.rawValue)]

Sources/BedrockService/Docs.docc/BedrockService.md

@@ -34,6 +34,7 @@ BedrockService is a lightweight layer on top of the AWS SDK for Swift that provi
 - <doc:Documents>
 - <doc:Reasoning>
 - <doc:Streaming>
+- <doc:Retrieval>
 
 ### Extending the Library