Skip to content

Keep non-text response segments when a stream finalizes #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ActuallyTaylor wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Keep non-text response segments when a stream finalizes #16

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 16 additions & 20 deletions Sources/AnyLanguageModel/Transcript.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,15 +60,18 @@ public struct Transcript: Sendable, Equatable, Codable {
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.
/// 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 {
guard case .response(var response) = entries.last else {
append(
Entry.response(
Response(
Expand All @@ -82,26 +85,19 @@ public struct Transcript: Sendable, Equatable, Codable {
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))
]
// 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 {
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)))
}
response.assetIDs = assetIDs

replace(index: entries.count - 1, with: newResponse)
// Replace the latest entry with the one we just updated.
replace(index: entries.count - 1, with: .response(response))
}

/// An entry in a transcript.
Expand Down
88 changes: 88 additions & 0 deletions Tests/AnyLanguageModelTests/TranscriptTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,94 @@ 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 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"))
Expand Down
Loading