diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index f256d046..3310fc4a 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/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,9 +47,9 @@ 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) ) @@ -56,52 +57,42 @@ public struct Transcript: Sendable, Equatable, Codable { 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. diff --git a/Tests/AnyLanguageModelTests/TranscriptTests.swift b/Tests/AnyLanguageModelTests/TranscriptTests.swift index 037c70a0..9495ca47 100644 --- a/Tests/AnyLanguageModelTests/TranscriptTests.swift +++ b/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"))