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..38bb4515 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,8 +83,8 @@ 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 { - enum Mode: Equatable { + public struct SamplingMode: Sendable, Equatable, Codable { + enum Mode: Equatable, Codable { case greedy case topK(Int, seed: UInt64?) case nucleus(Double, seed: UInt64?) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 965a46f2..3818050a 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) { @@ -89,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) } @@ -121,15 +171,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 +202,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 +248,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 } } @@ -250,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 ) } @@ -290,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 ) } @@ -333,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/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index 04392ba0..69b054b6 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. @@ -12,8 +12,18 @@ public struct Transcript: Sendable, Equatable { 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 { + public enum Entry: Sendable, Identifiable, Equatable, Codable { /// Instructions, typically provided by you, the developer. case instructions(Instructions) @@ -47,7 +57,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 +76,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 +89,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 +111,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 +143,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 +177,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 +212,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 +226,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 +244,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 +262,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 +284,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..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,51 +77,73 @@ 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") - - // Start consuming the stream in a task - let task = Task { - for try await _ in stream { - // Just consume the stream - } + let stream = streamSession.streamResponse(to: "Stream test") + + 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) + #expect(streamSession.isResponding == true) - // Wait for stream to complete - _ = try await task.value - - // Give time for endResponding to complete + _ = try await streamTask.value try await Task.sleep(for: .milliseconds(10)) - #expect(session.isResponding == false) + #expect(streamSession.isResponding == false) + #expect(streamSession.transcript.count == 2) + } + + @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 == 2) + + try await session.respond(to: "Second prompt") + let countAfterSecond = session.transcript.count + #expect(countAfterSecond == 4) + + try await session.respond(to: "Third prompt") + let countAfterThird = session.transcript.count + #expect(countAfterThird == 6) + + // Verify all entries are identifiable + for entry in session.transcript { + #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] ) }