minor API update

Author: sebsto

Sources/Converse/BedrockService+ConverseStreaming.swift

@@ -34,7 +34,8 @@ extension BedrockService {
     ///           BedrockLibraryError.invalidPrompt if the prompt is empty or too long
     ///           BedrockLibraryError.invalidModality for invalid modality from the selected model
     ///           BedrockLibraryError.invalidSDKResponse if the response body is missing
-    /// - Returns: A stream of ConverseResponseStreaming objects
+    /// - Returns: A ConverseReplyStream object that gives access to the high-level stream of ConverseStreamElements objects 
+    ///            or the low-level stream provided by the AWS SDK.
     public func converseStream(
         with model: BedrockModel,
         conversation: [Message],
@@ -46,7 +47,7 @@ extension BedrockService {
         tools: [Tool]? = nil,
         enableReasoning: Bool? = false,
         maxReasoningTokens: Int? = nil
-    ) async throws -> AsyncThrowingStream<ConverseStreamElement, any Error> {
+    ) async throws -> ConverseReplyStream {
         do {
             guard model.hasConverseStreamingModality() else {
                 throw BedrockLibraryError.invalidModality(
@@ -118,18 +119,18 @@ extension BedrockService {
             // - message metadata
             // see https://github.com/awslabs/aws-sdk-swift/blob/2697fb44f607b9c43ad0ce5ca79867d8d6c545c2/Sources/Services/AWSBedrockRuntime/Sources/AWSBedrockRuntime/Models.swift#L3478
             // it will be the responsibility of the user to handle the stream and re-assemble the messages and content
-            // TODO: should we expose the SDK ConverseStreamOutput from the SDK ? or wrap it (what's the added value) ?
 
-            let reply = ConverseReplyStream(sdkStream)
+            let reply = try ConverseReplyStream(sdkStream)
 
             // this time, a different stream is created from the previous one, this one has the following elements
-            // - content segment: this contains a ContentSegment, an enum which can be a .text(Int, String),
-            //   the integer is the id for the content block that the content segment is a part of,
-            //   the String is the part of text that is send from the model.
-            // - content block complete: this includes the id of the completed content block and the complete content block itself
+            // - messageStart: this is the start of a message, it contains the role (assistant or user)
+            // - text: this is a delta of the text content, it contains the partial text 
+            // - reasoning: this is a delta of the reasoning content, it contains the partial reasoning text
+            // - toolUse: this is a buffered tool use response, it contains the tool name and id, and the input parameters
             // - message complete: this includes the complete Message, ready to be added to the history and used for future requests
+            // - metaData: this is the metadata about the response, it contains statitics about the response, such as the number of tokens used and the latency
 
-            return reply.stream
+            return reply
 
         } catch {
             try handleCommonError(error, context: "invoking converse stream")
@@ -143,7 +144,7 @@ extension BedrockService {
     /// - Returns: A stream of ConverseResponseStreaming objects
     public func converseStream(
         with builder: ConverseRequestBuilder
-    ) async throws -> AsyncThrowingStream<ConverseStreamElement, any Error> {
+    ) async throws -> ConverseReplyStream {
         logger.trace("Conversing and streaming")
         do {
             var history = builder.history

Sources/Converse/Streaming/ConverseReplyStream.swift

@@ -30,15 +30,24 @@ public struct ConverseReplyStream: Sendable {
     package init(
         _ inputStream: AsyncThrowingStream<BedrockRuntimeClientTypes.ConverseStreamOutput, Error>,
         logger: Logger? = nil
-    ) {
+    ) throws {
 
         self.logger = logger ?? .init(label: "ConverseReplyStream")
 
         // store the sdk-provided stream to expose it to developers if needed
         self.sdkStream = inputStream
 
         // build a new stream that will convert the SDK stream output to our own ConverseStreamElement
-        self.stream = AsyncThrowingStream(ConverseStreamElement.self) { continuation in
+        self.stream = try ConverseReplyStream.convertToLibraryStream(inputStream, logger: self.logger)
+    }
+
+    /// Convert the SDK Stream to a highler level stream of ConverseStreamElement
+    private static func convertToLibraryStream(
+        _ inputStream: AsyncThrowingStream<BedrockRuntimeClientTypes.ConverseStreamOutput, Error>,
+        logger: Logger
+    ) throws -> AsyncThrowingStream<ConverseStreamElement, Error> {
+
+        AsyncThrowingStream(ConverseStreamElement.self) { continuation in
             let t = Task {
                 do {
                     var state: StreamState!
@@ -48,7 +57,7 @@ public struct ConverseReplyStream: Sendable {
 
                         switch output {
                         case .messagestart(let event):
-                            logger?.trace("Message Start", metadata: ["event": "\(event)"])
+                            logger.trace("Message Start", metadata: ["event": "\(event)"])
 
                             guard let sdkRole = event.role,
                                 let role = try? Role(from: sdkRole)
@@ -62,7 +71,7 @@ public struct ConverseReplyStream: Sendable {
                         // only received at the start of a tool use block
                         // https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-call.html#conversation-inference-call-response
                         case .contentblockstart(let event):
-                            logger?.trace("Content Block Start")
+                            logger.trace("Content Block Start")
                             guard state.currentBlockId == -1 else {
                                 // If we already have a block started, this is an error
                                 throw BedrockLibraryError.invalidSDKType(
@@ -79,7 +88,7 @@ public struct ConverseReplyStream: Sendable {
                         // do not yield an event here, wait for full ToolUse block to arrive
 
                         case .contentblockdelta(let event):
-                            logger?.trace("Content Block Delta")
+                            logger.trace("Content Block Delta")
                             guard let blockId = event.contentBlockIndex else {
                                 // when there is no blockId, this is an error
                                 throw BedrockLibraryError.invalidSDKType(
@@ -114,22 +123,22 @@ public struct ConverseReplyStream: Sendable {
                                     state.bufferReasoningData.append(redactedContent)
                                 // do not yield partial reasoning data, wait for full JSON data
                                 case .sdkUnknown(let output):
-                                    logger?.warning(
+                                    logger.warning(
                                         "Received unknown SDK Reasoning Delta",
                                         metadata: ["reasoning delta": "\(output)"]
                                     )
                                 }
                             case .sdkUnknown(let output):
-                                logger?.warning(
+                                logger.warning(
                                     "Received unknown SDK Event Delta",
                                     metadata: ["delta": "\(output)"]
                                 )
                             case .none:
-                                logger?.warning("Received none SDK Event Delta")
+                                logger.warning("Received none SDK Event Delta")
                             }
 
                         case .contentblockstop(let event):
-                            logger?.trace("Content Block Stop")
+                            logger.trace("Content Block Stop")
                             guard state.currentBlockId != -1 else {
                                 // If we don't have a block started, this is an error
                                 throw BedrockLibraryError.invalidSDKType(
@@ -165,7 +174,7 @@ public struct ConverseReplyStream: Sendable {
                             state.currentBlockId = -1
 
                         case .messagestop(let event):
-                            logger?.trace("Message Stop")
+                            logger.trace("Message Stop")
                             state.messageComplete = true
 
                             // create a Message with all content blocks
@@ -177,7 +186,7 @@ public struct ConverseReplyStream: Sendable {
                             continuation.yield(.messageComplete(message))
 
                         case .metadata(let event):
-                            logger?.trace("Metadata", metadata: ["event": "\(event)"])
+                            logger.trace("Metadata", metadata: ["event": "\(event)"])
 
                             // Convert the metadata event to our ResponseMetadata type
                             let metadata = try ResponseMetadata(from: event)
@@ -187,7 +196,7 @@ public struct ConverseReplyStream: Sendable {
                             // Handle unknown SDK output
                             // This is a catch-all for any future SDK output types that we don't handle yet
                             // We log it and continue, but we could also throw an error if desired
-                            logger?.warning(
+                            logger.warning(
                                 "Received unknown SDK ConverseStreamOutput",
                                 metadata: ["output": "\(output)"]
                             )
@@ -214,7 +223,7 @@ public struct ConverseReplyStream: Sendable {
                     t.cancel()  // Cancel the task when the stream is terminated
                 }
             }
-        }
+        }        
     }
 
     /// Flushes and processes the buffered content from the stream state

Tests/ConverseStream/ConverseReplyStreamTests.swift

@@ -33,7 +33,7 @@ struct ConverseReplyStreamTests {
     @Test("Test streaming text response")
     func testStreamingTextResponse() async throws {
         // Create the ConverseReplyStream from the simulated stream
-        let converseReplyStream = ConverseReplyStream(createSingleTextBlockStream())
+        let converseReplyStream = try ConverseReplyStream(createSingleTextBlockStream())
 
         // Collect all the stream elements
         var streamElements: [ConverseStreamElement] = []
@@ -77,7 +77,7 @@ struct ConverseReplyStreamTests {
     @Test("Test multiple content blocks")
     func testMultipleContentBlocks() async throws {
         // Create the ConverseReplyStream from the simulated stream
-        let converseReplyStream = ConverseReplyStream(createMultipleContentBlocksStream())
+        let converseReplyStream = try ConverseReplyStream(createMultipleContentBlocksStream())
 
         // Collect all the stream elements
         var streamElements: [ConverseStreamElement] = []
@@ -133,7 +133,7 @@ struct ConverseReplyStreamTests {
     @Test("Test cancellation of never-ending stream")
     func testCancellationOfNeverEndingStream() async throws {
         // Create the ConverseReplyStream from the simulated never-ending stream
-        let converseReplyStream = ConverseReplyStream(createNeverEndingStream())
+        let converseReplyStream = try ConverseReplyStream(createNeverEndingStream())
 
         // Create a task to consume the stream
         let consumptionTask = Task {

Tests/ConverseStream/ConverseStreamReasoningTests.swift

@@ -35,8 +35,8 @@ extension ConverseReplyStreamTests {
         #expect(builder.maxReasoningTokens == 4096)
         #expect(builder.history.count == 0)
 
-        var stream = try await bedrock.converseStream(with: builder)
-        var message: Message = try await validateStream(stream, elementsCount: 6)
+        var reply = try await bedrock.converseStream(with: builder)
+        var message: Message = try await validateStream(reply.stream, elementsCount: 6)
 
         try checkReasoningContent(message)
         try checkTextContent(message, prompt: prompt)
@@ -51,8 +51,8 @@ extension ConverseReplyStreamTests {
         #expect(builder.maxReasoningTokens == 4096)
         #expect(builder.history.count == 2)
 
-        stream = try await bedrock.converseStream(with: builder)
-        message = try await validateStream(stream, elementsCount: 6)
+        reply = try await bedrock.converseStream(with: builder)
+        message = try await validateStream(reply.stream, elementsCount: 6)
 
         try checkReasoningContent(message)
         try checkTextContent(message, prompt: prompt)
@@ -68,8 +68,8 @@ extension ConverseReplyStreamTests {
         #expect(builder.maxReasoningTokens == nil)
         #expect(builder.history.count == 4)
 
-        stream = try await bedrock.converseStream(with: builder)
-        message = try await validateStream(stream, elementsCount: 6, contentCount: 1)
+        reply = try await bedrock.converseStream(with: builder)
+        message = try await validateStream(reply.stream, elementsCount: 6, contentCount: 1)
         try checkTextContent(message, prompt: prompt)
         try checkReasoningContent(message, hasReasoningContent: false)
     }

Tests/ConverseStream/ConverseStreamToolTests.swift

@@ -44,11 +44,11 @@ extension ConverseReplyStreamTests {
         #expect(builder.systemPrompts == ["You are a helpful assistant."])
         #expect(builder.tools != nil)
 
-        var stream = try await bedrock.converseStream(with: builder)
+        var reply = try await bedrock.converseStream(with: builder)
 
         // Collect all the stream elements
         var streamElements: [ConverseStreamElement] = []
-        for try await element in stream {
+        for try await element in reply.stream {
             streamElements.append(element)
         }
 
@@ -88,10 +88,10 @@ extension ConverseReplyStreamTests {
         #expect(builder.history.count == 2)
         #expect(builder.tools != nil)
 
-        stream = try await bedrock.converseStream(with: builder)
+        reply = try await bedrock.converseStream(with: builder)
         // Collect all the stream elements
         streamElements = []
-        for try await element in stream {
+        for try await element in reply.stream {
             streamElements.append(element)
         }