From e49db05f6de79dc0b9dc1f93a0be7dead5530da5 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 06:09:58 -0700 Subject: [PATCH 1/7] Conform various types to Codable --- .../AnyLanguageModel/GeneratedContent.swift | 85 ++++++++++++++++++- Sources/AnyLanguageModel/GenerationID.swift | 2 +- .../AnyLanguageModel/GenerationOptions.swift | 55 +++++++++++- Sources/AnyLanguageModel/Transcript.swift | 26 +++--- .../MockLanguageModelTests.swift | 4 +- 5 files changed, 153 insertions(+), 19 deletions(-) diff --git a/Sources/AnyLanguageModel/GeneratedContent.swift b/Sources/AnyLanguageModel/GeneratedContent.swift index b54cffa0..88f473b1 100644 --- a/Sources/AnyLanguageModel/GeneratedContent.swift +++ b/Sources/AnyLanguageModel/GeneratedContent.swift @@ -4,7 +4,7 @@ import CoreFoundation /// A type that represents structured, generated content. /// /// Generated content may contain a single value, an array, or key-value pairs with unique keys. -public struct GeneratedContent: Sendable, Equatable, Generable, CustomDebugStringConvertible { +public struct GeneratedContent: Sendable, Equatable, Generable, CustomDebugStringConvertible, Codable { /// An instance of the generation schema. public static var generationSchema: GenerationSchema { // GeneratedContent is self-describing, it doesn't have a fixed schema @@ -391,3 +391,86 @@ public enum GeneratedContentError: Error { case typeMismatch case neverCannotBeInstantiated } + +// MARK: - Codable + +extension GeneratedContent { + private enum CodingKeys: String, CodingKey { + case id + case kind + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.id = try container.decodeIfPresent(GenerationID.self, forKey: .id) + self.kind = try container.decode(Kind.self, forKey: .kind) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encodeIfPresent(id, forKey: .id) + try container.encode(kind, forKey: .kind) + } +} + +extension GeneratedContent.Kind: Codable { + private enum CodingKeys: String, CodingKey { + case type + case value + case properties + case orderedKeys + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let type = try container.decode(String.self, forKey: .type) + + switch type { + case "null": + self = .null + case "bool": + self = .bool(try container.decode(Bool.self, forKey: .value)) + case "number": + self = .number(try container.decode(Double.self, forKey: .value)) + case "string": + self = .string(try container.decode(String.self, forKey: .value)) + case "array": + self = .array(try container.decode([GeneratedContent].self, forKey: .value)) + case "structure": + let properties = try container.decode([String: GeneratedContent].self, forKey: .properties) + let orderedKeys = try container.decode([String].self, forKey: .orderedKeys) + self = .structure(properties: properties, orderedKeys: orderedKeys) + default: + throw DecodingError.dataCorruptedError( + forKey: .type, + in: container, + debugDescription: "Unknown kind type: \(type)" + ) + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + + switch self { + case .null: + try container.encode("null", forKey: .type) + case .bool(let value): + try container.encode("bool", forKey: .type) + try container.encode(value, forKey: .value) + case .number(let value): + try container.encode("number", forKey: .type) + try container.encode(value, forKey: .value) + case .string(let value): + try container.encode("string", forKey: .type) + try container.encode(value, forKey: .value) + case .array(let elements): + try container.encode("array", forKey: .type) + try container.encode(elements, forKey: .value) + case .structure(let properties, let orderedKeys): + try container.encode("structure", forKey: .type) + try container.encode(properties, forKey: .properties) + try container.encode(orderedKeys, forKey: .orderedKeys) + } + } +} diff --git a/Sources/AnyLanguageModel/GenerationID.swift b/Sources/AnyLanguageModel/GenerationID.swift index e70692c1..08816b26 100644 --- a/Sources/AnyLanguageModel/GenerationID.swift +++ b/Sources/AnyLanguageModel/GenerationID.swift @@ -41,7 +41,7 @@ import struct Foundation.UUID /// } /// } /// ``` -public struct GenerationID: Sendable, Hashable { +public struct GenerationID: Sendable, Hashable, Codable { private let uuid: UUID /// Create a new, unique `GenerationID`. diff --git a/Sources/AnyLanguageModel/GenerationOptions.swift b/Sources/AnyLanguageModel/GenerationOptions.swift index c291cd3e..c0f47e6b 100644 --- a/Sources/AnyLanguageModel/GenerationOptions.swift +++ b/Sources/AnyLanguageModel/GenerationOptions.swift @@ -5,7 +5,7 @@ /// perform various adjustments on how the model chooses output tokens, /// to specify the penalties for repeating tokens or generating /// longer responses. -public struct GenerationOptions: Sendable, Equatable { +public struct GenerationOptions: Sendable, Equatable, Codable { /// A sampling strategy for how the model picks tokens when generating a /// response. @@ -83,7 +83,7 @@ extension GenerationOptions { /// loop the model produces a probability distribution for all the tokens in its /// vocabulary. The sampling mode controls how a token is selected from that /// distribution. - public struct SamplingMode: Sendable, Equatable { + public struct SamplingMode: Sendable, Equatable, Codable { enum Mode: Equatable { case greedy case topK(Int, seed: UInt64?) @@ -148,3 +148,54 @@ extension GenerationOptions { } } } + +// MARK: - Codable + +extension GenerationOptions.SamplingMode { + private enum CodingKeys: String, CodingKey { + case type + case value + case seed + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let type = try container.decode(String.self, forKey: .type) + + switch type { + case "greedy": + self.mode = .greedy + case "topK": + let k = try container.decode(Int.self, forKey: .value) + let seed = try container.decodeIfPresent(UInt64.self, forKey: .seed) + self.mode = .topK(k, seed: seed) + case "nucleus": + let threshold = try container.decode(Double.self, forKey: .value) + let seed = try container.decodeIfPresent(UInt64.self, forKey: .seed) + self.mode = .nucleus(threshold, seed: seed) + default: + throw DecodingError.dataCorruptedError( + forKey: .type, + in: container, + debugDescription: "Unknown sampling mode: \(type)" + ) + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + + switch mode { + case .greedy: + try container.encode("greedy", forKey: .type) + case .topK(let k, let seed): + try container.encode("topK", forKey: .type) + try container.encode(k, forKey: .value) + try container.encodeIfPresent(seed, forKey: .seed) + case .nucleus(let threshold, let seed): + try container.encode("nucleus", forKey: .type) + try container.encode(threshold, forKey: .value) + try container.encodeIfPresent(seed, forKey: .seed) + } + } +} diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index 04392ba0..bfd7bc82 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/Sources/AnyLanguageModel/Transcript.swift @@ -1,7 +1,7 @@ import struct Foundation.UUID /// A type that represents a conversation history between a user and a language model. -public struct Transcript: Sendable, Equatable { +public struct Transcript: Sendable, Equatable, Codable { private var entries: [Entry] /// Creates a transcript. @@ -13,7 +13,7 @@ public struct Transcript: Sendable, Equatable { } /// An entry in a transcript. - public enum Entry: Sendable, Identifiable, Equatable { + public enum Entry: Sendable, Identifiable, Equatable, Codable { /// Instructions, typically provided by you, the developer. case instructions(Instructions) @@ -47,7 +47,7 @@ public struct Transcript: Sendable, Equatable { } /// The types of segments that may be included in a transcript entry. - public enum Segment: Sendable, Identifiable, Equatable { + public enum Segment: Sendable, Identifiable, Equatable, Codable { /// A segment containing text. case text(TextSegment) @@ -66,7 +66,7 @@ public struct Transcript: Sendable, Equatable { } /// A segment containing text. - public struct TextSegment: Sendable, Identifiable, Equatable { + public struct TextSegment: Sendable, Identifiable, Equatable, Codable { /// The stable identity of the entity associated with this instance. public var id: String @@ -79,7 +79,7 @@ public struct Transcript: Sendable, Equatable { } /// A segment containing structured content. - public struct StructuredSegment: Sendable, Identifiable, Equatable { + public struct StructuredSegment: Sendable, Identifiable, Equatable, Codable { /// The stable identity of the entity associated with this instance. public var id: String @@ -101,7 +101,7 @@ public struct Transcript: Sendable, Equatable { /// Instructions are typically provided to define the role and behavior of the model. Apple trains the model /// to obey instructions over any commands it receives in prompts. This is a security mechanism to help /// mitigate prompt injection attacks. - public struct Instructions: Sendable, Identifiable, Equatable { + public struct Instructions: Sendable, Identifiable, Equatable, Codable { /// The stable identity of the entity associated with this instance. public var id: String @@ -133,7 +133,7 @@ public struct Transcript: Sendable, Equatable { } /// A prompt from the user asking the model. - public struct Prompt: Sendable, Identifiable, Equatable { + public struct Prompt: Sendable, Identifiable, Equatable, Codable { /// The identifier of the prompt. public var id: String @@ -167,7 +167,7 @@ public struct Transcript: Sendable, Equatable { } /// Specifies a response format that the model must conform its output to. - public struct ResponseFormat: Sendable { + public struct ResponseFormat: Sendable, Codable { private let schema: GenerationSchema /// A name associated with the response format. @@ -202,7 +202,7 @@ public struct Transcript: Sendable, Equatable { } /// A collection tool calls generated by the model. - public struct ToolCalls: Sendable, Identifiable, Equatable { + public struct ToolCalls: Sendable, Identifiable, Equatable, Codable { /// The stable identity of the entity associated with this instance. public var id: String @@ -216,7 +216,7 @@ public struct Transcript: Sendable, Equatable { } /// A tool call generated by the model containing the name of a tool and arguments to pass to it. - public struct ToolCall: Sendable, Identifiable, Equatable { + public struct ToolCall: Sendable, Identifiable, Equatable, Codable { /// The stable identity of the entity associated with this instance. public var id: String @@ -234,7 +234,7 @@ public struct Transcript: Sendable, Equatable { } /// A tool output provided back to the model. - public struct ToolOutput: Sendable, Identifiable, Equatable { + public struct ToolOutput: Sendable, Identifiable, Equatable, Codable { /// A unique id for this tool output. public var id: String @@ -252,7 +252,7 @@ public struct Transcript: Sendable, Equatable { } /// A response from the model. - public struct Response: Sendable, Identifiable, Equatable { + public struct Response: Sendable, Identifiable, Equatable, Codable { /// The stable identity of the entity associated with this instance. public var id: String @@ -274,7 +274,7 @@ public struct Transcript: Sendable, Equatable { } /// A definition of a tool. - public struct ToolDefinition: Sendable { + public struct ToolDefinition: Sendable, Codable { /// The tool's name. public var name: String diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index 92fc8be5..9d8b0173 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -84,7 +84,7 @@ struct MockLanguageModelTests { #expect(session.isResponding == false) let stream = session.streamResponse(to: "Test") - + // Start consuming the stream in a task let task = Task { for try await _ in stream { @@ -98,7 +98,7 @@ struct MockLanguageModelTests { // Wait for stream to complete _ = try await task.value - + // Give time for endResponding to complete try await Task.sleep(for: .milliseconds(10)) #expect(session.isResponding == false) From 7ed47acd0d6c59432922ca0958b88afa1ba4c2cc Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 07:19:07 -0700 Subject: [PATCH 2/7] Implement transcript bookkeeping for non-streaming responses --- .../LanguageModelSession.swift | 74 ++++++++++++++----- Sources/AnyLanguageModel/Transcript.swift | 10 +++ .../MockLanguageModelTests.swift | 73 ++++++++++++++++++ .../Shared/MockLanguageModel.swift | 9 ++- 4 files changed, 145 insertions(+), 21 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 965a46f2..d92e5c9f 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -53,7 +53,28 @@ public final class LanguageModelSession: @unchecked Sendable { self.model = model self.tools = tools self.instructions = instructions - self.transcript = transcript + + // Build transcript with instructions if provided and not already in transcript + var finalTranscript = transcript + if let instructions = instructions { + // Only add instructions if transcript doesn't already start with instructions + let hasInstructions = + finalTranscript.first.map { entry in + if case .instructions = entry { return true } else { return false } + } ?? false + + if !hasInstructions { + let instructionsEntry = Transcript.Entry.instructions( + Transcript.Instructions( + segments: [.text(.init(content: instructions.description))], + toolDefinitions: tools.map { Transcript.ToolDefinition(tool: $0) } + ) + ) + finalTranscript.append(instructionsEntry) + } + } + + self.transcript = finalTranscript } public func prewarm(promptPrefix: Prompt? = nil) { @@ -121,15 +142,12 @@ public final class LanguageModelSession: @unchecked Sendable { to prompt: Prompt, options: GenerationOptions = GenerationOptions() ) async throws -> Response { - try await wrapRespond { - try await model.respond( - within: self, - to: prompt, - generating: String.self, - includeSchemaInPrompt: true, - options: options - ) - } + try await respond( + to: prompt, + generating: String.self, + includeSchemaInPrompt: true, + options: options + ) } @discardableResult @@ -155,15 +173,12 @@ public final class LanguageModelSession: @unchecked Sendable { includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) async throws -> Response { - try await wrapRespond { - try await model.respond( - within: self, - to: prompt, - generating: GeneratedContent.self, - includeSchemaInPrompt: includeSchemaInPrompt, - options: options - ) - } + try await respond( + to: prompt, + generating: GeneratedContent.self, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) } @discardableResult @@ -204,13 +219,32 @@ public final class LanguageModelSession: @unchecked Sendable { options: GenerationOptions = GenerationOptions() ) async throws -> Response where Content: Generable { try await wrapRespond { - try await model.respond( + // Add prompt to transcript + let promptEntry = Transcript.Entry.prompt( + Transcript.Prompt( + segments: [.text(.init(content: prompt.description))], + options: options, + responseFormat: nil + ) + ) + await MainActor.run { + self.transcript.append(promptEntry) + } + + let response = try await model.respond( within: self, to: prompt, generating: type, includeSchemaInPrompt: includeSchemaInPrompt, options: options ) + + // Add response entries to transcript + await MainActor.run { + self.transcript.append(contentsOf: response.transcriptEntries) + } + + return response } } diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index bfd7bc82..69b054b6 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/Sources/AnyLanguageModel/Transcript.swift @@ -12,6 +12,16 @@ public struct Transcript: Sendable, Equatable, Codable { self.entries = Array(entries) } + /// Appends a single entry to the transcript. + mutating func append(_ entry: Entry) { + entries.append(entry) + } + + /// Appends multiple entries to the transcript. + mutating func append(contentsOf newEntries: S) where S: Sequence, S.Element == Entry { + entries.append(contentsOf: newEntries) + } + /// An entry in a transcript. public enum Entry: Sendable, Identifiable, Equatable, Codable { /// Instructions, typically provided by you, the developer. diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index 9d8b0173..c5384f16 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -103,4 +103,77 @@ struct MockLanguageModelTests { try await Task.sleep(for: .milliseconds(10)) #expect(session.isResponding == false) } + + @Test func transcriptStartsEmpty() async throws { + let model = MockLanguageModel.fixed("Hello") + let session = LanguageModelSession(model: model) + + #expect(session.transcript.count == 0) + } + + @Test func transcriptRecordsPromptAndResponse() async throws { + let model = MockLanguageModel.fixed("Hello, World!") + let session = LanguageModelSession(model: model) + + #expect(session.transcript.count == 0) + + let response = try await session.respond(to: "Say hello") + + let entries = Array(session.transcript) + #expect(entries.count > 0) + #expect(response.transcriptEntries.count > 0) + + #expect(response.content == "Hello, World!") + } + + @Test func transcriptGrowsWithMultipleInteractions() async throws { + let model = MockLanguageModel.echo + let session = LanguageModelSession(model: model) + + #expect(session.transcript.count == 0) + + try await session.respond(to: "First prompt") + let countAfterFirst = session.transcript.count + #expect(countAfterFirst > 0) + + try await session.respond(to: "Second prompt") + let countAfterSecond = session.transcript.count + #expect(countAfterSecond > countAfterFirst) + + try await session.respond(to: "Third prompt") + let countAfterThird = session.transcript.count + #expect(countAfterThird > countAfterSecond) + } + + @Test func transcriptIncludesInstructions() async throws { + let model = MockLanguageModel.fixed("Response") + let instructions = Instructions("Be helpful") + let session = LanguageModelSession( + model: model, + instructions: instructions + ) + + let entries = Array(session.transcript) + #expect(entries.count > 0) + + if case .instructions(let transcriptInstructions) = entries.first { + #expect(transcriptInstructions.segments.count > 0) + } else { + Issue.record("First transcript entry should be instructions") + } + } + + @Test func transcriptEntriesAreIdentifiable() async throws { + let model = MockLanguageModel.fixed("Response") + let session = LanguageModelSession(model: model) + + try await session.respond(to: "Test") + + let entries = Array(session.transcript) + #expect(entries.count > 0) + + for entry in entries { + #expect(!entry.id.isEmpty) + } + } } diff --git a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift index bcb2bd6f..3c96cf45 100644 --- a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift +++ b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift @@ -36,10 +36,17 @@ struct MockLanguageModel: LanguageModel { let promptWithInstructions = Prompt("Instructions: \(session.instructions?.description ?? "N/A")\n\(prompt)") let text = try await responseProvider(promptWithInstructions, options) + let responseEntry = Transcript.Entry.response( + Transcript.Response( + assetIDs: [], + segments: [.text(.init(content: text))] + ) + ) + return LanguageModelSession.Response( content: text as! Content, rawContent: GeneratedContent(text), - transcriptEntries: [] + transcriptEntries: [responseEntry] ) } From fee702e5497c5810813734d7a205c6bf281cff78 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 07:25:54 -0700 Subject: [PATCH 3/7] Implement transcript bookkeeping for streaming responses --- .../LanguageModelSession.swift | 71 ++++++++++++++----- .../MockLanguageModelTests.swift | 33 +++++++++ 2 files changed, 85 insertions(+), 19 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index d92e5c9f..3818050a 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -110,18 +110,47 @@ public final class LanguageModelSession: @unchecked Sendable { } nonisolated private func wrapStream( - _ upstream: sending ResponseStream + _ upstream: sending ResponseStream, + promptEntry: Transcript.Entry ) -> ResponseStream where Content: Generable, Content.PartiallyGenerated: Sendable { let session = self let relay = AsyncThrowingStream.Snapshot, any Error> { continuation in let stream = upstream Task { + // Add prompt to transcript when stream starts + await MainActor.run { + session.transcript.append(promptEntry) + } + await session.beginResponding() + var lastSnapshot: ResponseStream.Snapshot? do { for try await snapshot in stream { + lastSnapshot = snapshot continuation.yield(snapshot) } continuation.finish() + + // Add response to transcript after stream completes + if let lastSnapshot { + // Extract text content from the generated content + let textContent: String + if case .string(let str) = lastSnapshot.rawContent.kind { + textContent = str + } else { + textContent = lastSnapshot.rawContent.jsonString + } + + let responseEntry = Transcript.Entry.response( + Transcript.Response( + assetIDs: [], + segments: [.text(.init(content: textContent))] + ) + ) + await MainActor.run { + session.transcript.append(responseEntry) + } + } } catch { continuation.finish(throwing: error) } @@ -284,14 +313,11 @@ public final class LanguageModelSession: @unchecked Sendable { includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) -> sending ResponseStream { - wrapStream( - model.streamResponse( - within: self, - to: prompt, - generating: GeneratedContent.self, - includeSchemaInPrompt: includeSchemaInPrompt, - options: options - ) + streamResponse( + to: prompt, + generating: GeneratedContent.self, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options ) } @@ -324,14 +350,24 @@ public final class LanguageModelSession: @unchecked Sendable { includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) -> sending ResponseStream where Content: Generable { - wrapStream( + // Create prompt entry that will be added when stream starts + let promptEntry = Transcript.Entry.prompt( + Transcript.Prompt( + segments: [.text(.init(content: prompt.description))], + options: options, + responseFormat: nil + ) + ) + + return wrapStream( model.streamResponse( within: self, to: prompt, generating: type, includeSchemaInPrompt: includeSchemaInPrompt, options: options - ) + ), + promptEntry: promptEntry ) } @@ -367,14 +403,11 @@ public final class LanguageModelSession: @unchecked Sendable { to prompt: Prompt, options: GenerationOptions = GenerationOptions() ) -> sending ResponseStream { - wrapStream( - model.streamResponse( - within: self, - to: prompt, - generating: String.self, - includeSchemaInPrompt: true, - options: options - ) + streamResponse( + to: prompt, + generating: String.self, + includeSchemaInPrompt: true, + options: options ) } diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index c5384f16..b3de180a 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -176,4 +176,37 @@ struct MockLanguageModelTests { #expect(!entry.id.isEmpty) } } + + @Test func streamingRecordsToTranscript() async throws { + let model = MockLanguageModel.fixed("Streamed response") + let session = LanguageModelSession(model: model) + + #expect(session.transcript.count == 0) + + let stream = session.streamResponse(to: "Stream this") + + // Consume the stream + for try await _ in stream {} + + // Give time for transcript update to complete + try await Task.sleep(for: .milliseconds(10)) + + // Verify transcript has both prompt and response + let entries = Array(session.transcript) + #expect(entries.count == 2) + + // First entry should be prompt + if case .prompt(let prompt) = entries[0] { + #expect(prompt.segments.count > 0) + } else { + Issue.record("First entry should be prompt") + } + + // Second entry should be response + if case .response(let response) = entries[1] { + #expect(response.segments.count > 0) + } else { + Issue.record("Second entry should be response") + } + } } From 8a19df6bf314f6ee97deb146f43086612595ab2f Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 07:30:00 -0700 Subject: [PATCH 4/7] Consolidate mock tests --- .../MockLanguageModelTests.swift | 175 ++++++------------ 1 file changed, 56 insertions(+), 119 deletions(-) diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index b3de180a..5608d046 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -4,21 +4,30 @@ import Testing @Suite("MockLanguageModel") struct MockLanguageModelTests { - @Test func fixed() async throws { + @Test func fixedResponse() async throws { let model = MockLanguageModel.fixed("Hello, World!") let session = LanguageModelSession(model: model) + #expect(session.transcript.count == 0) + let response = try await session.respond(to: "Say hello") #expect(response.content == "Hello, World!") + + // Verify transcript was updated + #expect(session.transcript.count == 2) + #expect(response.transcriptEntries.count > 0) } - @Test func echo() async throws { + @Test func echoResponse() async throws { let model = MockLanguageModel.echo let session = LanguageModelSession(model: model) let prompt = Prompt("Echo this") let response = try await session.respond(to: prompt) #expect(response.content.contains(prompt.description)) + + // Verify transcript + #expect(session.transcript.count == 2) } @Test func withInstructions() async throws { @@ -34,18 +43,30 @@ struct MockLanguageModelTests { return "😐" } - for (prompt, expected) in [ + for (instructionText, expected) in [ ("Be helpful", "😇"), ("Be evil", "😈"), ("Meh", "😐"), ] { let session = LanguageModelSession( model: model, - instructions: Instructions(prompt) + instructions: Instructions(instructionText) ) + // Verify instructions are in transcript + let entriesBeforeResponse = Array(session.transcript) + #expect(entriesBeforeResponse.count == 1) + if case .instructions(let transcriptInstructions) = entriesBeforeResponse.first { + #expect(transcriptInstructions.segments.count > 0) + } else { + Issue.record("First entry should be instructions") + } + let response = try await session.respond(to: "Do what you want") #expect(response.content == expected) + + // Verify transcript has instructions, prompt, and response + #expect(session.transcript.count == 3) } } @@ -56,74 +77,50 @@ struct MockLanguageModelTests { #expect(model.isAvailable == false) } - @Test func isRespondingDuringAsyncResponse() async throws { - let model = MockLanguageModel { _, _ in + @Test func streamingResponse() async throws { + // Test async response with isResponding state + let asyncModel = MockLanguageModel { _, _ in try await Task.sleep(for: .milliseconds(100)) - return "Response" + return "Async Response" } - let session = LanguageModelSession(model: model) + let asyncSession = LanguageModelSession(model: asyncModel) - #expect(session.isResponding == false) + #expect(asyncSession.isResponding == false) + #expect(asyncSession.transcript.count == 0) - let task = Task { - try await session.respond(to: "Test") + let asyncTask = Task { + try await asyncSession.respond(to: "Async test") } try await Task.sleep(for: .milliseconds(50)) - #expect(session.isResponding == true) + #expect(asyncSession.isResponding == true) - _ = try await task.value + let response = try await asyncTask.value try await Task.sleep(for: .milliseconds(10)) - #expect(session.isResponding == false) - } + #expect(asyncSession.isResponding == false) + #expect(asyncSession.transcript.count == 2) + #expect(response.transcriptEntries.count > 0) - @Test func isRespondingDuringStreaming() async throws { - let model = MockLanguageModel.streamingMock() - let session = LanguageModelSession(model: model) + // Test streaming response with isResponding state + let streamModel = MockLanguageModel.streamingMock() + let streamSession = LanguageModelSession(model: streamModel) - #expect(session.isResponding == false) + #expect(streamSession.isResponding == false) + #expect(streamSession.transcript.count == 0) - let stream = session.streamResponse(to: "Test") + let stream = streamSession.streamResponse(to: "Stream test") - // Start consuming the stream in a task - let task = Task { - for try await _ in stream { - // Just consume the stream - } + let streamTask = Task { + for try await _ in stream {} } - // Give the streaming task time to start and call beginResponding try await Task.sleep(for: .milliseconds(50)) - #expect(session.isResponding == true) - - // Wait for stream to complete - _ = try await task.value + #expect(streamSession.isResponding == true) - // Give time for endResponding to complete + _ = try await streamTask.value try await Task.sleep(for: .milliseconds(10)) - #expect(session.isResponding == false) - } - - @Test func transcriptStartsEmpty() async throws { - let model = MockLanguageModel.fixed("Hello") - let session = LanguageModelSession(model: model) - - #expect(session.transcript.count == 0) - } - - @Test func transcriptRecordsPromptAndResponse() async throws { - let model = MockLanguageModel.fixed("Hello, World!") - let session = LanguageModelSession(model: model) - - #expect(session.transcript.count == 0) - - let response = try await session.respond(to: "Say hello") - - let entries = Array(session.transcript) - #expect(entries.count > 0) - #expect(response.transcriptEntries.count > 0) - - #expect(response.content == "Hello, World!") + #expect(streamSession.isResponding == false) + #expect(streamSession.transcript.count == 2) } @Test func transcriptGrowsWithMultipleInteractions() async throws { @@ -134,79 +131,19 @@ struct MockLanguageModelTests { try await session.respond(to: "First prompt") let countAfterFirst = session.transcript.count - #expect(countAfterFirst > 0) + #expect(countAfterFirst == 2) try await session.respond(to: "Second prompt") let countAfterSecond = session.transcript.count - #expect(countAfterSecond > countAfterFirst) + #expect(countAfterSecond == 4) try await session.respond(to: "Third prompt") let countAfterThird = session.transcript.count - #expect(countAfterThird > countAfterSecond) - } + #expect(countAfterThird == 6) - @Test func transcriptIncludesInstructions() async throws { - let model = MockLanguageModel.fixed("Response") - let instructions = Instructions("Be helpful") - let session = LanguageModelSession( - model: model, - instructions: instructions - ) - - let entries = Array(session.transcript) - #expect(entries.count > 0) - - if case .instructions(let transcriptInstructions) = entries.first { - #expect(transcriptInstructions.segments.count > 0) - } else { - Issue.record("First transcript entry should be instructions") - } - } - - @Test func transcriptEntriesAreIdentifiable() async throws { - let model = MockLanguageModel.fixed("Response") - let session = LanguageModelSession(model: model) - - try await session.respond(to: "Test") - - let entries = Array(session.transcript) - #expect(entries.count > 0) - - for entry in entries { + // Verify all entries are identifiable + for entry in session.transcript { #expect(!entry.id.isEmpty) } } - - @Test func streamingRecordsToTranscript() async throws { - let model = MockLanguageModel.fixed("Streamed response") - let session = LanguageModelSession(model: model) - - #expect(session.transcript.count == 0) - - let stream = session.streamResponse(to: "Stream this") - - // Consume the stream - for try await _ in stream {} - - // Give time for transcript update to complete - try await Task.sleep(for: .milliseconds(10)) - - // Verify transcript has both prompt and response - let entries = Array(session.transcript) - #expect(entries.count == 2) - - // First entry should be prompt - if case .prompt(let prompt) = entries[0] { - #expect(prompt.segments.count > 0) - } else { - Issue.record("First entry should be prompt") - } - - // Second entry should be response - if case .response(let response) = entries[1] { - #expect(response.segments.count > 0) - } else { - Issue.record("Second entry should be response") - } - } } From f7ee4d71977fd3bd0b842900e5041ef8ab7eab73 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 09:57:49 -0700 Subject: [PATCH 5/7] Simplify Codable implementation --- .../AnyLanguageModel/GenerationOptions.swift | 51 +------------------ 1 file changed, 1 insertion(+), 50 deletions(-) diff --git a/Sources/AnyLanguageModel/GenerationOptions.swift b/Sources/AnyLanguageModel/GenerationOptions.swift index c0f47e6b..80e0af18 100644 --- a/Sources/AnyLanguageModel/GenerationOptions.swift +++ b/Sources/AnyLanguageModel/GenerationOptions.swift @@ -84,7 +84,7 @@ extension GenerationOptions { /// vocabulary. The sampling mode controls how a token is selected from that /// distribution. public struct SamplingMode: Sendable, Equatable, Codable { - enum Mode: Equatable { + enum Mode: Equatable, Codable { case greedy case topK(Int, seed: UInt64?) case nucleus(Double, seed: UInt64?) @@ -149,53 +149,4 @@ extension GenerationOptions { } } -// MARK: - Codable -extension GenerationOptions.SamplingMode { - private enum CodingKeys: String, CodingKey { - case type - case value - case seed - } - - public init(from decoder: Decoder) throws { - let container = try decoder.container(keyedBy: CodingKeys.self) - let type = try container.decode(String.self, forKey: .type) - - switch type { - case "greedy": - self.mode = .greedy - case "topK": - let k = try container.decode(Int.self, forKey: .value) - let seed = try container.decodeIfPresent(UInt64.self, forKey: .seed) - self.mode = .topK(k, seed: seed) - case "nucleus": - let threshold = try container.decode(Double.self, forKey: .value) - let seed = try container.decodeIfPresent(UInt64.self, forKey: .seed) - self.mode = .nucleus(threshold, seed: seed) - default: - throw DecodingError.dataCorruptedError( - forKey: .type, - in: container, - debugDescription: "Unknown sampling mode: \(type)" - ) - } - } - - public func encode(to encoder: Encoder) throws { - var container = encoder.container(keyedBy: CodingKeys.self) - - switch mode { - case .greedy: - try container.encode("greedy", forKey: .type) - case .topK(let k, let seed): - try container.encode("topK", forKey: .type) - try container.encode(k, forKey: .value) - try container.encodeIfPresent(seed, forKey: .seed) - case .nucleus(let threshold, let seed): - try container.encode("nucleus", forKey: .type) - try container.encode(threshold, forKey: .value) - try container.encodeIfPresent(seed, forKey: .seed) - } - } -} From 65ee0af709d05b2fc6cd17138353c19a5eaf1630 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 10:01:05 -0700 Subject: [PATCH 6/7] Inline Transcript.append methods --- Sources/AnyLanguageModel/GenerationOptions.swift | 2 -- Sources/AnyLanguageModel/LanguageModelSession.swift | 10 +++++----- Sources/AnyLanguageModel/Transcript.swift | 12 +----------- 3 files changed, 6 insertions(+), 18 deletions(-) diff --git a/Sources/AnyLanguageModel/GenerationOptions.swift b/Sources/AnyLanguageModel/GenerationOptions.swift index 80e0af18..38bb4515 100644 --- a/Sources/AnyLanguageModel/GenerationOptions.swift +++ b/Sources/AnyLanguageModel/GenerationOptions.swift @@ -148,5 +148,3 @@ extension GenerationOptions { } } } - - diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 3818050a..bddc3158 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -70,7 +70,7 @@ public final class LanguageModelSession: @unchecked Sendable { toolDefinitions: tools.map { Transcript.ToolDefinition(tool: $0) } ) ) - finalTranscript.append(instructionsEntry) + finalTranscript.entries.append(instructionsEntry) } } @@ -119,7 +119,7 @@ public final class LanguageModelSession: @unchecked Sendable { Task { // Add prompt to transcript when stream starts await MainActor.run { - session.transcript.append(promptEntry) + session.transcript.entries.append(promptEntry) } await session.beginResponding() @@ -148,7 +148,7 @@ public final class LanguageModelSession: @unchecked Sendable { ) ) await MainActor.run { - session.transcript.append(responseEntry) + session.transcript.entries.append(responseEntry) } } } catch { @@ -257,7 +257,7 @@ public final class LanguageModelSession: @unchecked Sendable { ) ) await MainActor.run { - self.transcript.append(promptEntry) + self.transcript.entries.append(promptEntry) } let response = try await model.respond( @@ -270,7 +270,7 @@ public final class LanguageModelSession: @unchecked Sendable { // Add response entries to transcript await MainActor.run { - self.transcript.append(contentsOf: response.transcriptEntries) + self.transcript.entries.append(contentsOf: response.transcriptEntries) } return response diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index 69b054b6..c3216e9f 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/Sources/AnyLanguageModel/Transcript.swift @@ -2,7 +2,7 @@ import struct Foundation.UUID /// A type that represents a conversation history between a user and a language model. public struct Transcript: Sendable, Equatable, Codable { - private var entries: [Entry] + var entries: [Entry] /// Creates a transcript. /// @@ -12,16 +12,6 @@ public struct Transcript: Sendable, Equatable, Codable { self.entries = Array(entries) } - /// Appends a single entry to the transcript. - mutating func append(_ entry: Entry) { - entries.append(entry) - } - - /// Appends multiple entries to the transcript. - mutating func append(contentsOf newEntries: S) where S: Sequence, S.Element == Entry { - entries.append(contentsOf: newEntries) - } - /// An entry in a transcript. public enum Entry: Sendable, Identifiable, Equatable, Codable { /// Instructions, typically provided by you, the developer. From b7489d907147fa00a81b367cb2898af8fdc20f09 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Tue, 28 Oct 2025 10:15:02 -0700 Subject: [PATCH 7/7] Revert "Inline Transcript.append methods" This reverts commit 65ee0af709d05b2fc6cd17138353c19a5eaf1630. --- Sources/AnyLanguageModel/LanguageModelSession.swift | 10 +++++----- Sources/AnyLanguageModel/Transcript.swift | 12 +++++++++++- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index bddc3158..3818050a 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -70,7 +70,7 @@ public final class LanguageModelSession: @unchecked Sendable { toolDefinitions: tools.map { Transcript.ToolDefinition(tool: $0) } ) ) - finalTranscript.entries.append(instructionsEntry) + finalTranscript.append(instructionsEntry) } } @@ -119,7 +119,7 @@ public final class LanguageModelSession: @unchecked Sendable { Task { // Add prompt to transcript when stream starts await MainActor.run { - session.transcript.entries.append(promptEntry) + session.transcript.append(promptEntry) } await session.beginResponding() @@ -148,7 +148,7 @@ public final class LanguageModelSession: @unchecked Sendable { ) ) await MainActor.run { - session.transcript.entries.append(responseEntry) + session.transcript.append(responseEntry) } } } catch { @@ -257,7 +257,7 @@ public final class LanguageModelSession: @unchecked Sendable { ) ) await MainActor.run { - self.transcript.entries.append(promptEntry) + self.transcript.append(promptEntry) } let response = try await model.respond( @@ -270,7 +270,7 @@ public final class LanguageModelSession: @unchecked Sendable { // Add response entries to transcript await MainActor.run { - self.transcript.entries.append(contentsOf: response.transcriptEntries) + self.transcript.append(contentsOf: response.transcriptEntries) } return response diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index c3216e9f..69b054b6 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/Sources/AnyLanguageModel/Transcript.swift @@ -2,7 +2,7 @@ import struct Foundation.UUID /// A type that represents a conversation history between a user and a language model. public struct Transcript: Sendable, Equatable, Codable { - var entries: [Entry] + private var entries: [Entry] /// Creates a transcript. /// @@ -12,6 +12,16 @@ public struct Transcript: Sendable, Equatable, Codable { self.entries = Array(entries) } + /// Appends a single entry to the transcript. + mutating func append(_ entry: Entry) { + entries.append(entry) + } + + /// Appends multiple entries to the transcript. + mutating func append(contentsOf newEntries: S) where S: Sequence, S.Element == Entry { + entries.append(contentsOf: newEntries) + } + /// An entry in a transcript. public enum Entry: Sendable, Identifiable, Equatable, Codable { /// Instructions, typically provided by you, the developer.