From e4426af862332208f485ea97272908fbf1c34348 Mon Sep 17 00:00:00 2001 From: mattt Date: Thu, 22 Oct 2026 06:57:59 -0700 Subject: [PATCH 1/2] Send partial snapshots of structured streams Build each structured stream snapshot from the content's partial form, so a snapshot arrives before every required property has, and a usage update sends one too. A stopped tool round still uses the response text only when it decodes as the complete content. --- .../Shared/ResponseUsage.swift | 4 +- .../Shared/StreamingResponseState.swift | 5 +- .../ProviderUsageTests.swift | 93 ++++++++++++++----- 3 files changed, 77 insertions(+), 25 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/ResponseUsage.swift b/Sources/AnyLanguageModel/Shared/ResponseUsage.swift index 403379be..58fa1e47 100644 --- a/Sources/AnyLanguageModel/Shared/ResponseUsage.swift +++ b/Sources/AnyLanguageModel/Shared/ResponseUsage.swift @@ -222,8 +222,8 @@ extension LanguageModelSession.ResponseStream.Snapshot { (try? GeneratedContent(json: text)) ?? GeneratedContent(text), fromProperty: wrapperProperty ) - guard let content = try? Content(raw) else { return nil } - self.init(content: content.asPartiallyGenerated(), rawContent: raw, requests: requests) + guard let content = try? Content.PartiallyGenerated(raw) else { return nil } + self.init(content: content, rawContent: raw, requests: requests) } } } diff --git a/Sources/AnyLanguageModel/Shared/StreamingResponseState.swift b/Sources/AnyLanguageModel/Shared/StreamingResponseState.swift index 598088c2..354a8cb1 100644 --- a/Sources/AnyLanguageModel/Shared/StreamingResponseState.swift +++ b/Sources/AnyLanguageModel/Shared/StreamingResponseState.swift @@ -50,12 +50,13 @@ struct StreamingResponseState { /// A stopped tool call can have no response text, including for structured generation. /// - /// Without text, the snapshot uses the first empty value that the content type accepts: + /// When the text so far doesn't decode as the complete content type, + /// the snapshot uses the first empty value that the content type accepts: /// an empty object, an empty array, `null`, and then zero or `false` for scalar types. /// The value must decode as the complete content type, not only its partial form, /// so that `collect()` accepts the snapshot. func stoppedSnapshot() throws -> LanguageModelSession.ResponseStream.Snapshot { - if let snapshot = snapshot() { return snapshot } + if let snapshot = snapshot(), (try? Content(snapshot.rawContent)) != nil { return snapshot } let candidates: [GeneratedContent.Kind] = [ .structure(properties: [:], orderedKeys: []), .array([]), .null, .number(0), .bool(false), ] diff --git a/Tests/AnyLanguageModelTests/ProviderUsageTests.swift b/Tests/AnyLanguageModelTests/ProviderUsageTests.swift index ed234995..92fb58f1 100644 --- a/Tests/AnyLanguageModelTests/ProviderUsageTests.swift +++ b/Tests/AnyLanguageModelTests/ProviderUsageTests.swift @@ -262,22 +262,24 @@ import Testing } func stream(text: String = "Hello", includeUsage: Bool = true) throws -> String { + try stream(chunks: [text], includeUsage: includeUsage) + } + + /// A stream that sends the response text in `chunks`, one event each. + func stream(chunks: [String], includeUsage: Bool = true) throws -> String { + let text = chunks.joined() var events: [[String: Any]] switch self { case .chat: - events = [ - ["id": "test", "choices": [["delta": ["content": text]]]], - ["id": "test", "choices": [["delta": [:], "finish_reason": "stop"]]], - ] + events = chunks.map { ["id": "test", "choices": [["delta": ["content": $0]]]] } + events.append(["id": "test", "choices": [["delta": [:], "finish_reason": "stop"]]]) if includeUsage { events.append(["id": "test", "choices": [], "usage": counts]) } case .responses, .openResponses: - events = [ - ["type": "response.output_text.delta", "delta": text], - [ - "type": "response.completed", - "response": response(text: text, counts: includeUsage ? counts : nil), - ], - ] + events = chunks.map { ["type": "response.output_text.delta", "delta": $0] } + events.append([ + "type": "response.completed", + "response": response(text: text, counts: includeUsage ? counts : nil), + ]) case .anthropic: var startCounts = counts startCounts["output_tokens"] = 0 @@ -285,9 +287,13 @@ import Testing [ "type": "message_start", "message": response(text: "", counts: includeUsage ? startCounts : nil), - ], - ["type": "content_block_delta", "index": 0, "delta": ["type": "text_delta", "text": text]], + ] ] + for chunk in chunks { + events.append([ + "type": "content_block_delta", "index": 0, "delta": ["type": "text_delta", "text": chunk], + ]) + } for output in [7, 20] { var delta: [String: Any] = ["type": "message_delta", "delta": ["stop_reason": "end_turn"]] if includeUsage { delta["usage"] = ["output_tokens": output] } @@ -298,14 +304,17 @@ import Testing var early = counts early["candidatesTokenCount"] = 7 early["totalTokenCount"] = 122 - events = [response(text: text, counts: includeUsage ? early : nil)] + events = chunks.map { response(text: $0, counts: includeUsage ? early : nil) } if includeUsage { events.append(["usageMetadata": counts]) } case .ollama: - var first = response(text: text) - first["done"] = false + events = chunks.map { + var event = response(text: $0) + event["done"] = false + return event + } var last = response(counts: includeUsage ? counts : nil) last["message"] = ["role": "assistant"] - events = [first, last] + events.append(last) } let json = try events.map { try ProviderUsageTests.json($0) } if self == .ollama { return json.joined(separator: "\n") + "\n" } @@ -562,16 +571,18 @@ import Testing UsageURLProtocol.enqueue(json: try provider.toolStream()) let session = provider.makeSession(tools: [RecordingWeatherTool()]) session.toolExecutionDelegate = StopDelegate() - var snapshotCount = 0 + var skies: [Sky?] = [] var thrownError: (any Error)? do { - for try await _ in session.streamResponse(to: "Weather?", generating: Forecast.self) { - snapshotCount += 1 + for try await snapshot in session.streamResponse(to: "Weather?", generating: Forecast.self) { + skies.append(snapshot.content.sky) } } catch { thrownError = error } - #expect(snapshotCount == 0) + // The round's usage arrives in snapshots without content before the stream throws. + #expect(!skies.isEmpty) + #expect(skies.allSatisfy { $0 == nil }) #expect(thrownError as? GeneratedContentError == .typeMismatch) } @@ -786,6 +797,46 @@ import Testing #expect(collected.usage == provider.expected) } + @Generable + struct Article { + var title: String + var body: String + } + + /// A structured stream sends a snapshot for each chunk, + /// before every required property has arrived. + @Test(arguments: Provider.allCases) + func structuredStreamsSendPartialSnapshots(_ provider: Provider) async throws { + UsageURLProtocol.reset() + let chunks = ["{\"title\":\"Hel", "lo\",", "\"body\":\"Wo", "rld\"}"] + UsageURLProtocol.enqueue(json: try provider.stream(chunks: chunks)) + var progression: [String] = [] + var lastUsage: LanguageModelSession.Usage? + for try await snapshot in provider.makeSession().streamResponse(to: "Hi", generating: Article.self) { + let content: Article.PartiallyGenerated = snapshot.content + lastUsage = snapshot.usage + guard content.title != nil || content.body != nil else { continue } + let step = "\(content.title ?? "-")/\(content.body ?? "-")" + if progression.last != step { progression.append(step) } + } + #expect(progression == ["Hel/-", "Hello/-", "Hello/Wo", "Hello/World"]) + #expect(lastUsage == provider.expected) + } + + /// A usage update that arrives before any structured content still sends a snapshot. + @Test(arguments: Provider.allCases) + func structuredStreamsSendUsageOnlySnapshots(_ provider: Provider) async throws { + UsageURLProtocol.reset() + UsageURLProtocol.enqueue(json: try provider.stream(text: "")) + var usages: [LanguageModelSession.Usage] = [] + do { + for try await snapshot in provider.makeSession().streamResponse(to: "Hi", generating: Article.self) { + usages.append(snapshot.usage) + } + } catch {} + #expect(usages.last == provider.expected) + } + @Test(arguments: [Provider.chat, .responses, .openResponses, .gemini]) func rawGeneratedContentPreservesUnquotedText(_ provider: Provider) async throws { UsageURLProtocol.reset() From fbc099df6746da63b6102c56111e2d502ebd1449 Mon Sep 17 00:00:00 2001 From: mattt Date: Thu, 22 Oct 2026 07:23:25 -0700 Subject: [PATCH 2/2] Decode partial snapshots in a generic function for Swift 6.3 --- Sources/AnyLanguageModel/Shared/ResponseUsage.swift | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/Sources/AnyLanguageModel/Shared/ResponseUsage.swift b/Sources/AnyLanguageModel/Shared/ResponseUsage.swift index 58fa1e47..047bc69d 100644 --- a/Sources/AnyLanguageModel/Shared/ResponseUsage.swift +++ b/Sources/AnyLanguageModel/Shared/ResponseUsage.swift @@ -222,8 +222,19 @@ extension LanguageModelSession.ResponseStream.Snapshot { (try? GeneratedContent(json: text)) ?? GeneratedContent(text), fromProperty: wrapperProperty ) - guard let content = try? Content.PartiallyGenerated(raw) else { return nil } + guard let content = partiallyGenerated(Content.self, from: raw) else { return nil } self.init(content: content, rawContent: raw, requests: requests) } } } + +/// Decodes partially generated content, or returns `nil`. +/// +/// A generic function, because Swift 6.3 can't type-check +/// this initializer inside the snapshot's extension. +private func partiallyGenerated( + _: Value.Type, + from raw: GeneratedContent +) -> Value.PartiallyGenerated? { + try? Value.PartiallyGenerated(raw) +}