Keep non-text response segments when a stream finalizes

Sources/AnyLanguageModel/Transcript.swift

@@ -26,17 +26,18 @@ public struct Transcript: Sendable, Equatable, Codable {
         entries[index] = entry
     }
 
-    /// Updates a transcript with temporary text that is being streamed from a model.
-    /// Appends the assistant response to the end of entries, if the last entry is a response then that response is updated with the newest streamed text.
+    /// Replace a trailing response entry with `text`. If `assetIDs` is populated and there is no existing response these IDs will be used as the IDs for the response. If a response already exists, they will be appended to any existing IDs.
     ///
-    /// - Parameter text: The text to update the transcript with.
-    mutating func appendStreamingResponse(_ text: String) {
+    /// - Parameters:
+    ///   - text: The text to replace the trailing response with.
+    ///   - assetIDs: The asset IDs to use for the response.
+    private mutating func replaceOrAppendTrailingResponse(text: String, assetIDs: [String]) {
         // Make sure the last entry in the transcript is a response. If it is not, create a new response and append it to the end of the transcript.
         guard case .response(var response) = entries.last else {
             append(
                 Entry.response(
                     Response(
-                        assetIDs: [],
+                        assetIDs: assetIDs,
                         segments: [
                             Transcript.Segment.text(Transcript.TextSegment(content: text))
                         ]
@@ -46,62 +47,52 @@ public struct Transcript: Sendable, Equatable, Codable {
             return
         }
 
-        // If the last segment in the last response is text, replace it with the new content.
+        // Streamed text only ever lands in the trailing text segment, so replace that one
+        // in place — carrying its ID over — and leave every earlier segment untouched.
         if case .text(let last)? = response.segments.last {
-            // Keep the same ID as the last segment.
             response.segments[response.segments.count - 1] = Transcript.Segment.text(
                 Transcript.TextSegment(id: last.id, content: text)
             )
         } else {
             response.segments.append(Transcript.Segment.text(Transcript.TextSegment(content: text)))
         }
 
+        // Asset IDs identify everything used to generate the response, so keep the ones the
+        // in-progress response already recorded and add whatever the finalized one introduced.
+        for assetID in assetIDs where !response.assetIDs.contains(assetID) {
+            response.assetIDs.append(assetID)
+        }
+
         // Replace the latest entry with the one we just updated.
         replace(index: entries.count - 1, with: .response(response))
     }
 
-    /// Replaces the trailing response entry's text with the final text, or appends a new response entry if the last entry isn't a response.
+    /// Updates a transcript with temporary text that is being streamed from a model.
+    /// Appends the assistant response to the end of entries, if the last entry is a response then that response is updated with the newest streamed text.
+    ///
+    /// - Parameter text: The text to update the transcript with.
+    mutating func appendStreamingResponse(_ text: String) {
+        replaceOrAppendTrailingResponse(text: text, assetIDs: [])
+    }
+
+    /// Replaces the trailing response entry's streamed text with the final text, or appends a new response entry if the last entry isn't a response.
     /// Prevents streamed responses from having duplicate entries on completion.
     ///
+    /// Segments ahead of the trailing text are left alone, so anything else the model
+    /// produced during the same response survives the end of the stream.
+    ///
     /// - Parameters:
     ///   - text: The text to replace the final response with.
     ///   - assetIDs: The assetIDs for the response.
     mutating func finalizeStreamedTranscript(_ text: String, assetIDs: [String]) {
-        // Make sure the last entry in the transcript is a response. If it is not, create a new response and append it to the end of the transcript.
-        guard case .response(let response) = entries.last else {
-            append(
-                Entry.response(
-                    Response(
-                        assetIDs: assetIDs,
-                        segments: [
-                            Transcript.Segment.text(Transcript.TextSegment(content: text))
-                        ]
-                    )
-                )
-            )
+        // A turn that ended in tool calls already recorded its streamed text in the response
+        // entry ahead of them, so there is nothing left to finalize. Appending here would
+        // duplicate that text after the tool calls.
+        if case .toolCalls = entries.last {
             return
         }
 
-        // If the last segment is text we want to carry its ID over to the new text segment. Otherwise generate a new ID for it.
-        let id =
-            switch response.segments.last {
-            case .text(let last):
-                last.id
-            default:
-                UUID().uuidString
-            }
-
-        let newResponse: Entry = Entry.response(
-            Response(
-                id: response.id,
-                assetIDs: assetIDs,
-                segments: [
-                    Transcript.Segment.text(Transcript.TextSegment(id: id, content: text))
-                ]
-            )
-        )
-
-        replace(index: entries.count - 1, with: newResponse)
+        replaceOrAppendTrailingResponse(text: text, assetIDs: assetIDs)
     }
 
     /// An entry in a transcript.

Tests/AnyLanguageModelTests/TranscriptTests.swift

@@ -108,6 +108,163 @@ struct TranscriptTests {
         }
     }
 
+    @Test func finalizeStreamedTranscriptKeepsSegmentsAheadOfTheStreamedText() throws {
+        let structured = Transcript.StructuredSegment(
+            id: "structured-id",
+            source: "source",
+            content: try GeneratedContent(json: #"{"ok":true}"#)
+        )
+        var transcript = Transcript(entries: [
+            .response(
+                .init(
+                    id: "response-id",
+                    assetIDs: [],
+                    segments: [
+                        .structure(structured),
+                        .text(.init(id: "text-id", content: "partial")),
+                    ]
+                )
+            )
+        ])
+
+        transcript.finalizeStreamedTranscript("complete", assetIDs: ["asset"])
+
+        guard case .response(let response)? = transcript.last else {
+            Issue.record("Expected a trailing response entry")
+            return
+        }
+        #expect(response.id == "response-id")
+        #expect(response.assetIDs == ["asset"])
+        #expect(response.segments.count == 2)
+        #expect(response.segments.first == .structure(structured))
+        // The trailing text segment is replaced in place, keeping its ID stable.
+        #expect(response.segments.last == .text(.init(id: "text-id", content: "complete")))
+    }
+
+    @Test func finalizeStreamedTranscriptAppendsTextWhenTrailingSegmentIsNotText() throws {
+        let structured = Transcript.StructuredSegment(
+            id: "structured-id",
+            source: "source",
+            content: try GeneratedContent(json: #"{"ok":true}"#)
+        )
+        var transcript = Transcript(entries: [
+            .response(.init(id: "response-id", assetIDs: [], segments: [.structure(structured)]))
+        ])
+
+        transcript.finalizeStreamedTranscript("complete", assetIDs: [])
+
+        guard case .response(let response)? = transcript.last else {
+            Issue.record("Expected a trailing response entry")
+            return
+        }
+        #expect(response.segments.count == 2)
+        #expect(response.segments.first == .structure(structured))
+        #expect(response.segments.last?.description == "complete")
+    }
+
+    @Test func finalizeStreamedTranscriptAppendsResponseWhenLastEntryIsNotAResponse() {
+        var transcript = Transcript(entries: [
+            .prompt(.init(id: "prompt-id", segments: [.text(.init(content: "Hello"))]))
+        ])
+
+        transcript.finalizeStreamedTranscript("complete", assetIDs: ["asset"])
+
+        #expect(transcript.count == 2)
+        guard case .response(let response)? = transcript.last else {
+            Issue.record("Expected a trailing response entry")
+            return
+        }
+        #expect(response.assetIDs == ["asset"])
+        #expect(response.segments.count == 1)
+        #expect(response.segments.first?.description == "complete")
+    }
+
+    @Test func finalizeStreamedTranscriptLeavesTurnsThatEndedInToolCallsAlone() throws {
+        let arguments = try GeneratedContent(json: #"{"city":"Cupertino"}"#)
+        let toolCalls = Transcript.ToolCalls(
+            id: "tool-calls-id",
+            [Transcript.ToolCall(id: "call-id", toolName: "getWeather", arguments: arguments)]
+        )
+        var transcript = Transcript(entries: [
+            .prompt(.init(id: "prompt-id", segments: [.text(.init(content: "Hello"))]))
+        ])
+        transcript.appendStreamingResponse("Let me check")
+        transcript.append(.toolCalls(toolCalls))
+
+        transcript.finalizeStreamedTranscript("Let me check", assetIDs: [])
+
+        // The streamed text is already recorded ahead of the tool calls, so nothing is appended.
+        #expect(transcript.count == 3)
+        guard case .toolCalls(let trailing)? = transcript.last else {
+            Issue.record("Expected a trailing tool calls entry")
+            return
+        }
+        #expect(trailing.id == "tool-calls-id")
+        guard case .response(let response) = transcript[1] else {
+            Issue.record("Expected a response entry ahead of the tool calls")
+            return
+        }
+        #expect(response.segments.count == 1)
+        #expect(response.segments.first?.description == "Let me check")
+    }
+
+    @Test func finalizeStreamedTranscriptKeepsAssetIDsAlreadyRecorded() {
+        var transcript = Transcript(entries: [
+            .response(
+                .init(
+                    id: "response-id",
+                    assetIDs: ["existing"],
+                    segments: [.text(.init(id: "text-id", content: "partial"))]
+                )
+            )
+        ])
+
+        transcript.finalizeStreamedTranscript("complete", assetIDs: [])
+
+        guard case .response(let response)? = transcript.last else {
+            Issue.record("Expected a trailing response entry")
+            return
+        }
+        #expect(response.assetIDs == ["existing"])
+    }
+
+    @Test func finalizeStreamedTranscriptMergesNewAssetIDsWithoutDuplicating() {
+        var transcript = Transcript(entries: [
+            .response(
+                .init(
+                    id: "response-id",
+                    assetIDs: ["existing"],
+                    segments: [.text(.init(id: "text-id", content: "partial"))]
+                )
+            )
+        ])
+
+        transcript.finalizeStreamedTranscript("complete", assetIDs: ["existing", "added"])
+
+        guard case .response(let response)? = transcript.last else {
+            Issue.record("Expected a trailing response entry")
+            return
+        }
+        #expect(response.assetIDs == ["existing", "added"])
+    }
+
+    @Test func appendStreamingResponseGrowsTheTrailingTextSegmentInPlace() {
+        var transcript = Transcript(entries: [
+            .prompt(.init(id: "prompt-id", segments: [.text(.init(content: "Hello"))]))
+        ])
+
+        transcript.appendStreamingResponse("He")
+        transcript.appendStreamingResponse("Hello")
+
+        #expect(transcript.count == 2)
+        guard case .response(let response)? = transcript.last else {
+            Issue.record("Expected a trailing response entry")
+            return
+        }
+        #expect(response.segments.count == 1)
+        #expect(response.segments.first?.description == "Hello")
+    }
+
     @Test func responseFormatNameExtractsRefTypeNameOrFallsBack() {
         let refFormat = Transcript.ResponseFormat(type: Person.self)
         #expect(refFormat.name.contains("Person"))