Add basic support for agent runtime (#58)

* add basic support for agentruntime

* Add retrieve example to the CI

* Swift-format

* Update Sources/BedrockService/BedrockService+AgentRuntime.swift

Co-authored-by: Copilot <[email protected]>

* Update Sources/BedrockService/RetrieveResult.swift

Co-authored-by: Copilot <[email protected]>

* provide a convenience wrapper around retrieveinput

* add swift docc header

* swift-format

* move files to an agent runtime directory

* rearrange files in different directory structure

---------

Co-authored-by: Copilot <[email protected]>

Author: sebsto

.github/workflows/pull_request.yml

@@ -91,7 +91,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: "[ 'api-key', 'converse', 'converse-stream', 'embeddings', 'openai', 'text_chat' ]"
+      examples: "[ 'api-key', 'converse', 'converse-stream', 'embeddings', 'openai', 'retrieve', 'text_chat' ]"
 
   swift-6-language-mode:
     name: Swift 6 Language Mode

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/Protocols/BedrockClientProtocol.swift renamed to Sources/BedrockService/BedrockClient/BedrockClientProtocol.swift

@@ -14,6 +14,7 @@
 //===----------------------------------------------------------------------===//
 
 import AWSBedrock
+import AWSBedrockAgentRuntime
 import AWSBedrockRuntime
 import ClientRuntime
 import SmithyHTTPAuthAPI
@@ -38,3 +39,6 @@ extension BedrockClient.BedrockClientConfiguration: @retroactive @unchecked Send
 extension BedrockRuntimeClient.BedrockRuntimeClientConfiguration: @retroactive @unchecked Sendable,
     BedrockConfigProtocol
 {}
+extension BedrockAgentRuntimeClient.BedrockAgentRuntimeClientConfiguration: @retroactive @unchecked Sendable,
+    BedrockConfigProtocol
+{}

Sources/BedrockService/ListModels/ModelSummary.swift renamed to Sources/BedrockService/BedrockClient/ListModels/ModelSummary.swift

@@ -0,0 +1,37 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 AWSClientRuntime
+
+#if canImport(FoundationEssentials)
+import FoundationEssentials
+#else
+import Foundation
+#endif
+
+/// Protocol for Amazon Bedrock Runtime Agent operations
+///
+/// This protocol allows writing mocks for unit tests and provides a clean interface
+/// for knowledge base retrieval operations.
+public protocol BedrockRuntimeAgentProtocol: Sendable {
+    /// Retrieves information from a knowledge base
+    /// - Parameter input: The retrieve input containing query and configuration
+    /// - Returns: RetrieveOutput with the retrieved results
+    /// - Throws: Error if the retrieval operation fails
+    func retrieve(input: RetrieveInput) async throws -> RetrieveOutput
+}
+
+extension BedrockAgentRuntimeClient: @retroactive @unchecked Sendable, BedrockRuntimeAgentProtocol {}

Sources/BedrockService/Protocols/BedrockConfigProtocol.swift renamed to Sources/BedrockService/BedrockConfigProtocol.swift

@@ -0,0 +1,77 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 for RAG applications
+    ///
+    /// This method queries an Amazon Bedrock knowledge base to retrieve relevant information
+    /// that can be used for Retrieval-Augmented Generation (RAG) applications.
+    ///
+    /// - 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 with convenience methods
+    /// - 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 request = RetrieveRequest(
+            knowledgeBaseId: knowledgeBaseId,
+            retrievalQuery: retrievalQuery,
+            numberOfResults: numberOfResults
+        )
+
+        do {
+            let response = try await bedrockAgentRuntimeClient.retrieve(input: request.input)
+            logger.trace("Successfully retrieved from knowledge base")
+            return RetrieveResult(response)
+        } catch {
+            try handleCommonError(error, context: "retrieving from knowledge base")
+        }
+    }
+}