diff --git a/README.md b/README.md index c6065abd..0e7b4486 100644 --- a/README.md +++ b/README.md @@ -387,19 +387,22 @@ pass the model's multimodal projector with `mmprojPath:`. ### Prompt attachments -Image attachments use the Foundation Models 27 prompt construction syntax: +Image attachments use the Foundation Models 27 prompt construction syntax, +including ordinary Tool outputs: ```swift -let prompt = Prompt { - "Describe this image" - Attachment(imageURL: fileURL) +func call(arguments: Arguments) async throws -> Prompt { + Prompt { + "Image read from the selected file" + Attachment(imageURL: fileURL) + } } ``` On Apple platforms, `Attachment(CGImage, orientation:)` retains typed image -content and orientation. Prompt builders preserve the order of text and images, -and the OS 26 compatibility transcript continues to store images using its -existing `Transcript.Segment.image` Codable representation. +content and orientation. Prompts, builders, and Tool outputs preserve the order +of text and images, and the OS 26 compatibility transcript continues to store +images using its existing `Transcript.Segment.image` Codable representation. `SystemLanguageModel` and `FoundationLanguageModel` bridge prompt attachments to Foundation Models only in apps built with Swift 6.4 or later @@ -409,9 +412,18 @@ they throw an error when a prompt has an attachment. The system bridge requires local file URLs; other provider adapters retain URL images in their existing transcript representation. -This first compatibility surface covers prompt `CGImage` and `imageURL` -attachments. Tool outputs, instruction attachments, labels, `CIImage`, pixel -buffers, and newer data attachments remain outside this change. +Ordinary Tool outputs read local image files once when the output is created, +and retain the image bytes and MIME type in the compatibility transcript. +Continuations and saved transcript replay therefore do not depend on the original +file remaining at its URL. Remote image URLs are retained without downloading. + +Gemini supports images in Tool function responses only with Gemini 3 and later. +In Gemini Tool outputs, remote image URLs are sent as text at their original +position without downloading the image; the transcript retains the image segment. + +This compatibility surface covers `CGImage` and `imageURL` attachments in +prompts and ordinary Tool outputs. Instruction attachments, labels, `CIImage`, +pixel buffers, and newer data attachments remain outside this change. ### Tool Calling @@ -895,6 +907,10 @@ say which API they follow. [image inputs](#image-inputs). Foundation Models 27-style [prompt attachments](#prompt-attachments) are also available for ordered text and image prompt construction. +- `Tool.makeOutputSegments(from:)`, exposed only through the `Compatibility` SPI: + an integration extension for back-deployment provider adapters to lower normal + Tool outputs into the compatibility transcript. It is not an Apple Tool API; + OS 27 canonical executors use Foundation Models' own Tool execution instead. - `ToolExecutionDelegate`, `ToolExecutionDecision`, and `toolExecutionDelegate`: [observing and controlling tool calls](#tool-calling). - `transcriptErrorHandlingPolicy` and `waitForResponseCompletion()`: diff --git a/Sources/AnyLanguageModel/DynamicProfile.swift b/Sources/AnyLanguageModel/DynamicProfile.swift new file mode 100644 index 00000000..f1cdd3d8 --- /dev/null +++ b/Sources/AnyLanguageModel/DynamicProfile.swift @@ -0,0 +1,652 @@ +extension LanguageModelSession { + /// A declarative, per-model-request session profile. + /// + /// - Note: This API is exclusive to AnyLanguageModel on OS 26 and mirrors + /// Foundation Models 27. + @_typeEraser(AnyDynamicProfile) + public protocol DynamicProfile { + associatedtype Body: DynamicProfile + + @DynamicProfileBuilder + var body: Body { get } + } + + /// A concrete profile containing dynamic instructions and Tools. + public struct Profile: DynamicProfile { + public typealias Body = Never + + fileprivate let dynamicInstructions: AnyDynamicInstructions + fileprivate let identity: String + + public init( + @DynamicInstructionsBuilder _ dynamicInstructions: () -> some DynamicInstructions + ) { + let value = dynamicInstructions() + self.dynamicInstructions = AnyDynamicInstructions(value) + self.identity = String(reflecting: type(of: value)) + } + + public var body: Never { + fatalError("LanguageModelSession.Profile has no body") + } + } + + /// A type-erased dynamic profile. + public struct AnyDynamicProfile: DynamicProfile { + public typealias Body = Never + + fileprivate let resolveValue: () -> ResolvedDynamicProfile + + public init(_ dynamicProfile: any DynamicProfile) { + let identity = String(reflecting: type(of: dynamicProfile)) + resolveValue = { + var resolved = resolveDynamicProfile(dynamicProfile) + resolved.identity = identity + "/" + resolved.identity + return resolved + } + } + + public init(erasing dynamicProfile: some DynamicProfile) { + self.init(dynamicProfile) + } + + public var body: Never { + fatalError("LanguageModelSession.AnyDynamicProfile has no body") + } + + func resolveForRequest() -> ResolvedDynamicProfile { + resolveValue() + } + } + + public struct ConditionalDynamicProfile: DynamicProfile + where TrueContent: DynamicProfile, FalseContent: DynamicProfile { + enum Branch { + case trueContent(TrueContent) + case falseContent(FalseContent) + } + + public typealias Body = Never + fileprivate let branch: Branch + + init(_ branch: Branch) { + self.branch = branch + } + + public var body: Never { + fatalError("ConditionalDynamicProfile has no body") + } + } + + @resultBuilder + public struct DynamicProfileBuilder { + public static func buildBlock(_ content: Content) -> Content + where Content: DynamicProfile { + content + } + + public static func buildEither( + first content: TrueContent + ) -> ConditionalDynamicProfile + where TrueContent: DynamicProfile, FalseContent: DynamicProfile { + ConditionalDynamicProfile(.trueContent(content)) + } + + public static func buildEither( + second content: FalseContent + ) -> ConditionalDynamicProfile + where TrueContent: DynamicProfile, FalseContent: DynamicProfile { + ConditionalDynamicProfile(.falseContent(content)) + } + + @available(*, unavailable, message: "The body of a 'DynamicProfile' must evaluate to a single active profile") + public static func buildBlock(_ contents: repeat each Content) -> Never + where repeat each Content: DynamicProfile { + fatalError("A DynamicProfile must have one active profile") + } + + public static func buildLimitedAvailability( + _ component: some DynamicProfile + ) -> AnyDynamicProfile { + AnyDynamicProfile(component) + } + } + + public protocol DynamicProfileModifier { + associatedtype Body: DynamicProfile + typealias Content = DynamicProfileModifierContent + + @DynamicProfileBuilder + func body(content: Content) -> Body + } + + public struct DynamicProfileModifierContent: DynamicProfile + where Modifier: DynamicProfileModifier { + public typealias Body = Never + fileprivate let content: AnyDynamicProfile + + fileprivate init(_ content: Content) where Content: DynamicProfile { + self.content = AnyDynamicProfile(content) + } + + public var body: Never { + fatalError("DynamicProfileModifierContent has no body") + } + } + + public struct ModifiedDynamicProfile: DynamicProfile + where Content: DynamicProfile, Modifier: DynamicProfileModifier { + public typealias Body = Never + fileprivate let content: Content + fileprivate let modifier: Modifier + + fileprivate init(content: Content, modifier: Modifier) { + self.content = content + self.modifier = modifier + } + + public var body: Never { + fatalError("ModifiedDynamicProfile has no body") + } + } +} + +extension LanguageModelSession.DynamicProfile { + public typealias Profile = LanguageModelSession.Profile + public typealias DynamicProfile = LanguageModelSession.DynamicProfile + public typealias SessionProperty = LanguageModelSession.SessionProperty + + public func modifier( + _ modifier: Modifier + ) -> some LanguageModelSession.DynamicProfile + where Modifier: LanguageModelSession.DynamicProfileModifier { + LanguageModelSession.ModifiedDynamicProfile(content: self, modifier: modifier) + } + + public func model( + _ model: any LanguageModel + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .model(model)) + } + + public func model( + _ model: Model + ) -> some LanguageModelSession.DynamicProfile where Model: LanguageModel { + BuiltinModifiedDynamicProfile(content: self, modifier: .model(model)) + } + + public func temperature(_ temperature: Double?) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .temperature(temperature)) + } + + public func samplingMode( + _ samplingMode: GenerationOptions.SamplingMode? + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .samplingMode(samplingMode)) + } + + public func maximumResponseTokens( + _ maximumResponseTokens: Int? + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile( + content: self, + modifier: .maximumResponseTokens(maximumResponseTokens) + ) + } + + public func reasoningLevel( + _ reasoningLevel: ContextOptions.ReasoningLevel? + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .reasoningLevel(reasoningLevel)) + } + + @_spi(Compatibility) + public func includeSchemaInPrompt( + _ includeSchemaInPrompt: Bool? + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile( + content: self, + modifier: .includeSchemaInPrompt(includeSchemaInPrompt) + ) + } + + public func toolCallingMode( + _ toolCallingMode: GenerationOptions.ToolCallingMode? + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .toolCallingMode(toolCallingMode)) + } + + public func historyTransform( + _ transform: @escaping ([Transcript.Entry]) -> [Transcript.Entry] + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .historyTransform(transform)) + } + + public func transcriptErrorHandlingPolicy( + _ policy: TranscriptErrorHandlingPolicy? + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .transcriptErrorHandlingPolicy(policy)) + } + + public func onPrompt( + @_inheritActorContext perform action: + nonisolated(nonsending) sending @escaping ( + Transcript.Prompt + ) async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .onPrompt(action)) + } + + public func onPrompt( + @_inheritActorContext perform action: nonisolated(nonsending) sending @escaping () async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + onPrompt { _ in try await action() } + } + + public func onResponse( + @_inheritActorContext perform action: + nonisolated(nonsending) sending @escaping ( + Transcript.Response + ) async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .onResponse(action)) + } + + public func onResponse( + @_inheritActorContext perform action: nonisolated(nonsending) sending @escaping () async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + onResponse { _ in try await action() } + } + + public func onReasoning( + @_inheritActorContext perform action: + nonisolated(nonsending) sending @escaping ( + Transcript.Reasoning + ) async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .onReasoning(action)) + } + + public func onReasoning( + @_inheritActorContext perform action: nonisolated(nonsending) sending @escaping () async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + onReasoning { _ in try await action() } + } + + public func onToolCall( + @_inheritActorContext perform action: + nonisolated(nonsending) sending @escaping ( + Transcript.ToolCall + ) async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .onToolCall(action)) + } + + public func onToolCall( + @_inheritActorContext perform action: nonisolated(nonsending) sending @escaping () async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + onToolCall { _ in try await action() } + } + + public func onToolOutput( + @_inheritActorContext perform action: + nonisolated(nonsending) sending @escaping ( + Transcript.ToolCall, + Transcript.ToolOutput + ) async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + BuiltinModifiedDynamicProfile(content: self, modifier: .onToolOutput(action)) + } + + public func onToolOutput( + @_inheritActorContext perform action: nonisolated(nonsending) sending @escaping () async throws -> Void + ) -> some LanguageModelSession.DynamicProfile { + onToolOutput { _, _ in try await action() } + } + + public func onActivate( + @_inheritActorContext perform action: sending @escaping @isolated(any) () async -> Void + ) -> some LanguageModelSession.DynamicProfile { + let wrapped: () async -> Void = { await action() } + return BuiltinModifiedDynamicProfile(content: self, modifier: .onActivate(wrapped)) + } + + public func onDeactivate( + @_inheritActorContext perform action: sending @escaping @isolated(any) () async -> Void + ) -> some LanguageModelSession.DynamicProfile { + let wrapped: () async -> Void = { await action() } + return BuiltinModifiedDynamicProfile(content: self, modifier: .onDeactivate(wrapped)) + } +} + +extension LanguageModelSession.DynamicProfileModifier { + public typealias DynamicProfile = LanguageModelSession.DynamicProfile + public typealias SessionProperty = LanguageModelSession.SessionProperty +} + +extension Never: LanguageModelSession.DynamicProfile {} + +private protocol PrimitiveDynamicProfile { + func resolve() -> ResolvedDynamicProfile +} + +struct ResolvedDynamicProfile: @unchecked Sendable { + var identity: String + var model: (any LanguageModel)? + var instructions: Instructions? + var tools: [any Tool] + var samplingMode: GenerationOptions.SamplingMode? + var hasSamplingMode = false + var temperature: Double? + var hasTemperature = false + var maximumResponseTokens: Int? + var hasMaximumResponseTokens = false + var reasoningLevel: ContextOptions.ReasoningLevel? + var hasReasoningLevel = false + var includeSchemaInPrompt: Bool? + var hasIncludeSchemaInPrompt = false + var toolCallingMode: GenerationOptions.ToolCallingMode? + var hasToolCallingMode = false + var transcriptErrorHandlingPolicy: TranscriptErrorHandlingPolicy? + var hasTranscriptErrorHandlingPolicy = false + var historyTransforms: [([Transcript.Entry]) -> [Transcript.Entry]] = [] + var onPrompt: [(Transcript.Prompt) async throws -> Void] = [] + var onResponse: [(Transcript.Response) async throws -> Void] = [] + var onReasoning: [(Transcript.Reasoning) async throws -> Void] = [] + var onToolCall: [(Transcript.ToolCall) async throws -> Void] = [] + var onToolOutput: [(Transcript.ToolCall, Transcript.ToolOutput) async throws -> Void] = [] + var onActivate: [() async -> Void] = [] + var onDeactivate: [() async -> Void] = [] +} + +extension LanguageModelSession.Profile: PrimitiveDynamicProfile { + fileprivate func resolve() -> ResolvedDynamicProfile { + let history = SessionPropertyBinding.history.map { + SessionHistoryBinding($0.entries, isWritable: false) + } + let dynamic = SessionPropertyBinding.$history.withValue(history) { + dynamicInstructions.resolveForRequest() + } + return ResolvedDynamicProfile( + identity: identity, + model: nil, + instructions: dynamic.instructions, + tools: dynamic.tools + ) + } +} + +extension LanguageModelSession.AnyDynamicProfile: PrimitiveDynamicProfile { + fileprivate func resolve() -> ResolvedDynamicProfile { + resolveForRequest() + } +} + +extension LanguageModelSession.ConditionalDynamicProfile: PrimitiveDynamicProfile { + fileprivate func resolve() -> ResolvedDynamicProfile { + var resolved: ResolvedDynamicProfile + switch branch { + case .trueContent(let content): + resolved = resolveDynamicProfile(content) + resolved.identity = "true/" + resolved.identity + case .falseContent(let content): + resolved = resolveDynamicProfile(content) + resolved.identity = "false/" + resolved.identity + } + return resolved + } +} + +extension LanguageModelSession.DynamicProfileModifierContent: PrimitiveDynamicProfile { + fileprivate func resolve() -> ResolvedDynamicProfile { + content.resolveForRequest() + } +} + +extension LanguageModelSession.ModifiedDynamicProfile: PrimitiveDynamicProfile { + fileprivate func resolve() -> ResolvedDynamicProfile { + resolveDynamicProfile(modifier.body(content: .init(content))) + } +} + +private struct BuiltinModifiedDynamicProfile: LanguageModelSession.DynamicProfile, + PrimitiveDynamicProfile +where Content: LanguageModelSession.DynamicProfile { + typealias Body = Never + let content: Content + let modifier: BuiltinDynamicProfileModifier + + var body: Never { + fatalError("BuiltinModifiedDynamicProfile has no body") + } + + func resolve() -> ResolvedDynamicProfile { + var resolved = resolveDynamicProfile(content) + modifier.apply(to: &resolved) + return resolved + } +} + +private enum BuiltinDynamicProfileModifier: @unchecked Sendable { + case model(any LanguageModel) + case temperature(Double?) + case samplingMode(GenerationOptions.SamplingMode?) + case maximumResponseTokens(Int?) + case reasoningLevel(ContextOptions.ReasoningLevel?) + case includeSchemaInPrompt(Bool?) + case toolCallingMode(GenerationOptions.ToolCallingMode?) + case historyTransform(([Transcript.Entry]) -> [Transcript.Entry]) + case transcriptErrorHandlingPolicy(TranscriptErrorHandlingPolicy?) + case onPrompt((Transcript.Prompt) async throws -> Void) + case onResponse((Transcript.Response) async throws -> Void) + case onReasoning((Transcript.Reasoning) async throws -> Void) + case onToolCall((Transcript.ToolCall) async throws -> Void) + case onToolOutput((Transcript.ToolCall, Transcript.ToolOutput) async throws -> Void) + case onActivate(() async -> Void) + case onDeactivate(() async -> Void) + + func apply(to resolved: inout ResolvedDynamicProfile) { + switch self { + case .model(let model): + if resolved.model == nil { resolved.model = model } + case .temperature(let value): + if !resolved.hasTemperature { + resolved.temperature = value + resolved.hasTemperature = true + } + case .samplingMode(let value): + if !resolved.hasSamplingMode { + resolved.samplingMode = value + resolved.hasSamplingMode = true + } + case .maximumResponseTokens(let value): + if !resolved.hasMaximumResponseTokens { + resolved.maximumResponseTokens = value + resolved.hasMaximumResponseTokens = true + } + case .reasoningLevel(let value): + if !resolved.hasReasoningLevel { + resolved.reasoningLevel = value + resolved.hasReasoningLevel = true + } + case .includeSchemaInPrompt(let value): + if !resolved.hasIncludeSchemaInPrompt { + resolved.includeSchemaInPrompt = value + resolved.hasIncludeSchemaInPrompt = true + } + case .toolCallingMode(let value): + if !resolved.hasToolCallingMode { + resolved.toolCallingMode = value + resolved.hasToolCallingMode = true + } + case .transcriptErrorHandlingPolicy(let value): + if !resolved.hasTranscriptErrorHandlingPolicy { + resolved.transcriptErrorHandlingPolicy = value + resolved.hasTranscriptErrorHandlingPolicy = true + } + case .historyTransform(let transform): + resolved.historyTransforms.insert(transform, at: 0) + case .onPrompt(let action): + resolved.onPrompt.insert(action, at: 0) + case .onResponse(let action): + resolved.onResponse.insert(action, at: 0) + case .onReasoning(let action): + resolved.onReasoning.insert(action, at: 0) + case .onToolCall(let action): + resolved.onToolCall.insert(action, at: 0) + case .onToolOutput(let action): + resolved.onToolOutput.insert(action, at: 0) + case .onActivate(let action): + resolved.onActivate.insert(action, at: 0) + case .onDeactivate(let action): + resolved.onDeactivate.append(action) + } + } +} + +private func resolveDynamicProfile( + _ dynamicProfile: any LanguageModelSession.DynamicProfile +) -> ResolvedDynamicProfile { + func resolve(_ content: Content) -> ResolvedDynamicProfile + where Content: LanguageModelSession.DynamicProfile { + if let primitive = content as? any PrimitiveDynamicProfile { + return primitive.resolve() + } + return resolve(content.body) + } + + return resolve(dynamicProfile) +} + +actor ProfileLifecycle { + private var activeProfile: ResolvedDynamicProfile? + private var pendingPreparation: (id: Int, task: Task)? + private var nextPreparationID = 0 + + func prepare( + _ profile: ResolvedDynamicProfile, + properties: SessionPropertyValues, + history: [Transcript.Entry], + prompt: Transcript.Prompt?, + callbackState: ProfileCallbackState, + historyEdits: Locked<[Transcript.Entry]?>, + protectedEntryKeys: Set = [] + ) async throws { + let previous = pendingPreparation?.task + nextPreparationID += 1 + let id = nextPreparationID + let task = Task { + _ = try? await previous?.value + try await self.performPreparation( + profile, + properties: properties, + history: history, + prompt: prompt, + callbackState: callbackState, + historyEdits: historyEdits, + protectedEntryKeys: protectedEntryKeys + ) + } + pendingPreparation = (id, task) + do { + try await task.value + if pendingPreparation?.id == id { pendingPreparation = nil } + } catch { + if pendingPreparation?.id == id { pendingPreparation = nil } + throw error + } + } + + private func performPreparation( + _ profile: ResolvedDynamicProfile, + properties: SessionPropertyValues, + history: [Transcript.Entry], + prompt: Transcript.Prompt?, + callbackState: ProfileCallbackState, + historyEdits: Locked<[Transcript.Entry]?>, + protectedEntryKeys: Set + ) async throws { + let history = properties.historyBinding( + history, + isWritable: true, + protecting: protectedEntryKeys, + onChange: { updated in historyEdits.withLock { $0 = updated } } + ) + if activeProfile?.identity != profile.identity { + if let previous = activeProfile { + await withBindings( + properties: properties, + history: history + ) { + for action in previous.onDeactivate { await action() } + } + } + await withBindings( + properties: properties, + history: history + ) { + for action in profile.onActivate { await action() } + } + activeProfile = profile + } else { + activeProfile = profile + } + + guard let prompt, callbackState.reserve(prompt.id) else { return } + do { + try await withBindings( + properties: properties, + history: history + ) { + for action in profile.onPrompt { try await action(prompt) } + } + } catch { + callbackState.release(prompt.id) + throw error + } + } + + private func withBindings( + properties: SessionPropertyValues, + history: SessionHistoryBinding, + operation: @Sendable () async throws -> Result + ) async rethrows -> Result { + return try await SessionPropertyBinding.$values.withValue(properties) { + try await SessionPropertyBinding.$history.withValue(history) { + try await operation() + } + } + } +} + +struct MissingDynamicProfileModel: LanguageModel { + typealias UnavailableReason = Never + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + throw MissingDynamicProfileModelError.missingModelModifier + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + let stream = AsyncThrowingStream.Snapshot, any Error> { + continuation in + continuation.finish(throwing: MissingDynamicProfileModelError.missingModelModifier) + } + return LanguageModelSession.ResponseStream(stream: stream) + } +} + +enum MissingDynamicProfileModelError: Error { + case missingModelModifier +} diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 4b57d92d..569b7409 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -66,6 +66,7 @@ public final class LanguageModelSession: @unchecked Sendable { @ObservationIgnored private let state: Locked @ObservationIgnored private let responseRelays = Locked<[UUID: Task]>([:]) + @ObservationIgnored private let lastResponseModel = Locked<(any LanguageModel)?>(nil) /// Waits for transcript cleanup of all streaming relays registered when this call begins. /// @@ -110,13 +111,21 @@ public final class LanguageModelSession: @unchecked Sendable { public let instructions: Instructions? private let dynamicInstructions: AnyDynamicInstructions? + private let dynamicProfile: AnyDynamicProfile? private let dynamicInstructionsLock = NSLock() + private let dynamicProfileLock = NSLock() + private let profileLifecycle = ProfileLifecycle() + + /// Session-scoped values available to profiles and dynamic instructions. + public let properties: SessionPropertyValues /// Whether the session resolves its instructions and tools before each request. nonisolated var usesDynamicInstructions: Bool { - dynamicInstructions != nil + dynamicInstructions != nil || dynamicProfile != nil } + nonisolated var usesDynamicProfile: Bool { dynamicProfile != nil } + /// A delegate that observes and controls tool execution. /// /// Set this property to intercept tool calls, provide custom output, @@ -150,16 +159,97 @@ public final class LanguageModelSession: @unchecked Sendable { /// The tools that the model can call in response to the request. public let tools: [any Tool] - fileprivate init(transcript: Transcript, instructions: Instructions?, tools: [any Tool]) { + func toolsForExecution(in session: LanguageModelSession, offeredTools: [any Tool]? = nil) -> [any Tool] { + (offeredTools ?? tools).map { + $0.bindingSessionProperties(session.properties, history: canonicalTranscript) + } + } + + /// The model selected for this request. + @_spi(Compatibility) + public let model: any LanguageModel + + /// Generation options after profile defaults are combined with call-site options. + @_spi(Compatibility) + public let options: GenerationOptions + + /// Context options selected by the active profile. + @_spi(Compatibility) + public let contextOptions: ContextOptions + + /// Transcript failure behavior captured for this request. + @_spi(Compatibility) + public let transcriptErrorHandlingPolicy: TranscriptErrorHandlingPolicy? + + fileprivate let profile: ResolvedDynamicProfile? + fileprivate let canonicalTranscript: [Transcript.Entry] + fileprivate let pendingEntryKeys: Set + fileprivate let callbackState: ProfileCallbackState + fileprivate let callbackHistory = Locked<[Transcript.Entry]?>(nil) + + var hasDynamicProfile: Bool { profile != nil } + + fileprivate init( + transcript: Transcript, + instructions: Instructions?, + tools: [any Tool], + model: any LanguageModel, + options: GenerationOptions = .init(), + contextOptions: ContextOptions = .init(), + transcriptErrorHandlingPolicy: TranscriptErrorHandlingPolicy? = nil, + profile: ResolvedDynamicProfile? = nil, + canonicalTranscript: [Transcript.Entry]? = nil, + pendingEntryKeys: Set = [], + callbackState: ProfileCallbackState = ProfileCallbackState() + ) { self.transcript = transcript self.instructions = instructions self.tools = tools + self.model = model + self.options = options + self.contextOptions = contextOptions + self.transcriptErrorHandlingPolicy = transcriptErrorHandlingPolicy + self.profile = profile + self.canonicalTranscript = canonicalTranscript ?? Array(transcript) + self.pendingEntryKeys = pendingEntryKeys + self.callbackState = callbackState + } + + fileprivate func sharingCallbackState(with context: RequestContext) -> RequestContext { + RequestContext( + transcript: transcript, + instructions: instructions, + tools: tools, + model: model, + options: options, + contextOptions: contextOptions, + transcriptErrorHandlingPolicy: transcriptErrorHandlingPolicy, + profile: profile, + canonicalTranscript: canonicalTranscript, + pendingEntryKeys: pendingEntryKeys, + callbackState: context.callbackState + ) + } + + fileprivate func withCanonicalTranscript(_ entries: [Transcript.Entry]) -> RequestContext { + RequestContext( + transcript: transcript, + instructions: instructions, + tools: tools, + model: model, + options: options, + contextOptions: contextOptions, + transcriptErrorHandlingPolicy: transcriptErrorHandlingPolicy, + profile: profile, + canonicalTranscript: entries, + callbackState: callbackState + ) } } /// Returns the transcript, instructions, and tools for the next model request. /// - /// Calling this method doesn't change the session. + /// Profile resolution may mutate session properties and canonical history. /// /// - Note: This API is exclusive to AnyLanguageModel /// and using it means your code is no longer drop-in compatible @@ -167,29 +257,402 @@ public final class LanguageModelSession: @unchecked Sendable { /// It's public so that language models outside this module /// can read the inputs for each request. nonisolated public func resolvedRequestContext() -> RequestContext { - guard let dynamicInstructions else { - return RequestContext(transcript: transcript, instructions: instructions, tools: tools) - } - - // Evaluate the body for this request. - // The resolved instructions go into the request's transcript, - // never into the session's. - let resolved = dynamicInstructionsLock.withLock { - dynamicInstructions.resolveForRequest() - } - var requestTranscript = transcript - if let instructions = resolved.instructions { - let instructionsEntry = Transcript.Entry.instructions( - Transcript.Instructions( - segments: [.text(Transcript.TextSegment(content: instructions.description))], - toolDefinitions: resolved.tools - .filter(\.includesSchemaInInstructions) - .map { Transcript.ToolDefinition(tool: $0) } + if let prepared = ProfileRequestBinding.prepared, + prepared.session === self + { + return prepared.context + } + return resolveRequestContext(additionalEntries: [], options: .init()) + } + + nonisolated func toolRoundLimit(for provider: String, options: GenerationOptions) -> ToolRoundLimit { + guard let prepared = ProfileRequestBinding.prepared, + prepared.session === self + else { + return ToolRoundLimit(provider: provider, options: options) + } + return prepared.toolRoundLimit(for: provider) + } + + nonisolated func recordToolRound(_ limit: ToolRoundLimit, completed: [Transcript.Entry]) throws { + if let prepared = ProfileRequestBinding.prepared, prepared.session === self { + try limit.record(completed: prepared.pendingEntries(including: completed)) + } else { + try limit.record(completed: completed) + } + } + + nonisolated func profileIncludeSchemaInPrompt( + for context: RequestContext, + callSiteValue: Bool + ) -> Bool { + let originalValue: Bool + if let prepared = ProfileRequestBinding.prepared, prepared.session === self { + originalValue = prepared.callSiteIncludeSchemaInPrompt + } else { + originalValue = callSiteValue + } + return context.contextOptions.includeSchemaInPrompt ?? originalValue + } + + /// Resolves the active profile for one actual model request. + /// + /// `additionalEntries` supplies canonical entries produced in the current Tool + /// round but not yet committed to the durable Session transcript. A profile's + /// history transform sees those entries before selecting the request projection. + /// + /// - Note: This overload is exclusive to AnyLanguageModel. Compatibility + /// executors use it to preserve Foundation Models 27 request timing while + /// AnyLanguageModel still exposes its legacy model interface. + @_spi(Compatibility) + nonisolated public func resolvedRequestContext( + including additionalEntries: [Transcript.Entry], + options: GenerationOptions = .init() + ) async throws -> RequestContext { + let prepared = ProfileRequestBinding.prepared + if additionalEntries.isEmpty, let prepared, prepared.session === self { + return prepared.context + } + let callSiteOptions = + if let prepared, prepared.session === self { + prepared.callSiteOptions + } else { + options + } + let pendingEntries = + if let prepared, prepared.session === self { + prepared.pendingEntries(including: additionalEntries) + } else { + additionalEntries + } + var context = resolveRequestContext( + additionalEntries: pendingEntries, + options: callSiteOptions + ) + if let prepared, prepared.session === self { + context = context.sharingCallbackState(with: prepared.context) + } + if let profile = context.profile { + if let prepared, prepared.session === self { + prepared.context = context + } + try await profileLifecycle.prepare( + profile, + properties: properties, + history: context.canonicalTranscript, + prompt: prepared?.session === self + ? prepared?.prompt + : context.canonicalTranscript.reversed().compactMap { entry in + if case .prompt(let prompt) = entry { return prompt } + return nil + }.first, + callbackState: context.callbackState, + historyEdits: context.callbackHistory, + protectedEntryKeys: context.pendingEntryKeys + ) + } + if let prepared, prepared.session === self { + prepared.context = context + } + return context + } + + nonisolated private func resolveRequestContext( + additionalEntries: [Transcript.Entry], + options: GenerationOptions + ) -> RequestContext { + let committed = Array(transcript) + var entryKeys = Set(keyedTranscriptEntries(committed).map(\.key)) + let pending = keyedTranscriptEntries(additionalEntries) + .filter { entryKeys.insert($0.key).inserted } + let protectedEntryKeys = Set(pending.map(\.key)) + var canonical = committed + pending.map(\.entry) + let sessionErrorPolicy = transcriptErrorHandlingPolicy + guard let dynamicProfile else { + guard let dynamicInstructions else { + return RequestContext( + transcript: Transcript(entries: canonical), + instructions: instructions, + tools: tools, + model: model, + options: options, + transcriptErrorHandlingPolicy: sessionErrorPolicy, + canonicalTranscript: canonical, + pendingEntryKeys: protectedEntryKeys ) + } + + // Dynamic instructions are request-only. Their body can observe + // session properties and canonical history, but cannot rewrite history. + let history = properties.historyBinding(canonical, isWritable: false) + let resolved = dynamicInstructionsLock.withLock { + SessionPropertyBinding.$values.withValue(properties) { + SessionPropertyBinding.$history.withValue(history) { + dynamicInstructions.resolveForRequest() + } + } + } + var requestTranscript = Transcript(entries: canonical) + if let currentInstructions = resolved.instructions { + requestTranscript = Transcript( + entries: [instructionsEntry(currentInstructions, tools: resolved.tools)] + canonical + ) + } + return RequestContext( + transcript: requestTranscript, + instructions: resolved.instructions, + tools: resolved.tools, + model: model, + options: options, + transcriptErrorHandlingPolicy: sessionErrorPolicy, + canonicalTranscript: canonical, + pendingEntryKeys: protectedEntryKeys ) - requestTranscript = Transcript(entries: [instructionsEntry] + requestTranscript) } - return RequestContext(transcript: requestTranscript, instructions: resolved.instructions, tools: resolved.tools) + + let history = properties.historyBinding( + canonical, + isWritable: true, + protecting: protectedEntryKeys + ) + let resolved = dynamicProfileLock.withLock { + SessionPropertyBinding.$values.withValue(properties) { + SessionPropertyBinding.$history.withValue(history) { + dynamicProfile.resolveForRequest() + } + } + } + canonical = history.entries + var projected = canonical + for transform in resolved.historyTransforms { + let transformHistory = properties.historyBinding( + canonical, + isWritable: true, + protecting: protectedEntryKeys + ) + projected = SessionPropertyBinding.$values.withValue(properties) { + SessionPropertyBinding.$history.withValue(transformHistory) { + transform(projected) + } + } + canonical = transformHistory.entries + } + var requestTranscript = Transcript(entries: projected) + if let currentInstructions = resolved.instructions { + requestTranscript = Transcript( + entries: [instructionsEntry(currentInstructions, tools: resolved.tools)] + projected + ) + } + + var resolvedOptions = options + if resolvedOptions.sampling == nil, resolved.hasSamplingMode { + resolvedOptions.sampling = resolved.samplingMode + } + if resolvedOptions.temperature == nil, resolved.hasTemperature { + resolvedOptions.temperature = resolved.temperature + } + if resolvedOptions.maximumResponseTokens == nil, resolved.hasMaximumResponseTokens { + resolvedOptions.maximumResponseTokens = resolved.maximumResponseTokens + } + if resolvedOptions.toolCallingMode == nil, resolved.hasToolCallingMode { + resolvedOptions.toolCallingMode = resolved.toolCallingMode + } + return RequestContext( + transcript: requestTranscript, + instructions: resolved.instructions, + tools: resolved.tools, + model: resolved.model ?? model, + options: resolvedOptions, + contextOptions: .init( + includeSchemaInPrompt: resolved.includeSchemaInPrompt, + reasoningLevel: resolved.reasoningLevel + ), + transcriptErrorHandlingPolicy: resolved.hasTranscriptErrorHandlingPolicy + ? resolved.transcriptErrorHandlingPolicy + : sessionErrorPolicy, + profile: resolved, + canonicalTranscript: canonical, + pendingEntryKeys: protectedEntryKeys + ) + } + + nonisolated private func instructionsEntry( + _ instructions: Instructions, + tools: [any Tool] + ) -> Transcript.Entry { + .instructions( + .init( + segments: [.text(.init(content: instructions.description))], + toolDefinitions: + tools + .filter(\.includesSchemaInInstructions) + .map { Transcript.ToolDefinition(tool: $0) } + ) + ) + } + + /// Runs the active request snapshot's Tool-call callbacks. + @_spi(Compatibility) + nonisolated public func profileWillExecuteToolCall( + _ call: Transcript.ToolCall, + requestContext: RequestContext, + currentRoundEntries: [Transcript.Entry] + ) async throws { + guard let profile = requestContext.profile else { return } + let history = profileHistory( + requestContext: requestContext, + currentRoundEntries: currentRoundEntries + ) + try await withSessionProperties( + history: history.entries, + isHistoryWritable: true, + protectedEntryKeys: history.protectedEntryKeys, + onHistoryChange: { updated in requestContext.callbackHistory.withLock { $0 = updated } } + ) { + for action in profile.onToolCall { try await action(call) } + } + } + + /// Runs the active request snapshot's Tool-output callbacks. + @_spi(Compatibility) + nonisolated public func profileDidProduceToolOutput( + for call: Transcript.ToolCall, + output: Transcript.ToolOutput, + requestContext: RequestContext, + currentRoundEntries: [Transcript.Entry] + ) async throws { + guard let profile = requestContext.profile else { return } + let history = profileHistory( + requestContext: requestContext, + currentRoundEntries: currentRoundEntries + ) + try await withSessionProperties( + history: history.entries, + isHistoryWritable: true, + protectedEntryKeys: history.protectedEntryKeys, + onHistoryChange: { updated in requestContext.callbackHistory.withLock { $0 = updated } } + ) { + for action in profile.onToolOutput { try await action(call, output) } + } + } + + /// Runs reasoning and response callbacks from the profile snapshot that + /// produced the entries. Each stable entry ID is dispatched at most once + /// for the active request, including when streaming snapshots repeat it. + /// + /// - Note: This API is exclusive to AnyLanguageModel's compatibility model + /// seam. Ordinary callers compose callbacks through ``DynamicProfile``. + @_spi(Compatibility) + nonisolated public func profileDidProduceEntries( + _ entries: [Transcript.Entry], + requestContext: RequestContext, + currentRoundEntries: [Transcript.Entry] + ) async throws { + guard let profile = requestContext.profile else { return } + let history = profileHistory( + requestContext: requestContext, + currentRoundEntries: currentRoundEntries + ) + let binding = properties.historyBinding( + history.entries, + isWritable: true, + protecting: history.protectedEntryKeys, + onChange: { updated in requestContext.callbackHistory.withLock { $0 = updated } } + ) + for entry in entries { + switch entry { + case .reasoning(let reasoning): + guard requestContext.callbackState.reserve(entry.id) else { continue } + do { + try await withSessionProperties(history: binding) { + for action in profile.onReasoning { try await action(reasoning) } + } + } catch { + requestContext.callbackState.release(entry.id) + throw error + } + case .response(let response): + guard requestContext.callbackState.reserve(entry.id) else { continue } + do { + try await withSessionProperties(history: binding) { + for action in profile.onResponse { try await action(response) } + } + } catch { + requestContext.callbackState.release(entry.id) + throw error + } + default: + continue + } + } + } + + /// Executes compatibility-model Tool work with session properties bound to + /// the producing request. Custom properties remain mutable, while the + /// canonical history is read-only in a Tool, matching Foundation Models 27. + @_spi(Compatibility) + nonisolated public func withProfileToolExecution( + requestContext: RequestContext, + currentRoundEntries: [Transcript.Entry], + operation: @Sendable () async throws -> Result + ) async rethrows -> Result { + let history = profileHistory( + requestContext: requestContext, + currentRoundEntries: currentRoundEntries + ) + return try await withSessionProperties( + history: history.entries, + isHistoryWritable: false, + protectedEntryKeys: history.protectedEntryKeys + ) { + try await operation() + } + } + + nonisolated private func profileHistory( + requestContext: RequestContext, + currentRoundEntries: [Transcript.Entry] + ) -> (entries: [Transcript.Entry], protectedEntryKeys: Set) { + let canonicalKeys = Set(keyedTranscriptEntries(requestContext.canonicalTranscript).map(\.key)) + let pending = keyedTranscriptEntries(currentRoundEntries).filter { + !canonicalKeys.contains($0.key) || requestContext.pendingEntryKeys.contains($0.key) + } + let history = requestContext.callbackHistory.withLock { history in + let entries = mergingTranscriptEntries(history ?? requestContext.canonicalTranscript, with: pending) + history = entries + return entries + } + return ( + history, + requestContext.pendingEntryKeys.union(pending.map(\.key)) + ) + } + + nonisolated private func withSessionProperties( + history: [Transcript.Entry], + isHistoryWritable: Bool, + protectedEntryKeys: Set = [], + onHistoryChange: (@Sendable ([Transcript.Entry]) -> Void)? = nil, + operation: @Sendable () async throws -> Result + ) async rethrows -> Result { + let history = properties.historyBinding( + history, + isWritable: isHistoryWritable, + protecting: protectedEntryKeys, + onChange: onHistoryChange + ) + return try await withSessionProperties(history: history, operation: operation) + } + + nonisolated private func withSessionProperties( + history: SessionHistoryBinding, + operation: @Sendable () async throws -> Result + ) async rethrows -> Result { + return try await SessionPropertyBinding.$values.withValue(properties) { + try await SessionPropertyBinding.$history.withValue(history) { + try await operation() + } + } } /// Creates a session with a model, tools, @@ -284,7 +747,29 @@ public final class LanguageModelSession: @unchecked Sendable { tools: [], instructions: nil, transcript: Transcript(entries: history), - dynamicInstructions: AnyDynamicInstructions(dynamicInstructions) + dynamicInstructions: AnyDynamicInstructions(dynamicInstructions), + dynamicProfile: nil + ) + } + + /// Creates a session whose active model, instructions, Tools, and request + /// projection are resolved before every model request. + @_spi(Compatibility) + public convenience init( + profile: sending some DynamicProfile, + history: some Collection = [] + ) { + let history = history.filter { entry in + if case .instructions = entry { return false } + return true + } + self.init( + model: MissingDynamicProfileModel(), + tools: [], + instructions: nil, + transcript: Transcript(entries: history), + dynamicInstructions: nil, + dynamicProfile: AnyDynamicProfile(profile) ) } @@ -293,11 +778,13 @@ public final class LanguageModelSession: @unchecked Sendable { tools: [any Tool], instructions: Instructions?, transcript: Transcript, - dynamicInstructions: AnyDynamicInstructions? = nil + dynamicInstructions: AnyDynamicInstructions? = nil, + dynamicProfile: AnyDynamicProfile? = nil ) { self.model = model self.tools = tools self.dynamicInstructions = dynamicInstructions + self.dynamicProfile = dynamicProfile let resolvedInstructions = instructions ?? Self.instructions(from: transcript) self.instructions = resolvedInstructions @@ -324,7 +811,50 @@ public final class LanguageModelSession: @unchecked Sendable { } } - self.state = .init(.init(finalTranscript)) + let state = Locked(State(finalTranscript)) + self.state = state + let setHistory: @Sendable ([Transcript.Entry]) -> Void = { history in + let history = history.filter { entry in + if case .instructions = entry { return false } + return true + } + state.withLock { locked in + let instructions = locked.transcript.prefix { + if case .instructions = $0 { return true } + return false + } + locked.transcript = Transcript(entries: Array(instructions) + history) + } + } + let observedHistorySetter = Locked<(@Sendable ([Transcript.Entry]) -> Void)?>(nil) + self.properties = SessionPropertyValues( + historyGetter: { + state.withLock { locked in + locked.transcript.filter { entry in + if case .instructions = entry { return false } + return true + } + } + }, + historySetter: { history in + if let setter = observedHistorySetter.withLock({ $0 }) { + setter(history) + } else { + setHistory(history) + } + } + ) + observedHistorySetter.withLock { setter in + setter = { [weak self] history in + guard let self else { + setHistory(history) + return + } + self.withMutation(keyPath: \.transcript) { + setHistory(history) + } + } + } } private static func instructions(from transcript: Transcript) -> Instructions? { @@ -338,6 +868,7 @@ public final class LanguageModelSession: @unchecked Sendable { } public func prewarm(promptPrefix: Prompt? = nil) { + let model = usesDynamicProfile ? resolvedRequestContext().model : self.model model.prewarm(for: self, promptPrefix: promptPrefix) } @@ -373,17 +904,28 @@ public final class LanguageModelSession: @unchecked Sendable { } } + nonisolated private func withTranscriptMutation( + _ mutation: () throws -> Result + ) rethrows -> Result { + try properties._withSessionPropertyMutation(keyPath: \.history) { + try withMutation(keyPath: \.transcript, mutation) + } + } + // The prompt is the only entry committed before a request succeeds. Checkpoints and // the final response commit atomically at completion, so rollback owns only this ID. nonisolated private func removePrompt(id: String) { - withMutation(keyPath: \.transcript) { + withTranscriptMutation { state.withLock { state in state.transcript = Transcript(entries: state.transcript.filter { $0.id != id }) } } } - nonisolated private func wrapRespond(_ operation: (String) async throws -> T) async throws -> T { + nonisolated private func wrapRespond( + profileTracker: PreparedProfileRequestTracker? = nil, + _ operation: (String) async throws -> T + ) async throws -> T { let promptID = UUID().uuidString beginResponding() do { @@ -391,7 +933,13 @@ public final class LanguageModelSession: @unchecked Sendable { endResponding() return result } catch { - if transcriptErrorHandlingPolicy == .revertTranscript { + let policy: TranscriptErrorHandlingPolicy? + if let prepared = profileTracker?.prepared { + policy = prepared.context.transcriptErrorHandlingPolicy + } else { + policy = transcriptErrorHandlingPolicy + } + if policy == .revertTranscript { removePrompt(id: promptID) } endResponding() @@ -420,7 +968,8 @@ public final class LanguageModelSession: @unchecked Sendable { nonisolated private func wrapStream( _ upstream: sending ResponseStream, - promptEntry: Transcript.Entry + promptEntry: Transcript.Entry, + profileTracker: PreparedProfileRequestTracker? = nil ) -> ResponseStream where Content: Generable, Content.PartiallyGenerated: Sendable { let session = self let relay = AsyncThrowingStream.Snapshot, any Error> { continuation in @@ -434,6 +983,7 @@ public final class LanguageModelSession: @unchecked Sendable { session.beginResponding() var lastSnapshot: ResponseStream.Snapshot? var accountedUsage = Usage.zero + var committedEntryIDs: Set = [] do { for try await snapshot in stream { lastSnapshot = snapshot @@ -485,20 +1035,50 @@ public final class LanguageModelSession: @unchecked Sendable { providerMetadata: lastSnapshot.providerMetadata ) ) - session.withMutation(keyPath: \.transcript) { - session.state.withLock { - $0.transcript.append(contentsOf: lastSnapshot.transcriptEntries) - $0.transcript.append(responseEntry) + let committed = session.withTranscriptMutation { + session.state.withLock { state in + guard state.transcript.contains(where: { $0.id == promptEntry.id }) else { + return false + } + state.transcript.append(contentsOf: lastSnapshot.transcriptEntries) + state.transcript.append(responseEntry) + return true } } + if committed { + committedEntryIDs = Set(lastSnapshot.transcriptEntries.map(\.id) + [responseEntry.id]) + } + if committed, let prepared = profileTracker?.prepared { + let context = prepared.context.withCanonicalTranscript(Array(session.transcript)) + try await session.profileDidProduceEntries( + Array(lastSnapshot.transcriptEntries) + [responseEntry], + requestContext: context, + currentRoundEntries: [] + ) + session.lastResponseModel.withLock { $0 = prepared.context.model } + } session.endResponding() continuation.finish() } catch { - session.withMutation(keyPath: \.transcript) { + session.withTranscriptMutation { session.state.withLock { state in - if state.transcriptErrorHandlingPolicy == .preserveTranscript { + if !committedEntryIDs.isEmpty { + state.transcript = Transcript( + entries: state.transcript.filter { + !committedEntryIDs.contains($0.id) + } + ) + } + guard state.transcript.contains(where: { $0.id == promptEntry.id }) else { return } + let policy: TranscriptErrorHandlingPolicy? + if let prepared = profileTracker?.prepared { + policy = prepared.context.transcriptErrorHandlingPolicy + } else { + policy = state.transcriptErrorHandlingPolicy + } + if policy == .preserveTranscript { state.transcript.append(contentsOf: lastSnapshot?.transcriptEntries ?? []) - } else if state.transcriptErrorHandlingPolicy == .revertTranscript { + } else if policy == .revertTranscript { state.transcript = Transcript( entries: state.transcript.filter { $0.id != promptEntry.id } ) @@ -720,43 +1300,101 @@ public final class LanguageModelSession: @unchecked Sendable { includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) async throws -> Response where Content: Generable { - try await respond( + let profileTracker = dynamicProfile == nil ? nil : PreparedProfileRequestTracker() + return try await respond( to: prompt, responseFormat: type == String.self ? nil : .init(type: type), - options: options - ) { - try await model.respond( - within: self, - to: prompt, - generating: type, + options: options, + profileTracker: profileTracker + ) { requestPrompt in + try await profiledResponse( + prompt: requestPrompt, includeSchemaInPrompt: includeSchemaInPrompt, - options: options + options: options, + tracker: profileTracker + ) { + model, + resolvedOptions, + contextOptions in + try await model.respond( + within: self, + to: prompt, + generating: type, + includeSchemaInPrompt: contextOptions.includeSchemaInPrompt + ?? includeSchemaInPrompt, + options: resolvedOptions + ) + } + } + } + + nonisolated private func profiledResponse( + prompt: Transcript.Prompt, + includeSchemaInPrompt: Bool, + options: GenerationOptions, + tracker: PreparedProfileRequestTracker?, + generate: ( + any LanguageModel, + GenerationOptions, + ContextOptions + ) async throws -> Response + ) async throws -> ProfiledModelResponse where Content: Generable { + guard dynamicProfile != nil else { + return ProfiledModelResponse( + response: try await countingRequests { try await generate(model, options, .init()) }, + prepared: nil ) } + let context = resolveRequestContext(additionalEntries: [], options: options) + let prepared = PreparedProfileRequest( + session: self, + context: context, + prompt: prompt, + callSiteOptions: options, + callSiteIncludeSchemaInPrompt: includeSchemaInPrompt + ) + tracker?.prepared = prepared + let response = try await ProfileRequestBinding.$prepared.withValue(prepared) { + if let profile = context.profile { + try await profileLifecycle.prepare( + profile, + properties: properties, + history: context.canonicalTranscript, + prompt: prepared.prompt, + callbackState: context.callbackState, + historyEdits: context.callbackHistory + ) + } + return try await countingRequests { + try await generate(context.model, context.options, context.contextOptions) + } + } + return ProfiledModelResponse(response: response, prepared: prepared) } nonisolated private func respond( to prompt: Prompt, responseFormat: Transcript.ResponseFormat?, options: GenerationOptions, - generate: () async throws -> Response + profileTracker: PreparedProfileRequestTracker?, + generate: (Transcript.Prompt) async throws -> ProfiledModelResponse ) async throws -> Response { - try await wrapRespond { promptID in + try await wrapRespond(profileTracker: profileTracker) { promptID in // Add prompt to transcript - let promptEntry = Transcript.Entry.prompt( - Transcript.Prompt( - id: promptID, - segments: try prompt.makeTranscriptSegments(occurrenceID: promptID), - options: options, - responseFormat: responseFormat - ) + let requestPrompt = Transcript.Prompt( + id: promptID, + segments: try prompt.makeTranscriptSegments(occurrenceID: promptID), + options: options, + responseFormat: responseFormat ) - withMutation(keyPath: \.transcript) { + let promptEntry = Transcript.Entry.prompt(requestPrompt) + withTranscriptMutation { state.withLock { $0.transcript.append(promptEntry) } } + let profiled = try await generate(requestPrompt) let response = Self.representingNilExplicitly( - in: try await countingRequests(of: generate), + in: profiled.response, responseFormat: responseFormat ) @@ -776,15 +1414,51 @@ public final class LanguageModelSession: @unchecked Sendable { ) ) - // Add tool entries and response to transcript - withMutation(keyPath: \.transcript) { + try await commitResponse( + response, + entry: responseEntry, + promptID: promptID, + prepared: profiled.prepared + ) + + return response + } + } + + nonisolated private func commitResponse( + _ response: Response, + entry responseEntry: Transcript.Entry, + promptID: String, + prepared: PreparedProfileRequest? + ) async throws where Content: Generable { + let committed = withTranscriptMutation { + state.withLock { lockedState in + guard lockedState.transcript.contains(where: { $0.id == promptID }) else { return false } + lockedState.transcript.append(contentsOf: response.transcriptEntries) + lockedState.transcript.append(responseEntry) + return true + } + } + + guard committed, let prepared else { return } + let context = prepared.context.withCanonicalTranscript(Array(transcript)) + do { + try await profileDidProduceEntries( + Array(response.transcriptEntries) + [responseEntry], + requestContext: context, + currentRoundEntries: [] + ) + lastResponseModel.withLock { $0 = prepared.context.model } + } catch { + let ids = Set(response.transcriptEntries.map(\.id) + [responseEntry.id]) + withTranscriptMutation { state.withLock { lockedState in - lockedState.transcript.append(contentsOf: response.transcriptEntries) - lockedState.transcript.append(responseEntry) + lockedState.transcript = Transcript( + entries: lockedState.transcript.filter { !ids.contains($0.id) } + ) } } - - return response + throw error } } @@ -802,22 +1476,33 @@ public final class LanguageModelSession: @unchecked Sendable { } catch { return ResponseStream(stream: AsyncThrowingStream { $0.finish(throwing: error) }) } - let promptEntry = Transcript.Entry.prompt( - Transcript.Prompt( - id: promptID, - segments: segments, - options: options, - responseFormat: type == String.self ? nil : .init(type: type) - ) + let requestPrompt = Transcript.Prompt( + id: promptID, + segments: segments, + options: options, + responseFormat: type == String.self ? nil : .init(type: type) ) - withMutation(keyPath: \.transcript) { + let promptEntry = Transcript.Entry.prompt(requestPrompt) + withTranscriptMutation { state.withLock { $0.transcript.append(promptEntry) } } - // A stream counts its usage from its snapshots: its requests aren't counted for a response - // it's made within, such as by a tool. - return wrapStream( - RequestUsageChannel.outside { + let upstream: ResponseStream + let profileTracker: PreparedProfileRequestTracker? + if dynamicProfile != nil { + let tracker = PreparedProfileRequestTracker() + profileTracker = tracker + upstream = profileDispatchedStream( + prompt: prompt, + requestPrompt: requestPrompt, + type: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options, + tracker: tracker + ) + } else { + profileTracker = nil + upstream = RequestUsageChannel.outside { model.streamResponse( within: self, to: prompt, @@ -825,10 +1510,101 @@ public final class LanguageModelSession: @unchecked Sendable { includeSchemaInPrompt: includeSchemaInPrompt, options: options ) - }, - promptEntry: promptEntry + } + } + return wrapStream( + upstream, + promptEntry: promptEntry, + profileTracker: profileTracker ) } + + nonisolated private func profileDispatchedStream( + prompt: Prompt, + requestPrompt: Transcript.Prompt, + type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions, + tracker: PreparedProfileRequestTracker + ) -> ResponseStream where Content: Generable { + profileDispatchedStream( + requestPrompt: requestPrompt, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options, + tracker: tracker + ) { + model, + resolvedOptions, + contextOptions in + model.streamResponse( + within: self, + to: prompt, + generating: type, + includeSchemaInPrompt: contextOptions.includeSchemaInPrompt + ?? includeSchemaInPrompt, + options: resolvedOptions + ) + } + } + + nonisolated private func profileDispatchedStream( + requestPrompt: Transcript.Prompt, + includeSchemaInPrompt: Bool, + options: GenerationOptions, + tracker: PreparedProfileRequestTracker, + generate: + @Sendable @escaping ( + any LanguageModel, + GenerationOptions, + ContextOptions + ) -> ResponseStream + ) -> ResponseStream where Content: Generable { + let stream = RequestUsageChannel.outside { + AsyncThrowingStream.Snapshot, any Error> { continuation in + let task = Task { + do { + let context = self.resolveRequestContext( + additionalEntries: [], + options: options + ) + let prepared = PreparedProfileRequest( + session: self, + context: context, + prompt: requestPrompt, + callSiteOptions: options, + callSiteIncludeSchemaInPrompt: includeSchemaInPrompt + ) + tracker.prepared = prepared + try await ProfileRequestBinding.$prepared.withValue(prepared) { + if let profile = context.profile { + try await self.profileLifecycle.prepare( + profile, + properties: self.properties, + history: context.canonicalTranscript, + prompt: prepared.prompt, + callbackState: context.callbackState, + historyEdits: context.callbackHistory + ) + } + let upstream = generate( + context.model, + context.options, + context.contextOptions + ) + for try await snapshot in upstream { + continuation.yield(snapshot) + } + } + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in task.cancel() } + } + } + return ResponseStream(stream: stream) + } } // MARK: - String Response Convenience Methods @@ -902,14 +1678,31 @@ extension LanguageModelSession { includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) async throws -> Response { - try await respond(to: prompt, responseFormat: .init(schema: schema), options: options) { - try await model.respond( - within: self, - to: prompt, - schema: schema, + let profileTracker = dynamicProfile == nil ? nil : PreparedProfileRequestTracker() + return try await respond( + to: prompt, + responseFormat: .init(schema: schema), + options: options, + profileTracker: profileTracker + ) { requestPrompt in + try await profiledResponse( + prompt: requestPrompt, includeSchemaInPrompt: includeSchemaInPrompt, - options: options - ) + options: options, + tracker: profileTracker + ) { + model, + resolvedOptions, + contextOptions in + try await model.respond( + within: self, + to: prompt, + schema: schema, + includeSchemaInPrompt: contextOptions.includeSchemaInPrompt + ?? includeSchemaInPrompt, + options: resolvedOptions + ) + } } } @@ -962,22 +1755,43 @@ extension LanguageModelSession { } catch { return ResponseStream(stream: AsyncThrowingStream { $0.finish(throwing: error) }) } - let promptEntry = Transcript.Entry.prompt( - Transcript.Prompt( - id: promptID, - segments: segments, - options: options, - responseFormat: .init(schema: schema) - ) + let requestPrompt = Transcript.Prompt( + id: promptID, + segments: segments, + options: options, + responseFormat: .init(schema: schema) ) - withMutation(keyPath: \.transcript) { + let promptEntry = Transcript.Entry.prompt(requestPrompt) + withTranscriptMutation { state.withLock { $0.transcript.append(promptEntry) } } - // A stream counts its usage from its snapshots: its requests aren't counted for a response - // it's made within, such as by a tool. - return wrapStream( - RequestUsageChannel.outside { + let upstream: ResponseStream + let profileTracker: PreparedProfileRequestTracker? + if dynamicProfile != nil { + let tracker = PreparedProfileRequestTracker() + profileTracker = tracker + upstream = profileDispatchedStream( + requestPrompt: requestPrompt, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options, + tracker: tracker + ) { + model, + resolvedOptions, + contextOptions in + model.streamResponse( + within: self, + to: prompt, + schema: schema, + includeSchemaInPrompt: contextOptions.includeSchemaInPrompt + ?? includeSchemaInPrompt, + options: resolvedOptions + ) + } + } else { + profileTracker = nil + upstream = RequestUsageChannel.outside { model.streamResponse( within: self, to: prompt, @@ -985,8 +1799,12 @@ extension LanguageModelSession { includeSchemaInPrompt: includeSchemaInPrompt, options: options ) - }, - promptEntry: promptEntry + } + } + return wrapStream( + upstream, + promptEntry: promptEntry, + profileTracker: profileTracker ) } @@ -1156,7 +1974,8 @@ extension LanguageModelSession { includeSchemaInPrompt: Bool = true, options: GenerationOptions = GenerationOptions() ) async throws -> Response where Content: Generable { - try await wrapRespond { promptID in + let profileTracker = dynamicProfile == nil ? nil : PreparedProfileRequestTracker() + return try await wrapRespond(profileTracker: profileTracker) { promptID in // Build segments from text and images var segments: [Transcript.Segment] = [] if !prompt.isEmpty { @@ -1165,31 +1984,40 @@ extension LanguageModelSession { segments.append(contentsOf: images.map { .image($0) }) // Add prompt to transcript - let promptEntry = Transcript.Entry.prompt( - Transcript.Prompt( - id: promptID, - segments: segments, - options: options, - responseFormat: type == String.self ? nil : .init(type: type) - ) + let requestPrompt = Transcript.Prompt( + id: promptID, + segments: segments, + options: options, + responseFormat: type == String.self ? nil : .init(type: type) ) - withMutation(keyPath: \.transcript) { + let promptEntry = Transcript.Entry.prompt(requestPrompt) + withTranscriptMutation { state.withLock { $0.transcript.append(promptEntry) } } // Extract text content for the Prompt parameter let textPrompt = Prompt(prompt) + let profiled = try await profiledResponse( + prompt: requestPrompt, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options, + tracker: profileTracker + ) { + model, + resolvedOptions, + contextOptions in + try await model.respond( + within: self, + to: textPrompt, + generating: type, + includeSchemaInPrompt: contextOptions.includeSchemaInPrompt + ?? includeSchemaInPrompt, + options: resolvedOptions + ) + } let response = Self.representingNilExplicitly( - in: try await countingRequests { - try await model.respond( - within: self, - to: textPrompt, - generating: type, - includeSchemaInPrompt: includeSchemaInPrompt, - options: options - ) - }, + in: profiled.response, responseFormat: type == String.self ? nil : .init(type: type) ) @@ -1209,13 +2037,12 @@ extension LanguageModelSession { ) ) - // Add tool entries and response to transcript - withMutation(keyPath: \.transcript) { - state.withLock { lockedState in - lockedState.transcript.append(contentsOf: response.transcriptEntries) - lockedState.transcript.append(responseEntry) - } - } + try await commitResponse( + response, + entry: responseEntry, + promptID: promptID, + prepared: profiled.prepared + ) return response } @@ -1298,24 +2125,45 @@ extension LanguageModelSession { segments.append(contentsOf: images.map { .image($0) }) // Create prompt entry that will be added when stream starts - let promptEntry = Transcript.Entry.prompt( - Transcript.Prompt( - segments: segments, - options: options, - responseFormat: type == String.self ? nil : .init(type: type) - ) + let requestPrompt = Transcript.Prompt( + segments: segments, + options: options, + responseFormat: type == String.self ? nil : .init(type: type) ) - withMutation(keyPath: \.transcript) { + let promptEntry = Transcript.Entry.prompt(requestPrompt) + withTranscriptMutation { state.withLock { $0.transcript.append(promptEntry) } } // Extract text content for the Prompt parameter let textPrompt = Prompt(prompt) - // A stream counts its usage from its snapshots: its requests aren't counted for a response - // it's made within, such as by a tool. - return wrapStream( - RequestUsageChannel.outside { + let upstream: ResponseStream + let profileTracker: PreparedProfileRequestTracker? + if dynamicProfile != nil { + let tracker = PreparedProfileRequestTracker() + profileTracker = tracker + upstream = profileDispatchedStream( + requestPrompt: requestPrompt, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options, + tracker: tracker + ) { + model, + resolvedOptions, + contextOptions in + model.streamResponse( + within: self, + to: textPrompt, + generating: type, + includeSchemaInPrompt: contextOptions.includeSchemaInPrompt + ?? includeSchemaInPrompt, + options: resolvedOptions + ) + } + } else { + profileTracker = nil + upstream = RequestUsageChannel.outside { model.streamResponse( within: self, to: textPrompt, @@ -1323,8 +2171,12 @@ extension LanguageModelSession { includeSchemaInPrompt: includeSchemaInPrompt, options: options ) - }, - promptEntry: promptEntry + } + } + return wrapStream( + upstream, + promptEntry: promptEntry, + profileTracker: profileTracker ) } } @@ -1338,7 +2190,10 @@ extension LanguageModelSession { issues: [LanguageModelFeedback.Issue] = [], desiredOutput: Transcript.Entry? = nil ) -> Data { - model.logFeedbackAttachment( + let model = + lastResponseModel.withLock { $0 } + ?? (usesDynamicProfile ? resolvedRequestContext().model : self.model) + return model.logFeedbackAttachment( within: self, sentiment: sentiment, issues: issues, @@ -1719,6 +2574,85 @@ private enum ResponseStreamError: Error, LocalizedError { } } +private struct ProfiledModelResponse where Content: Generable { + let response: LanguageModelSession.Response + let prepared: PreparedProfileRequest? +} + +final class ProfileCallbackState: @unchecked Sendable { + private let dispatchedEntryIDs = Locked>([]) + + func reserve(_ id: String) -> Bool { + dispatchedEntryIDs.withLock { $0.insert(id).inserted } + } + + func release(_ id: String) { + dispatchedEntryIDs.withLock { _ = $0.remove(id) } + } +} + +private final class PreparedProfileRequest: @unchecked Sendable { + let session: LanguageModelSession + let prompt: Transcript.Prompt + let callSiteOptions: GenerationOptions + let callSiteIncludeSchemaInPrompt: Bool + private let state: Locked + private let toolRoundLimitStorage = Locked(nil) + + init( + session: LanguageModelSession, + context: LanguageModelSession.RequestContext, + prompt: Transcript.Prompt, + callSiteOptions: GenerationOptions, + callSiteIncludeSchemaInPrompt: Bool + ) { + self.session = session + self.prompt = prompt + self.callSiteOptions = callSiteOptions + self.callSiteIncludeSchemaInPrompt = callSiteIncludeSchemaInPrompt + state = Locked(State(context: context)) + } + + var context: LanguageModelSession.RequestContext { + get { state.withLock { $0.context } } + set { state.withLock { $0.context = newValue } } + } + + func pendingEntries(including entries: [Transcript.Entry]) -> [Transcript.Entry] { + state.withLock { state in + state.pendingEntries = mergingTranscriptEntries(state.pendingEntries, with: keyedTranscriptEntries(entries)) + return state.pendingEntries + } + } + + func toolRoundLimit(for provider: String) -> ToolRoundLimit { + toolRoundLimitStorage.withLock { limit in + if let limit { return limit } + let created = ToolRoundLimit(provider: provider, options: callSiteOptions) + limit = created + return created + } + } + + private struct State: Sendable { + var context: LanguageModelSession.RequestContext + var pendingEntries: [Transcript.Entry] = [] + } +} + +private final class PreparedProfileRequestTracker: @unchecked Sendable { + private let storage = Locked(nil) + + var prepared: PreparedProfileRequest? { + get { storage.withLock { $0 } } + set { storage.withLock { $0 = newValue } } + } +} + +private enum ProfileRequestBinding { + @TaskLocal static var prepared: PreparedProfileRequest? +} + // MARK: - private struct State: Equatable, Sendable { diff --git a/Sources/AnyLanguageModel/Models/AnthropicLanguageModel.swift b/Sources/AnyLanguageModel/Models/AnthropicLanguageModel.swift index 6d2e1de1..7204e87a 100644 --- a/Sources/AnyLanguageModel/Models/AnthropicLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/AnthropicLanguageModel.swift @@ -527,7 +527,7 @@ public struct AnthropicLanguageModel: LanguageModel { if !toolUses.isEmpty { let resolution = try await resolveToolUses( toolUses, - tools: requestContext.tools, + tools: requestContext.toolsForExecution(in: session), session: session ) switch resolution { @@ -631,7 +631,7 @@ public struct AnthropicLanguageModel: LanguageModel { type == String.self ? nil : try convertSchemaToAnthropicFormat(schema) var inFlightMessages: [AnthropicMessage] = [] var state = StreamingResponseState() - var toolRounds = ToolRoundLimit(provider: "Anthropic", options: options) + let toolRounds = ToolRoundLimit(provider: "Anthropic", options: options) while true { try Task.checkCancellation() // The entries of the rounds that ran, before this round adds its reasoning. @@ -742,7 +742,7 @@ public struct AnthropicLanguageModel: LanguageModel { try toolRounds.record(completed: completedEntries) switch try await resolveToolUses( toolUses, - tools: requestContext.tools, + tools: requestContext.toolsForExecution(in: session), session: session ) { case .stop(let calls): @@ -760,7 +760,10 @@ public struct AnthropicLanguageModel: LanguageModel { .toolResult( .init( toolUseId: invocation.call.id, - content: convertSegmentsToAnthropicContent(invocation.output.segments) + content: try convertSegmentsToAnthropicContent( + invocation.output.segments, + isToolOutput: true + ) ) ) ) @@ -1117,14 +1120,14 @@ extension Transcript { messages.append( .init( role: .user, - content: convertSegmentsToAnthropicContent(instructions.segments) + content: try convertSegmentsToAnthropicContent(instructions.segments) ) ) case .prompt(let prompt): messages.append( .init( role: .user, - content: convertSegmentsToAnthropicContent(prompt.segments) + content: try convertSegmentsToAnthropicContent(prompt.segments) ) ) case .reasoning(let reasoning): @@ -1151,7 +1154,7 @@ extension Transcript { case .response(let response): // Anthropic rejects text blocks without non-whitespace text, // such as the empty response of a turn that only called tools. - let content = convertSegmentsToAnthropicContent(response.segments).filter { block in + let content = try convertSegmentsToAnthropicContent(response.segments).filter { block in guard case .text(let text) = block else { return true } return !text.text.allSatisfy(\.isWhitespace) } @@ -1179,7 +1182,10 @@ extension Transcript { .toolResult( AnthropicToolResult( toolUseId: toolOutput.id, - content: convertSegmentsToAnthropicContent(toolOutput.segments) + content: try convertSegmentsToAnthropicContent( + toolOutput.segments, + isToolOutput: true + ) ) ) ] @@ -1333,7 +1339,10 @@ private extension AnthropicThinking { } } -private func convertSegmentsToAnthropicContent(_ segments: [Transcript.Segment]) -> [AnthropicContent] { +private func convertSegmentsToAnthropicContent( + _ segments: [Transcript.Segment], + isToolOutput: Bool = false +) throws -> [AnthropicContent] { var blocks: [AnthropicContent] = [] blocks.reserveCapacity(segments.count) for segment in segments { @@ -1345,8 +1354,36 @@ private func convertSegmentsToAnthropicContent(_ segments: [Transcript.Segment]) case .image(let img): switch img.source { case .url(let url): - blocks.append(.image(AnthropicImage(url: url.absoluteString))) + if url.isFileURL { + let mimeType: String + switch url.pathExtension.lowercased() { + case "jpg", "jpeg": mimeType = "image/jpeg" + case "png": mimeType = "image/png" + case "gif": mimeType = "image/gif" + case "webp": mimeType = "image/webp" + default: + if isToolOutput { + blocks.append(.text(AnthropicText(text: ""))) + continue + } + throw AnthropicImageAttachmentError.unsupportedImageType(url) + } + blocks.append( + .image( + AnthropicImage( + base64Data: try Data(contentsOf: url).base64EncodedString(), + mimeType: mimeType + ) + ) + ) + } else { + blocks.append(.image(AnthropicImage(url: url.absoluteString))) + } case .data(let data, let mimeType): + if isToolOutput && !["image/jpeg", "image/png", "image/gif", "image/webp"].contains(mimeType) { + blocks.append(.text(AnthropicText(text: ""))) + continue + } blocks.append(.image(AnthropicImage(base64Data: data.base64EncodedString(), mimeType: mimeType))) } } @@ -1354,6 +1391,17 @@ private func convertSegmentsToAnthropicContent(_ segments: [Transcript.Segment]) return blocks } +private enum AnthropicImageAttachmentError: LocalizedError { + case unsupportedImageType(URL) + + var errorDescription: String? { + switch self { + case .unsupportedImageType(let url): + "Anthropic image attachments support PNG, JPEG, GIF, and WebP files; cannot determine a supported MIME type for \(url.lastPathComponent)." + } + } +} + private struct AnthropicToolUse: Codable, Sendable { let type: String let id: String diff --git a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift index 39c54903..103cc9af 100644 --- a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift @@ -9,7 +9,7 @@ import OrderedCollections import FoundationNetworking #endif -public struct GeminiLanguageModel: LanguageModel { +public struct GeminiLanguageModel: LanguageModel, ProfileRequestIdentifiedLanguageModel { public typealias UnavailableReason = Never public static let defaultBaseURL = URL(string: "https://generativelanguage.googleapis.com")! @@ -154,6 +154,7 @@ public struct GeminiLanguageModel: LanguageModel { public let responseStartTimeout: Duration? private let transport: any HTTPTransport + let profileRequestIdentity = UUID() /// Creates a new Gemini language model. /// @@ -242,47 +243,54 @@ public struct GeminiLanguageModel: LanguageModel { includeSchemaInPrompt: Bool, options: GenerationOptions ) async throws -> LanguageModelSession.Response where Content: Generable { - // Extract effective configuration from custom options or fall back to model defaults - let customOptions = options[custom: GeminiLanguageModel.self] - let effectiveThinking = customOptions?.thinking ?? .disabled - let effectiveServerTools = customOptions?.serverTools ?? [] - let effectiveJsonMode = customOptions?.jsonMode // Gemini's response schemas can be objects, primitives, or arrays, but not unions. let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .unionRoots) - let url = baseURL .appendingPathComponent(apiVersion) .appendingPathComponent("models/\(model):generateContent") let headers = buildHeaders() - var inFlightEntries: [Transcript.Entry] = [] - - // The entries this call adds, which is what the response reports. `inFlightEntries` - // preserves tool rounds while each iteration rebuilds the request from a fresh context. + // The entries this call adds, which is what the response reports. var entries: [Transcript.Entry] = [] var usage = RequestsUsage() // The text of earlier tool rounds, which string responses include. var earlierText = "" - var toolRounds = ToolRoundLimit(provider: "Gemini", options: options) + let toolRounds = session.toolRoundLimit(for: "Gemini", options: options) // Multi-turn conversation loop for tool calling while true { - let requestContext = session.resolvedRequestContext() + let requestContext = try await session.resolvedRequestContext( + including: entries, + options: options + ) + if let continuation = try await session.profileContinuationResponse( + ifSelectedModelDiffersFrom: profileRequestIdentity, + requestContext: requestContext, + prompt: prompt, + generating: type, + schema: schema, + includeSchemaInPrompt: includeSchemaInPrompt, + prefixEntries: entries, + prefixUsage: usage.total, + prefixText: earlierText + ) { + return continuation + } + let customOptions = requestContext.options[custom: GeminiLanguageModel.self] + let effectiveThinking = customOptions?.thinking ?? .disabled + let effectiveServerTools = customOptions?.serverTools ?? [] + let effectiveJsonMode = customOptions?.jsonMode let geminiTools = try buildTools( from: requestContext.tools, serverTools: effectiveServerTools ) - var requestTranscript = requestContext.transcript - for entry in inFlightEntries { - requestTranscript.append(entry) - } let params = try createGenerateContentParams( - contents: requestTranscript.toGeminiContent(), + contents: try requestContext.transcript.toGeminiContent(), tools: geminiTools, generating: type, schema: responseSchema.schema, - options: options, + options: requestContext.options, thinking: effectiveThinking, jsonMode: effectiveJsonMode ) @@ -312,16 +320,18 @@ public struct GeminiLanguageModel: LanguageModel { if !functionCalls.isEmpty { // Resolve function calls - try toolRounds.record(completed: entries) - let resolution = try await resolveFunctionCalls( + try session.recordToolRound(toolRounds, completed: entries) + let (callsEntry, resolution) = try await resolveFunctionCalls( functionCalls, - tools: requestContext.tools, + requestContext: requestContext, + currentRoundEntries: entries, + providerMetadata: providerMetadata, session: session ) switch resolution { case .stop(let calls): if !calls.isEmpty { - entries.append(.toolCalls(Transcript.ToolCalls(calls, providerMetadata: providerMetadata))) + entries.append(.toolCalls(callsEntry)) } let empty = try emptyResponseContent(for: type) return LanguageModelSession.Response( @@ -332,15 +342,11 @@ public struct GeminiLanguageModel: LanguageModel { ) case .invocations(let invocations): if !invocations.isEmpty { - let calls = Transcript.Entry.toolCalls( - Transcript.ToolCalls(invocations.map(\.call), providerMetadata: providerMetadata) - ) - inFlightEntries.append(calls) + let calls = Transcript.Entry.toolCalls(callsEntry) entries.append(calls) for invocation in invocations { let output = Transcript.Entry.toolOutput(invocation.output) - inFlightEntries.append(output) entries.append(output) } } @@ -419,14 +425,8 @@ public struct GeminiLanguageModel: LanguageModel { includeSchemaInPrompt: Bool, options: GenerationOptions ) -> sending LanguageModelSession.ResponseStream where Content: Generable { - // Extract effective configuration from custom options or fall back to model defaults - let customOptions = options[custom: GeminiLanguageModel.self] - let effectiveThinking = customOptions?.thinking ?? .disabled - let effectiveServerTools = customOptions?.serverTools ?? [] - let effectiveJsonMode = customOptions?.jsonMode // Gemini's response schemas can be objects, primitives, or arrays, but not unions. let responseSchema = ProviderResponseSchema(schema, generating: type, wrapping: .unionRoots) - var streamURL = baseURL .appendingPathComponent(apiVersion) @@ -439,26 +439,43 @@ public struct GeminiLanguageModel: LanguageModel { let task = Task { @Sendable in do { let headers = buildHeaders() - var inFlightEntries: [Transcript.Entry] = [] var state = StreamingResponseState(wrapperProperty: responseSchema.wrapperProperty) - var toolRounds = ToolRoundLimit(provider: "Gemini", options: options) + let toolRounds = session.toolRoundLimit(for: "Gemini", options: options) while true { try Task.checkCancellation() - let requestContext = session.resolvedRequestContext() + let requestContext = try await session.resolvedRequestContext( + including: state.entries, + options: options + ) + if let handoff = session.profileContinuationStream( + ifSelectedModelDiffersFrom: profileRequestIdentity, + requestContext: requestContext, + prompt: prompt, + generating: type, + schema: schema, + includeSchemaInPrompt: includeSchemaInPrompt, + prefixEntries: state.entries, + prefixUsage: state.requests.total, + prefixText: state.responseText + ) { + for try await snapshot in handoff { continuation.yield(snapshot) } + continuation.finish() + return + } + let customOptions = requestContext.options[custom: GeminiLanguageModel.self] + let effectiveThinking = customOptions?.thinking ?? .disabled + let effectiveServerTools = customOptions?.serverTools ?? [] + let effectiveJsonMode = customOptions?.jsonMode let geminiTools = try buildTools( from: requestContext.tools, serverTools: effectiveServerTools ) - var requestTranscript = requestContext.transcript - for entry in inFlightEntries { - requestTranscript.append(entry) - } let params = try createGenerateContentParams( - contents: requestTranscript.toGeminiContent(), + contents: try requestContext.transcript.toGeminiContent(), tools: geminiTools, generating: type, schema: responseSchema.schema, - options: options, + options: requestContext.options, thinking: effectiveThinking, jsonMode: effectiveJsonMode ) @@ -498,26 +515,25 @@ public struct GeminiLanguageModel: LanguageModel { guard !functionCalls.isEmpty else { break } try Task.checkCancellation() let metadata = try textPartMetadata(parts, includeUnsignedText: true) - try toolRounds.record(completed: state.entries) - switch try await resolveFunctionCalls( + try session.recordToolRound(toolRounds, completed: state.entries) + let (callsEntry, resolution) = try await resolveFunctionCalls( functionCalls, - tools: requestContext.tools, + requestContext: requestContext, + currentRoundEntries: state.entries, + providerMetadata: metadata, session: session - ) { - case .stop(let calls): - state.entries.append(.toolCalls(Transcript.ToolCalls(calls, providerMetadata: metadata))) + ) + switch resolution { + case .stop: + state.entries.append(.toolCalls(callsEntry)) continuation.yield(try state.stoppedSnapshot()) continuation.finish() return case .invocations(let invocations): - let calls = Transcript.Entry.toolCalls( - Transcript.ToolCalls(invocations.map(\.call), providerMetadata: metadata) - ) - inFlightEntries.append(calls) + let calls = Transcript.Entry.toolCalls(callsEntry) state.entries.append(calls) for invocation in invocations { let output = Transcript.Entry.toolOutput(invocation.output) - inFlightEntries.append(output) state.entries.append(output) } } @@ -675,18 +691,11 @@ private enum ToolResolutionOutcome { private func resolveFunctionCalls( _ functionCalls: [GeminiFunctionCall], - tools: [any Tool], + requestContext: LanguageModelSession.RequestContext, + currentRoundEntries: [Transcript.Entry], + providerMetadata: [String: String]? = nil, session: LanguageModelSession -) async throws -> ToolResolutionOutcome { - if functionCalls.isEmpty { return .invocations([]) } - - var toolsByName: [String: any Tool] = [:] - for tool in tools { - if toolsByName[tool.name] == nil { - toolsByName[tool.name] = tool - } - } - +) async throws -> (calls: Transcript.ToolCalls, outcome: ToolResolutionOutcome) { var transcriptCalls: [Transcript.ToolCall] = [] transcriptCalls.reserveCapacity(functionCalls.count) for call in functionCalls { @@ -702,82 +711,18 @@ private func resolveFunctionCalls( ) } - if let delegate = session.toolExecutionDelegate { - await delegate.didGenerateToolCalls(transcriptCalls, in: session) - } - - guard !transcriptCalls.isEmpty else { return .invocations([]) } - - var decisions: [ToolExecutionDecision] = [] - decisions.reserveCapacity(transcriptCalls.count) - - if let delegate = session.toolExecutionDelegate { - for call in transcriptCalls { - let decision = await delegate.toolCallDecision(for: call, in: session) - if case .stop = decision { - return .stop(calls: transcriptCalls) - } - decisions.append(decision) - } - } else { - decisions = Array(repeating: .execute, count: transcriptCalls.count) - } - - var results: [ToolInvocationResult] = [] - results.reserveCapacity(transcriptCalls.count) - - for (index, call) in transcriptCalls.enumerated() { - switch decisions[index] { - case .stop: - // This branch should be unreachable because `.stop` returns during decision collection. - // Keep it as a defensive guard in case that logic changes. - return .stop(calls: transcriptCalls) - case .provideOutput(let segments): - let output = Transcript.ToolOutput( - id: call.id, - toolName: call.toolName, - segments: segments - ) - if let delegate = session.toolExecutionDelegate { - await delegate.didExecuteToolCall(call, output: output, in: session) - } - results.append(ToolInvocationResult(call: call, output: output)) - case .execute: - guard let tool = toolsByName[call.toolName] else { - let message = Transcript.Segment.text(.init(content: "Tool not found: \(call.toolName)")) - let output = Transcript.ToolOutput( - id: call.id, - toolName: call.toolName, - segments: [message] - ) - if let delegate = session.toolExecutionDelegate { - await delegate.didExecuteToolCall(call, output: output, in: session) - } - results.append(ToolInvocationResult(call: call, output: output)) - continue - } - - do { - let segments = try await tool.makeOutputSegments(from: call.arguments) - let output = Transcript.ToolOutput( - id: call.id, - toolName: tool.name, - segments: segments - ) - if let delegate = session.toolExecutionDelegate { - await delegate.didExecuteToolCall(call, output: output, in: session) - } - results.append(ToolInvocationResult(call: call, output: output)) - } catch { - if let delegate = session.toolExecutionDelegate { - await delegate.didFailToolCall(call, error: error, in: session) - } - throw LanguageModelSession.ToolCallError(tool: tool, underlyingError: error) - } - } + let callsEntry = Transcript.ToolCalls(transcriptCalls, providerMetadata: providerMetadata) + let entries = currentRoundEntries + [.toolCalls(callsEntry)] + switch try await session.resolveProfileToolCalls( + transcriptCalls, + requestContext: requestContext, + currentRoundEntries: entries + ) { + case .stop(let calls): + return (callsEntry, .stop(calls: calls)) + case .invocations(let invocations): + return (callsEntry, .invocations(invocations.map { .init(call: $0.call, output: $0.output) })) } - - return .invocations(results) } /// Joins the text parts of a Gemini response. @@ -809,32 +754,79 @@ private func emptyResponseContent( } } -private func toJSONValue(_ toolOutput: Transcript.ToolOutput) throws -> [String: JSONValue] { - var result: [String: JSONValue] = [:] +private func toGeminiFunctionResponse( + _ toolOutput: Transcript.ToolOutput +) throws -> GeminiFunctionResponse { + var resultValues: [String] = [] + var parts: [GeminiPart] = [] for segment in toolOutput.segments { switch segment { case .text(let text): - result["result"] = .string(text.content) + resultValues.append(text.content) case .structure(let structured): - // For structured segments, encode the content let data = try JSONEncoder().encode(structured.content) if let jsonString = String(data: data, encoding: .utf8) { - result["result"] = .string(jsonString) + resultValues.append(jsonString) + } + case .image(let image): + switch image.source { + case .data(let data, let mimeType): + guard ["image/png", "image/jpeg", "image/webp"].contains(mimeType) else { + resultValues.append("") + continue + } + parts.append( + .inlineData( + GeminiInlineData( + mimeType: mimeType, + data: data.base64EncodedString() + ) + ) + ) + case .url(let url): + if url.isFileURL { + if let data = try geminiFunctionResponseData(from: url) { + parts.append(.inlineData(data)) + } else { + resultValues.append("") + } + } else { + resultValues.append(url.absoluteString) + } } - case .image: - // Ignore images in tool outputs for Gemini conversion - break } } - return result + var response: [String: JSONValue] = [:] + if !resultValues.isEmpty { + response["result"] = .string(resultValues.joined(separator: "\n")) + } + return GeminiFunctionResponse( + name: toolOutput.toolName, + response: response, + parts: parts.isEmpty ? nil : parts + ) +} + +private func geminiFunctionResponseData(from url: URL) throws -> GeminiInlineData? { + let mimeType: String + switch url.pathExtension.lowercased() { + case "png": mimeType = "image/png" + case "jpg", "jpeg": mimeType = "image/jpeg" + case "webp": mimeType = "image/webp" + default: return nil + } + return try GeminiInlineData( + mimeType: mimeType, + data: Data(contentsOf: url).base64EncodedString() + ) } // MARK: - Supporting Types extension Transcript { - fileprivate func toGeminiContent() -> [GeminiContent] { + fileprivate func toGeminiContent() throws -> [GeminiContent] { var messages = [GeminiContent]() for item in self { switch item { @@ -884,11 +876,7 @@ extension Transcript { ) case .toolOutput(let toolOutput): // Add function response as a user message (Gemini API expects function responses from user role) - let response = try? toJSONValue(toolOutput) - let functionResponse = GeminiFunctionResponse( - name: toolOutput.toolName, - response: response ?? [:] - ) + let functionResponse = try toGeminiFunctionResponse(toolOutput) messages.append( .init( role: .user, @@ -1162,6 +1150,7 @@ private struct GeminiFunctionCall: Codable, Sendable { private struct GeminiFunctionResponse: Codable, Sendable { let name: String let response: [String: JSONValue] + let parts: [GeminiPart]? } private struct GeminiGenerateContentResponse: Codable, Sendable { diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index f0a1eb0e..13f4bd5b 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -850,7 +850,7 @@ import Foundation if type == String.self { let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 let outputFormat = currentToolCallFormat() - var toolRounds = ToolRoundLimit(provider: "Llama", options: options) + let toolRounds = session.toolRoundLimit(for: "Llama", options: options) var allEntries: [Transcript.Entry] = [] var text = "" var usage = RequestsUsage() @@ -930,11 +930,11 @@ import Foundation break generationLoop } - try toolRounds.record(completed: allEntries) + try session.recordToolRound(toolRounds, completed: allEntries) let resolution = try await resolveToolCalls( parsedCalls, - tools: tools, + tools: requestContext.toolsForExecution(in: session, offeredTools: tools), session: session ) switch resolution { @@ -1085,7 +1085,7 @@ import Foundation let runtimeOptions = resolvedOptions(from: options) let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 let outputFormat = self.currentToolCallFormat() - var toolRounds = ToolRoundLimit(provider: "Llama", options: options) + let toolRounds = session.toolRoundLimit(for: "Llama", options: options) var accumulatedEntries: [Transcript.Entry] = [] var emittedBase = "" var usage = RequestsUsage() @@ -1194,11 +1194,11 @@ import Foundation break generationLoop } - try toolRounds.record(completed: accumulatedEntries) + try session.recordToolRound(toolRounds, completed: accumulatedEntries) let resolution = try await self.resolveToolCalls( parsedCalls, - tools: tools, + tools: requestContext.toolsForExecution(in: session, offeredTools: tools), session: session ) switch resolution { diff --git a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift index ef8d5008..ee2b292b 100644 --- a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift @@ -1128,7 +1128,7 @@ import Foundation var usage = RequestsUsage() var allTextChunks: [String] = [] var allEntries: [Transcript.Entry] = [] - var toolRounds = ToolRoundLimit(provider: "MLX", options: options) + let toolRounds = session.toolRoundLimit(for: "MLX", options: options) // Loop until no more tool calls while true { @@ -1197,11 +1197,11 @@ import Foundation // If there are tool calls, execute them and continue if !collectedToolCalls.isEmpty { - try toolRounds.record(completed: allEntries) + try session.recordToolRound(toolRounds, completed: allEntries) let resolution = try await resolveToolCalls( collectedToolCalls, - tools: tools, + tools: requestContext.toolsForExecution(in: session, offeredTools: tools), session: session ) switch resolution { @@ -1356,7 +1356,7 @@ import Foundation var usage = RequestsUsage() var accumulatedText = "" var accumulatedEntries: [Transcript.Entry] = [] - var toolRounds = ToolRoundLimit(provider: "MLX", options: options) + let toolRounds = session.toolRoundLimit(for: "MLX", options: options) // Yields a snapshot carrying the cumulative text and tool entries so far. func yieldSnapshot() { @@ -1440,11 +1440,11 @@ import Foundation guard !collectedToolCalls.isEmpty else { break } - try toolRounds.record(completed: accumulatedEntries) + try session.recordToolRound(toolRounds, completed: accumulatedEntries) let resolution = try await resolveToolCalls( collectedToolCalls, - tools: tools, + tools: requestContext.toolsForExecution(in: session, offeredTools: tools), session: session ) switch resolution { diff --git a/Sources/AnyLanguageModel/Models/OllamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/OllamaLanguageModel.swift index 48fb0523..39959daa 100644 --- a/Sources/AnyLanguageModel/Models/OllamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/OllamaLanguageModel.swift @@ -187,7 +187,7 @@ public struct OllamaLanguageModel: LanguageModel { if let toolCalls = chatResponse.message.toolCalls, !toolCalls.isEmpty { let resolution = try await resolveToolCalls( toolCalls, - tools: tools, + tools: requestContext.toolsForExecution(in: session, offeredTools: tools), session: session ) switch resolution { @@ -281,7 +281,7 @@ public struct OllamaLanguageModel: LanguageModel { let format = type == String.self ? nil : try JSONValue(convertSchemaToOllamaFormat(schema)) var inFlightMessages: [OllamaMessage] = [] var state = StreamingResponseState() - var toolRounds = ToolRoundLimit(provider: "Ollama", options: options) + let toolRounds = ToolRoundLimit(provider: "Ollama", options: options) while true { try Task.checkCancellation() let requestContext = session.resolvedRequestContext() @@ -338,7 +338,7 @@ public struct OllamaLanguageModel: LanguageModel { } switch try await resolveToolCalls( toolCalls, - tools: offeredTools, + tools: requestContext.toolsForExecution(in: session, offeredTools: offeredTools), session: session ) { case .stop(let calls): diff --git a/Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift b/Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift index 5b0196f3..7b95603d 100644 --- a/Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift @@ -527,7 +527,7 @@ public struct OpenAILanguageModel: LanguageModel { var roundText = "" var inFlightMessages: [OpenAIMessage] = [] - var toolRounds = ToolRoundLimit(provider: "OpenAI", options: options) + let toolRounds = ToolRoundLimit(provider: "OpenAI", options: options) // Loop until no more tool calls while true { let requestContext = session.resolvedRequestContext() @@ -576,7 +576,7 @@ public struct OpenAILanguageModel: LanguageModel { try toolRounds.record(completed: entries) let resolution = try await resolveToolCalls( toolCalls, - tools: requestContext.tools, + tools: requestContext.toolsForExecution(in: session), session: session ) switch resolution { @@ -669,7 +669,7 @@ public struct OpenAILanguageModel: LanguageModel { let url = baseURL.appendingPathComponent("responses") - var toolRounds = ToolRoundLimit(provider: "OpenAI", options: options) + let toolRounds = ToolRoundLimit(provider: "OpenAI", options: options) // Loop until no more tool calls while true { let requestContext = session.resolvedRequestContext() @@ -707,7 +707,7 @@ public struct OpenAILanguageModel: LanguageModel { try toolRounds.record(completed: entries) let resolution = try await resolveToolCalls( toolCalls, - tools: requestContext.tools, + tools: requestContext.toolsForExecution(in: session), session: session ) switch resolution { @@ -821,7 +821,7 @@ public struct OpenAILanguageModel: LanguageModel { do { var inFlightMessages: [OpenAIMessage] = [] var state = StreamingResponseState(wrapperProperty: responseSchema.wrapperProperty) - var toolRounds = ToolRoundLimit(provider: "OpenAI", options: options) + let toolRounds = ToolRoundLimit(provider: "OpenAI", options: options) while true { try Task.checkCancellation() let requestContext = session.resolvedRequestContext() @@ -944,7 +944,7 @@ public struct OpenAILanguageModel: LanguageModel { try toolRounds.record(completed: state.entries) switch try await resolveToolCalls( toolCalls, - tools: requestContext.tools, + tools: requestContext.toolsForExecution(in: session), session: session ) { case .stop(let calls): diff --git a/Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift b/Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift index 71b5859e..5a1058c8 100644 --- a/Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/OpenResponsesLanguageModel.swift @@ -533,7 +533,7 @@ public struct OpenResponsesLanguageModel: LanguageModel { do { var inFlightMessages: [OpenResponsesMessage] = [] var state = StreamingResponseState(wrapperProperty: responseSchema.wrapperProperty) - var toolRounds = ToolRoundLimit(provider: "Open Responses", options: options) + let toolRounds = ToolRoundLimit(provider: "Open Responses", options: options) while true { try Task.checkCancellation() let requestContext = session.resolvedRequestContext() @@ -595,7 +595,7 @@ public struct OpenResponsesLanguageModel: LanguageModel { try toolRounds.record(completed: state.entries) switch try await resolveToolCalls( toolCalls, - tools: requestContext.tools, + tools: requestContext.toolsForExecution(in: session), session: session ) { case .stop(let calls): @@ -648,7 +648,7 @@ public struct OpenResponsesLanguageModel: LanguageModel { var inFlightMessages: [OpenResponsesMessage] = [] let url = baseURL.appendingPathComponent("responses") - var toolRounds = ToolRoundLimit(provider: "Open Responses", options: options) + let toolRounds = ToolRoundLimit(provider: "Open Responses", options: options) while true { let requestContext = session.resolvedRequestContext() let tools: [OpenResponsesTool]? = @@ -685,7 +685,7 @@ public struct OpenResponsesLanguageModel: LanguageModel { try toolRounds.record(completed: entries) let resolution = try await resolveToolCalls( toolCalls, - tools: requestContext.tools, + tools: requestContext.toolsForExecution(in: session), session: session ) switch resolution { diff --git a/Sources/AnyLanguageModel/Models/PrivateCloudComputeLanguageModel.swift b/Sources/AnyLanguageModel/Models/PrivateCloudComputeLanguageModel.swift index d76657c2..b8c6bc3c 100644 --- a/Sources/AnyLanguageModel/Models/PrivateCloudComputeLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/PrivateCloudComputeLanguageModel.swift @@ -130,7 +130,7 @@ // Attach the feedback to the session's conversation, including its latest response. let fmSession = FoundationModels.LanguageModelSession( model: pccModel, - tools: requestContext.tools.toFoundationModels(), + tools: requestContext.toolsForExecution(in: session).toFoundationModels(), transcript: requestContext.transcript.toFoundationModels( instructions: requestContext.instructions, toolDefinitions: requestContext.tools diff --git a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift index 634bbcef..989cec38 100644 --- a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift @@ -277,7 +277,7 @@ // Attach the feedback to the session's conversation, including its latest response. let fmSession = FoundationModels.LanguageModelSession( model: systemModel, - tools: requestContext.tools.toFoundationModels(in: session), + tools: requestContext.toolsForExecution(in: session).toFoundationModels(in: session), transcript: requestContext.transcript.toFoundationModels( instructions: requestContext.instructions, toolDefinitions: requestContext.tools @@ -319,7 +319,7 @@ let requestContext = session.resolvedRequestContext() return FoundationModels.LanguageModelSession( model: systemModel, - tools: requestContext.tools.toFoundationModels(in: session), + tools: requestContext.toolsForExecution(in: session).toFoundationModels(in: session), transcript: fmTranscriptDroppingDuplicatePrompt( requestContext.transcript, prompt: prompt @@ -388,7 +388,7 @@ if let instructions = requestContext.instructions { instructions.toFoundationModels() } - requestContext.tools.toFoundationModels(in: session) + requestContext.toolsForExecution(in: session).toFoundationModels(in: session) } } @@ -420,7 +420,7 @@ let requestContext = session.resolvedRequestContext() return FoundationModels.LanguageModelSession( model: model, - tools: requestContext.tools.toFoundationModels(in: session), + tools: requestContext.toolsForExecution(in: session).toFoundationModels(in: session), transcript: fmTranscriptDroppingDuplicatePrompt( requestContext.transcript, prompt: prompt @@ -744,9 +744,10 @@ Transcript.ToolOutput(id: call.id, toolName: tool.name, segments: segments) } guard let session, let delegate = session.toolExecutionDelegate else { - let text = try await tool.callFunction(content: arguments).promptRepresentation.description - await ToolCallRecord.current?.ran(call, output: output([.text(.init(content: text))])) - return FoundationModels.Prompt(text) + let segments = try await tool.makeOutputSegments(from: arguments) + let prompt = try segments.toFoundationModelsToolOutput() + await ToolCallRecord.current?.ran(call, output: output(segments)) + return prompt } await delegate.didGenerateToolCalls([call], in: session) switch await delegate.toolCallDecision(for: call, in: session) { @@ -761,10 +762,11 @@ return prompt case .execute: do { - let text = try await tool.callFunction(content: arguments).promptRepresentation.description - await delegate.didExecuteToolCall(call, output: output([.text(.init(content: text))]), in: session) - await ToolCallRecord.current?.ran(call, output: output([.text(.init(content: text))])) - return FoundationModels.Prompt(text) + let segments = try await tool.makeOutputSegments(from: arguments) + let prompt = try segments.toFoundationModelsToolOutput() + await delegate.didExecuteToolCall(call, output: output(segments), in: session) + await ToolCallRecord.current?.ran(call, output: output(segments)) + return prompt } catch { await delegate.didFailToolCall(call, error: error, in: session) throw error diff --git a/Sources/AnyLanguageModel/ProfileToolExecution.swift b/Sources/AnyLanguageModel/ProfileToolExecution.swift new file mode 100644 index 00000000..f476b65b --- /dev/null +++ b/Sources/AnyLanguageModel/ProfileToolExecution.swift @@ -0,0 +1,254 @@ +import Foundation + +protocol ProfileRequestIdentifiedLanguageModel { + var profileRequestIdentity: UUID { get } +} + +struct ProfileToolInvocation: Sendable { + let call: Transcript.ToolCall + let output: Transcript.ToolOutput +} + +enum ProfileToolResolution: Sendable { + case stop(calls: [Transcript.ToolCall]) + case invocations([ProfileToolInvocation]) +} + +extension LanguageModelSession { + func profileContinuationResponse( + ifSelectedModelDiffersFrom currentModelIdentity: UUID, + requestContext: RequestContext, + prompt: Prompt, + generating type: Content.Type, + schema: GenerationSchema, + includeSchemaInPrompt: Bool, + prefixEntries: [Transcript.Entry], + prefixUsage: Usage, + prefixText: String + ) async throws -> Response? { + guard requestContext.hasDynamicProfile, !prefixEntries.isEmpty else { return nil } + if let selected = requestContext.model as? any ProfileRequestIdentifiedLanguageModel, + selected.profileRequestIdentity == currentModelIdentity + { + return nil + } + + let includeSchemaInPrompt = profileIncludeSchemaInPrompt( + for: requestContext, + callSiteValue: includeSchemaInPrompt + ) + let response: Response + if type == GeneratedContent.self { + let generated = try await requestContext.model.respond( + within: self, + to: prompt, + schema: schema, + includeSchemaInPrompt: includeSchemaInPrompt, + options: requestContext.options + ) + guard let typed = generated as? Response else { + preconditionFailure("GeneratedContent response type mismatch") + } + response = typed + } else { + response = try await requestContext.model.respond( + within: self, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: requestContext.options + ) + } + + var earlierUsage = prefixUsage + earlierUsage.add(response.requests.earlier) + var content = response.content + var rawContent = response.rawContent + if type == String.self, case .string(let suffix) = response.rawContent.kind { + let combined = prefixText + suffix + content = combined as! Content + rawContent = GeneratedContent(combined) + } + return Response( + content: content, + rawContent: rawContent, + transcriptEntries: ArraySlice(prefixEntries + response.transcriptEntries), + requests: RequestsUsage(earlier: earlierUsage, final: response.requests.final), + providerMetadata: response.providerMetadata + ) + } + + func profileContinuationStream( + ifSelectedModelDiffersFrom currentModelIdentity: UUID, + requestContext: RequestContext, + prompt: Prompt, + generating type: Content.Type, + schema: GenerationSchema, + includeSchemaInPrompt: Bool, + prefixEntries: [Transcript.Entry], + prefixUsage: Usage, + prefixText: String + ) -> ResponseStream? { + guard requestContext.hasDynamicProfile, !prefixEntries.isEmpty else { return nil } + if let selected = requestContext.model as? any ProfileRequestIdentifiedLanguageModel, + selected.profileRequestIdentity == currentModelIdentity + { + return nil + } + + let includeSchemaInPrompt = profileIncludeSchemaInPrompt( + for: requestContext, + callSiteValue: includeSchemaInPrompt + ) + let upstream: ResponseStream + if type == GeneratedContent.self { + let generated = requestContext.model.streamResponse( + within: self, + to: prompt, + schema: schema, + includeSchemaInPrompt: includeSchemaInPrompt, + options: requestContext.options + ) + guard let typed = generated as? ResponseStream else { + preconditionFailure("GeneratedContent response stream type mismatch") + } + upstream = typed + } else { + upstream = requestContext.model.streamResponse( + within: self, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: requestContext.options + ) + } + + let stream = AsyncThrowingStream.Snapshot, any Error> { continuation in + let task = Task { + do { + for try await snapshot in upstream { + var earlierUsage = prefixUsage + earlierUsage.add(snapshot.requests.earlier) + var content = snapshot.content + var rawContent = snapshot.rawContent + if type == String.self, case .string(let suffix) = snapshot.rawContent.kind { + let combined = prefixText + suffix + content = combined as! Content.PartiallyGenerated + rawContent = GeneratedContent(combined) + } + continuation.yield( + .init( + content: content, + rawContent: rawContent, + transcriptEntries: ArraySlice(prefixEntries + snapshot.transcriptEntries), + requests: RequestsUsage(earlier: earlierUsage, final: snapshot.requests.final), + providerMetadata: snapshot.providerMetadata + ) + ) + } + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in task.cancel() } + } + return ResponseStream(stream: stream) + } + + func resolveProfileToolCalls( + _ calls: [Transcript.ToolCall], + requestContext: RequestContext, + currentRoundEntries: [Transcript.Entry], + offeredTools: [any Tool]? = nil, + missingToolOutput: @Sendable (Transcript.ToolCall) -> [Transcript.Segment] = { call in + [.text(.init(content: "Tool not found: \(call.toolName)"))] + } + ) async throws -> ProfileToolResolution { + guard !calls.isEmpty else { return .invocations([]) } + + if let delegate = toolExecutionDelegate { + await delegate.didGenerateToolCalls(calls, in: self) + } + + var decisions: [ToolExecutionDecision] = [] + decisions.reserveCapacity(calls.count) + for call in calls { + try await profileWillExecuteToolCall( + call, + requestContext: requestContext, + currentRoundEntries: currentRoundEntries + ) + let decision = + await toolExecutionDelegate?.toolCallDecision(for: call, in: self) ?? .execute + if case .stop = decision { + return .stop(calls: calls) + } + decisions.append(decision) + } + + var toolsByName: [String: any Tool] = [:] + for tool in requestContext.toolsForExecution(in: self, offeredTools: offeredTools) + where toolsByName[tool.name] == nil { + toolsByName[tool.name] = tool + } + + var invocations: [ProfileToolInvocation] = [] + invocations.reserveCapacity(calls.count) + var completedEntries: [Transcript.Entry] = [] + for (call, decision) in zip(calls, decisions) { + try Task.checkCancellation() + let output: Transcript.ToolOutput + switch decision { + case .stop: + return .stop(calls: calls) + case .provideOutput(let segments): + output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: segments + ) + case .execute: + guard let tool = toolsByName[call.toolName] else { + output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: missingToolOutput(call) + ) + break + } + do { + let segments = try await withProfileToolExecution( + requestContext: requestContext, + currentRoundEntries: currentRoundEntries + completedEntries + ) { + try await tool.makeOutputSegments(from: call.arguments) + } + output = Transcript.ToolOutput( + id: call.id, + toolName: tool.name, + segments: segments + ) + } catch { + if let delegate = toolExecutionDelegate { + await delegate.didFailToolCall(call, error: error, in: self) + } + throw ToolCallError(tool: tool, underlyingError: error) + } + } + + invocations.append(.init(call: call, output: output)) + completedEntries.append(.toolOutput(output)) + try await profileDidProduceToolOutput( + for: call, + output: output, + requestContext: requestContext, + currentRoundEntries: currentRoundEntries + completedEntries + ) + if let delegate = toolExecutionDelegate { + await delegate.didExecuteToolCall(call, output: output, in: self) + } + } + return .invocations(invocations) + } +} diff --git a/Sources/AnyLanguageModel/SessionProperties.swift b/Sources/AnyLanguageModel/SessionProperties.swift new file mode 100644 index 00000000..54a9893c --- /dev/null +++ b/Sources/AnyLanguageModel/SessionProperties.swift @@ -0,0 +1,295 @@ +import Foundation +import Observation + +extension Transcript { + /// The mutable, instruction-free history exposed through session properties. + /// + /// - Note: This API is exclusive to AnyLanguageModel on OS 26. It mirrors + /// Foundation Models 27's `Transcript.HistoryView` role with value semantics. + public struct HistoryView: RandomAccessCollection, MutableCollection, RangeReplaceableCollection, + ExpressibleByArrayLiteral, Sendable + { + public typealias Element = Entry + public typealias SubSequence = HistoryView + public typealias Indices = Range + private var entries: ArraySlice + + public struct Index: Comparable, Hashable, Strideable, Sendable { + fileprivate let offset: Int + + public static func < (lhs: Self, rhs: Self) -> Bool { lhs.offset < rhs.offset } + public func distance(to other: Self) -> Int { other.offset - offset } + public func advanced(by n: Int) -> Self { Self(offset: offset + n) } + } + + public init() { entries = [] } + public init(_ elements: S) where S.Element == Entry { entries = Array(elements)[...] } + public init(arrayLiteral elements: Entry...) { entries = elements[...] } + private init(entries: ArraySlice) { self.entries = entries } + + public var startIndex: Index { Index(offset: entries.startIndex) } + public var endIndex: Index { Index(offset: entries.endIndex) } + public var indices: Indices { startIndex ..< endIndex } + public func index(after i: Index) -> Index { i.advanced(by: 1) } + public func index(before i: Index) -> Index { i.advanced(by: -1) } + public func index(_ i: Index, offsetBy distance: Int) -> Index { i.advanced(by: distance) } + public func distance(from start: Index, to end: Index) -> Int { start.distance(to: end) } + + public subscript(position: Index) -> Entry { + get { entries[position.offset] } + set { entries[position.offset] = newValue } + } + + public subscript(bounds: Range) -> HistoryView { + get { HistoryView(entries: entries[bounds.lowerBound.offset ..< bounds.upperBound.offset]) } + set { replaceSubrange(bounds, with: newValue) } + } + + public mutating func replaceSubrange(_ subrange: Range, with newElements: C) + where C.Element == Entry { + entries.replaceSubrange(subrange.lowerBound.offset ..< subrange.upperBound.offset, with: newElements) + } + } +} + +/// A type-keyed session property. +/// +/// - Note: This API is exclusive to AnyLanguageModel on OS 26 and mirrors the +/// Foundation Models 27 session-property contract. +public protocol SessionPropertyKey: SendableMetatype { + associatedtype Value + static var defaultValue: Value { get } +} + +/// Declares a stored value on ``SessionPropertyValues``. +@attached(accessor) +@attached(peer, names: prefixed(__Key_)) +public macro SessionPropertyEntry() = + #externalMacro(module: "AnyLanguageModelMacros", type: "SessionPropertyEntryMacro") + +/// Session-scoped storage shared by profiles, dynamic instructions, and tools. +/// +/// - Note: This API is exclusive to AnyLanguageModel on OS 26 and mirrors +/// Foundation Models 27. +@Observable +public final class SessionPropertyValues: @unchecked Sendable { + @ObservationIgnored private let sessionPropertyRegistrar = ObservationRegistrar() + @ObservationIgnored private let customValues = Locked<[ObjectIdentifier: Any]>([:]) + @ObservationIgnored private let historyGetter: @Sendable () -> [Transcript.Entry] + @ObservationIgnored private let historySetter: @Sendable ([Transcript.Entry]) -> Void + + init( + historyGetter: @escaping @Sendable () -> [Transcript.Entry], + historySetter: @escaping @Sendable ([Transcript.Entry]) -> Void + ) { + self.historyGetter = historyGetter + self.historySetter = historySetter + } + + /// The canonical session history after dynamic instructions. + public var history: Transcript.HistoryView { + get { + _accessSessionProperty(keyPath: \.history) + if SessionPropertyBinding.values === self, let history = SessionPropertyBinding.history { + return Transcript.HistoryView(history.entries) + } + return Transcript.HistoryView(historyGetter()) + } + set { + _withSessionPropertyMutation(keyPath: \.history) { + let newValue = Array(newValue).filter { entry in + if case .instructions = entry { return false } + return true + } + if SessionPropertyBinding.values === self, let history = SessionPropertyBinding.history { + precondition( + history.isWritable, + "Session history is read-only while dynamic instructions or a Tool is active" + ) + history.entries = newValue + history.persist(history.entries) + } else { + historySetter(newValue) + } + } + } + } + + public subscript(key: Key.Type) -> Key.Value where Key: SessionPropertyKey { + get { + if let value = customValues.withLock({ $0[ObjectIdentifier(key)] as? Key.Value }) { + return value + } + let defaultValue = Key.defaultValue + return customValues.withLock { values in + if let value = values[ObjectIdentifier(key)] as? Key.Value { return value } + values[ObjectIdentifier(key)] = defaultValue + return defaultValue + } + } + set { + customValues.withLock { values in + values[ObjectIdentifier(key)] = newValue + } + } + } + + func historyBinding( + _ entries: [Transcript.Entry], + isWritable: Bool, + protecting protectedEntryKeys: Set = [], + onChange: (@Sendable ([Transcript.Entry]) -> Void)? = nil + ) -> SessionHistoryBinding { + let protected = keyedTranscriptEntries(entries).filter { protectedEntryKeys.contains($0.key) }.map(\.entry) + let persist: (@Sendable ([Transcript.Entry]) -> Void)? + if isWritable { + persist = { [historySetter] updated in + onChange?(updated) + historySetter( + removingProtectedTranscriptEntries(from: updated, protecting: protected) + ) + } + } else { + persist = nil + } + return SessionHistoryBinding( + entries, + isWritable: isWritable, + protectedEntryKeys: protectedEntryKeys, + persist: persist + ) + } + + @_documentation(visibility: internal) + public func _accessSessionProperty( + keyPath: KeyPath + ) { + sessionPropertyRegistrar.access(self, keyPath: keyPath) + } + + @_documentation(visibility: internal) + public func _withSessionPropertyMutation( + keyPath: KeyPath, + _ mutation: () throws -> Result + ) rethrows -> Result { + try sessionPropertyRegistrar.withMutation( + of: self, + keyPath: keyPath, + mutation + ) + } +} + +enum SessionPropertyBinding { + @TaskLocal static var values: SessionPropertyValues? + @TaskLocal static var history: SessionHistoryBinding? +} + +final class SessionHistoryBinding: @unchecked Sendable { + private let storage: Locked<[Transcript.Entry]> + let isWritable: Bool + private let protectedEntries: [Transcript.Entry] + private let persistValue: (@Sendable ([Transcript.Entry]) -> Void)? + + init( + _ entries: [Transcript.Entry], + isWritable: Bool, + protectedEntryKeys: Set = [], + persist: (@Sendable ([Transcript.Entry]) -> Void)? = nil + ) { + storage = Locked( + entries.filter { entry in + if case .instructions = entry { return false } + return true + } + ) + self.isWritable = isWritable + self.protectedEntries = keyedTranscriptEntries(entries).filter { protectedEntryKeys.contains($0.key) }.map( + \.entry + ) + self.persistValue = persist + } + + var entries: [Transcript.Entry] { + get { storage.withLock { $0 } } + set { + storage.withLock { current in + current = + removingProtectedTranscriptEntries(from: newValue, protecting: protectedEntries) + + protectedEntries + } + } + } + + func persist(_ entries: [Transcript.Entry]) { + persistValue?(entries) + } +} + +extension LanguageModelSession { + /// Accesses a value in the currently evaluating session. + @propertyWrapper + public struct SessionProperty { + private let keyPath: ReferenceWritableKeyPath + + public init(_ keyPath: ReferenceWritableKeyPath) { + self.keyPath = keyPath + } + + public var wrappedValue: Value { + get { + guard let values = SessionPropertyBinding.values else { + preconditionFailure( + "SessionProperty can only be accessed while a session profile, dynamic instructions, or a Tool is active" + ) + } + return values[keyPath: keyPath] + } + nonmutating set { + guard let values = SessionPropertyBinding.values else { + preconditionFailure( + "SessionProperty can only be accessed while a session profile, dynamic instructions, or a Tool is active" + ) + } + values[keyPath: keyPath] = newValue + } + } + } +} + +extension LanguageModelSession.SessionProperty: @unchecked Sendable where Value: Sendable {} + +extension DynamicInstructions { + public typealias SessionProperty = LanguageModelSession.SessionProperty +} + +extension Tool { + public typealias SessionProperty = LanguageModelSession.SessionProperty + + func bindingSessionProperties(_ properties: SessionPropertyValues, history: [Transcript.Entry]) -> any Tool { + SessionPropertyTool(wrapped: self, properties: properties, history: history) + } +} + +private struct SessionPropertyTool: Tool { + let wrapped: Wrapped + let properties: SessionPropertyValues + let history: [Transcript.Entry] + + var name: String { wrapped.name } + var description: String { wrapped.description } + var parameters: GenerationSchema { wrapped.parameters } + var includesSchemaInInstructions: Bool { wrapped.includesSchemaInInstructions } + + func call(arguments: Wrapped.Arguments) async throws -> Wrapped.Output { + let entries = + SessionPropertyBinding.values === properties + ? SessionPropertyBinding.history?.entries ?? history : history + let history = SessionHistoryBinding(entries, isWritable: false) + return try await SessionPropertyBinding.$values.withValue(properties) { + try await SessionPropertyBinding.$history.withValue(history) { + try await wrapped.call(arguments: arguments) + } + } + } +} diff --git a/Sources/AnyLanguageModel/Shared/ToolContinuation.swift b/Sources/AnyLanguageModel/Shared/ToolContinuation.swift new file mode 100644 index 00000000..ae84d9de --- /dev/null +++ b/Sources/AnyLanguageModel/Shared/ToolContinuation.swift @@ -0,0 +1,79 @@ +enum TranscriptEntryKey: Hashable, Sendable { + case entry(String) + case toolOutput(round: String?, call: String) +} + +// Provider call IDs are scoped to their Tool round, not to the whole transcript. +func keyedTranscriptEntries(_ entries: [Transcript.Entry]) -> [(key: TranscriptEntryKey, entry: Transcript.Entry)] { + var round: String? + return entries.map { entry in + let key: TranscriptEntryKey + switch entry { + case .toolCalls(let calls): + round = calls.id + key = .entry(entry.id) + case .toolOutput: + key = .toolOutput(round: round, call: entry.id) + case .instructions, .prompt, .response: + round = nil + key = .entry(entry.id) + case .reasoning: + key = .entry(entry.id) + } + return (key, entry) + } +} + +func mergingTranscriptEntries( + _ original: [Transcript.Entry], + with additional: [(key: TranscriptEntryKey, entry: Transcript.Entry)] +) -> [Transcript.Entry] { + var result = original + var indices: [TranscriptEntryKey: Int] = [:] + for (index, item) in keyedTranscriptEntries(original).enumerated() { indices[item.key] = index } + for item in additional { + if let index = indices[item.key] { + result[index] = item.entry + } else { + indices[item.key] = result.count + result.append(item.entry) + } + } + return result +} + +func removingProtectedTranscriptEntries( + from entries: [Transcript.Entry], + protecting protected: [Transcript.Entry] +) -> [Transcript.Entry] { + var result = entries + for entry in protected.reversed() { + if let index = result.lastIndex(of: entry) { result.remove(at: index) } + } + return result +} + +func toolContinuationMessages( + from transcript: Transcript, + replacing currentRoundEntries: [Transcript.Entry], + with continuation: [Message], + hasDynamicProfile: Bool, + convert: (Transcript) throws -> [Message] +) rethrows -> [Message] { + guard hasDynamicProfile else { return try convert(transcript) + continuation } + guard !continuation.isEmpty, !currentRoundEntries.isEmpty else { return try convert(transcript) } + let projected = Array(transcript) + let count = currentRoundEntries.count + let ranges = projected.indices.compactMap { start -> Range? in + let end = start + count + guard end <= projected.count, Array(projected[start ..< end]) == currentRoundEntries else { return nil } + return start ..< end + } + guard ranges.count == 1, let range = ranges.first else { + // A profile may remove or reorder the current round. Its projection is authoritative. + return try convert(transcript) + } + return try convert(Transcript(entries: projected[.. maximumRounds { - throw LanguageModelSession.ToolRoundLimitExceeded( - limit: maximumRounds, - provider: provider, - transcriptEntries: ArraySlice(completed) - ) + func record(completed: [Transcript.Entry]) throws { + try roundCount.withLock { rounds in + rounds += 1 + if rounds > maximumRounds { + throw LanguageModelSession.ToolRoundLimitExceeded( + limit: maximumRounds, + provider: provider, + transcriptEntries: ArraySlice(completed) + ) + } } } } diff --git a/Sources/AnyLanguageModel/Tool.swift b/Sources/AnyLanguageModel/Tool.swift index bdca23df..22fcd645 100644 --- a/Sources/AnyLanguageModel/Tool.swift +++ b/Sources/AnyLanguageModel/Tool.swift @@ -1,3 +1,9 @@ +import Foundation + +#if canImport(UniformTypeIdentifiers) + import UniformTypeIdentifiers +#endif + /// A tool that a model can call to gather information at runtime or perform side effects. /// /// Tool calling gives the model the ability to call your code to incorporate @@ -110,7 +116,10 @@ extension Tool where Self.Arguments: Generable { // MARK: - Helpers extension Tool { - internal func makeOutputSegments(from arguments: GeneratedContent) async throws -> [Transcript.Segment] { + /// Compatibility integration SPI for provider adapters outside this module. + /// This is not part of Apple's public Tool contract. + @_spi(Compatibility) + public func makeOutputSegments(from arguments: GeneratedContent) async throws -> [Transcript.Segment] { let parsedArguments = try Arguments(arguments) // What a tool asks of a model is its own: it isn't counted for the response that called it. let output = try await RequestUsageChannel.outside { try await call(arguments: parsedArguments) } @@ -125,7 +134,47 @@ extension Tool { return [Transcript.Segment.text(.init(content: stringOutput))] } - let fallback = output.promptRepresentation.description - return [Transcript.Segment.text(.init(content: fallback))] + return try output.promptRepresentation.makeTranscriptSegments().map { segment in + guard case .image(let image) = segment, + case .url(let url) = image.source, + url.isFileURL + else { return segment } + + let mimeType: String? + #if canImport(UniformTypeIdentifiers) + guard let type = UTType(filenameExtension: url.pathExtension) else { + throw ToolAttachmentError.unknownImageType(url) + } + guard type.conforms(to: .image) else { + throw ToolAttachmentError.unsupportedImageType(url) + } + mimeType = type.preferredMIMEType + #else + mimeType = + switch url.pathExtension.lowercased() { + case "jpg", "jpeg": "image/jpeg" + case "png": "image/png" + case "gif": "image/gif" + case "webp": "image/webp" + default: nil + } + #endif + guard let mimeType else { throw ToolAttachmentError.unknownImageType(url) } + return .image(.init(id: image.id, data: try Data(contentsOf: url), mimeType: mimeType)) + } + } +} + +private enum ToolAttachmentError: LocalizedError { + case unknownImageType(URL) + case unsupportedImageType(URL) + + var errorDescription: String? { + switch self { + case .unknownImageType(let url): + "Cannot determine the MIME type of the Tool output image \(url.lastPathComponent)." + case .unsupportedImageType(let url): + "The Tool output attachment \(url.lastPathComponent) is not an image file." + } } } diff --git a/Sources/AnyLanguageModelMacros/Plugin.swift b/Sources/AnyLanguageModelMacros/Plugin.swift index 917d2deb..dee43755 100644 --- a/Sources/AnyLanguageModelMacros/Plugin.swift +++ b/Sources/AnyLanguageModelMacros/Plugin.swift @@ -6,5 +6,6 @@ struct AnyLanguageModelMacrosPlugin: CompilerPlugin { let providingMacros: [any Macro.Type] = [ GenerableMacro.self, GuideMacro.self, + SessionPropertyEntryMacro.self, ] } diff --git a/Sources/AnyLanguageModelMacros/SessionPropertyEntryMacro.swift b/Sources/AnyLanguageModelMacros/SessionPropertyEntryMacro.swift new file mode 100644 index 00000000..e6625660 --- /dev/null +++ b/Sources/AnyLanguageModelMacros/SessionPropertyEntryMacro.swift @@ -0,0 +1,74 @@ +import SwiftSyntax +import SwiftSyntaxBuilder +import SwiftSyntaxMacros + +/// Synthesizes a `SessionPropertyKey` and forwards the declared property through +/// `SessionPropertyValues`' type-keyed storage. +public struct SessionPropertyEntryMacro: AccessorMacro, PeerMacro { + public static func expansion( + of node: AttributeSyntax, + providingAccessorsOf declaration: some DeclSyntaxProtocol, + in context: some MacroExpansionContext + ) throws -> [AccessorDeclSyntax] { + guard let name = propertyName(in: declaration) else { return [] } + return [ + """ + get { + self._accessSessionProperty(keyPath: \\.\(raw: name)) + return self[__Key_\(raw: name).self] + } + """, + """ + set { + self._withSessionPropertyMutation(keyPath: \\.\(raw: name)) { + self[__Key_\(raw: name).self] = newValue + } + } + """, + ] + } + + public static func expansion( + of node: AttributeSyntax, + providingPeersOf declaration: some DeclSyntaxProtocol, + in context: some MacroExpansionContext + ) throws -> [DeclSyntax] { + guard + let variable = declaration.as(VariableDeclSyntax.self), + variable.bindings.count == 1, + let binding = variable.bindings.first, + let name = binding.pattern.as(IdentifierPatternSyntax.self)?.identifier.text, + let initializer = binding.initializer?.value + else { return [] } + + let defaultValue: DeclSyntax + if let typeAnnotation = binding.typeAnnotation { + defaultValue = + """ + private struct __Key_\(raw: name): SessionPropertyKey { + static let defaultValue \(typeAnnotation) = \(initializer) + } + """ + } else { + defaultValue = + """ + private struct __Key_\(raw: name): SessionPropertyKey { + static let defaultValue = \(initializer) + } + """ + } + return [defaultValue] + } + + private static func propertyName(in declaration: some DeclSyntaxProtocol) -> String? { + guard + let variable = declaration.as(VariableDeclSyntax.self), + variable.bindings.count == 1, + let binding = variable.bindings.first, + binding.accessorBlock == nil, + let name = binding.pattern.as(IdentifierPatternSyntax.self)?.identifier.text, + binding.initializer != nil + else { return nil } + return name + } +} diff --git a/Tests/AnyLanguageModelTests/APICompatibilityAnyLanguageModelTests.swift b/Tests/AnyLanguageModelTests/APICompatibilityAnyLanguageModelTests.swift index 0d7dc600..9d2eb3bb 100644 --- a/Tests/AnyLanguageModelTests/APICompatibilityAnyLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/APICompatibilityAnyLanguageModelTests.swift @@ -10,18 +10,78 @@ import Testing return false }() + @Test func anyLanguageModelHistoryViewCollectionCompatibility() { + func requireContract(_ history: C) + where + C: RandomAccessCollection & MutableCollection & RangeReplaceableCollection + & ExpressibleByArrayLiteral & Sendable, + C.Element == Transcript.Entry, C.SubSequence == C, C.Index: Hashable & Strideable, + C.Index.Stride == Int + {} + var history: Transcript.HistoryView = [] + let empty: Transcript.HistoryView = history[history.startIndex ..< history.endIndex] + history[history.startIndex ..< history.endIndex] = empty + history.append(contentsOf: empty) + history.replaceSubrange(history.startIndex ..< history.endIndex, with: empty) + requireContract(history) + #expect(history.isEmpty) + } + + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) + extension SessionPropertyValues { + @SessionPropertyEntry var compatibilityRevision = 0 + } + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) private struct CompatibilityDynamicInstructions: DynamicInstructions { let includeDetail: Bool + @SessionProperty(\.compatibilityRevision) private var revision var body: some DynamicInstructions { - Instructions("You are a helpful assistant.") + Instructions("You are a helpful assistant. Revision \(revision).") if includeDetail { Instructions("Include useful detail.") } } } + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) + private struct CompatibilityDynamicProfile: LanguageModelSession.DynamicProfile { + let model: any LanguageModel + @SessionProperty(\.compatibilityRevision) private var revision + + var body: some LanguageModelSession.DynamicProfile { + Profile { + Instructions("Revision \(revision)") + } + .model(model) + .samplingMode(.greedy) + .temperature(0.2) + .maximumResponseTokens(32) + .reasoningLevel(.light) + .toolCallingMode(.allowed) + .historyTransform { $0 } + .transcriptErrorHandlingPolicy(nil) + } + } + + @Test func anyLanguageModelNonSendableHistoryTransformCompatibility() { + let captured = NonSendableHistoryCapture() + let transform: ([Transcript.Entry]) -> [Transcript.Entry] = { entries in + captured.calls += 1 + return entries + } + _ = LanguageModelSession.AnyDynamicProfile( + LanguageModelSession.Profile { Instructions("Projection") } + .historyTransform(transform) + ) + #expect(captured.calls == 0) + } + + private final class NonSendableHistoryCapture { + var calls = 0 + } + @available(macOS 26.0, iOS 26.0, tvOS 26.0, visionOS 26.0, *) @Test("AnyLanguageModel Drop-In Compatibility", .enabled(if: isSystemLanguageModelAvailable)) func anyLanguageModelCompatibility() async throws { @@ -32,11 +92,14 @@ import Testing ) if #available(macOS 27.0, iOS 27.0, visionOS 27.0, *) { - _ = LanguageModelSession( + let dynamicSession = LanguageModelSession( model: model, dynamicInstructions: CompatibilityDynamicInstructions(includeDetail: true), history: session.transcript ) + dynamicSession.properties.compatibilityRevision = 1 + _ = dynamicSession.properties.history + _ = LanguageModelSession.AnyDynamicProfile(CompatibilityDynamicProfile(model: model)) } let options = GenerationOptions(temperature: 0.7) diff --git a/Tests/AnyLanguageModelTests/APICompatibilityFoundationModelsTests.swift b/Tests/AnyLanguageModelTests/APICompatibilityFoundationModelsTests.swift index 79161f72..8901e486 100644 --- a/Tests/AnyLanguageModelTests/APICompatibilityFoundationModelsTests.swift +++ b/Tests/AnyLanguageModelTests/APICompatibilityFoundationModelsTests.swift @@ -12,17 +12,79 @@ import Testing #if compiler(>=6.4) #if os(macOS) || os(iOS) || os(visionOS) + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) + @Test func foundationModelsHistoryViewCollectionCompatibility() { + func requireContract(_ history: C) + where + C: RandomAccessCollection & MutableCollection & RangeReplaceableCollection + & ExpressibleByArrayLiteral & Sendable, + C.Element == Transcript.Entry, C.SubSequence == C, C.Index: Hashable & Strideable, + C.Index.Stride == Int + {} + var history: Transcript.HistoryView = [] + let empty: Transcript.HistoryView = history[history.startIndex ..< history.endIndex] + history[history.startIndex ..< history.endIndex] = empty + history.append(contentsOf: empty) + history.replaceSubrange(history.startIndex ..< history.endIndex, with: empty) + requireContract(history) + #expect(history.isEmpty) + } + + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) + extension SessionPropertyValues { + @SessionPropertyEntry var compatibilityRevision = 0 + } + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) private struct CompatibilityDynamicInstructions: DynamicInstructions { let includeDetail: Bool + @SessionProperty(\.compatibilityRevision) private var revision var body: some DynamicInstructions { - Instructions("You are a helpful assistant.") + Instructions("You are a helpful assistant. Revision \(revision).") if includeDetail { Instructions("Include useful detail.") } } } + + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) + private struct CompatibilityDynamicProfile: LanguageModelSession.DynamicProfile { + let model: any LanguageModel + @SessionProperty(\.compatibilityRevision) private var revision + + var body: some LanguageModelSession.DynamicProfile { + Profile { + Instructions("Revision \(revision)") + } + .model(model) + .samplingMode(.greedy) + .temperature(0.2) + .maximumResponseTokens(32) + .reasoningLevel(.light) + .toolCallingMode(.allowed) + .historyTransform { $0 } + .transcriptErrorHandlingPolicy(nil) + } + } + + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) + @Test func foundationModelsNonSendableHistoryTransformCompatibility() { + let captured = NonSendableHistoryCapture() + let transform: ([Transcript.Entry]) -> [Transcript.Entry] = { entries in + captured.calls += 1 + return entries + } + _ = LanguageModelSession.AnyDynamicProfile( + LanguageModelSession.Profile { Instructions("Projection") } + .historyTransform(transform) + ) + #expect(captured.calls == 0) + } + + private final class NonSendableHistoryCapture { + var calls = 0 + } #endif #endif @@ -41,11 +103,14 @@ import Testing #if compiler(>=6.4) #if os(macOS) || os(iOS) || os(visionOS) if #available(macOS 27.0, iOS 27.0, visionOS 27.0, *) { - _ = LanguageModelSession( + let dynamicSession = LanguageModelSession( model: model, dynamicInstructions: CompatibilityDynamicInstructions(includeDetail: true), history: session.transcript ) + dynamicSession.properties.compatibilityRevision = 1 + _ = dynamicSession.properties.history + _ = LanguageModelSession.AnyDynamicProfile(CompatibilityDynamicProfile(model: model)) } #endif #endif diff --git a/Tests/AnyLanguageModelTests/DynamicInstructionsTests.swift b/Tests/AnyLanguageModelTests/DynamicInstructionsTests.swift index af43ad37..8943439b 100644 --- a/Tests/AnyLanguageModelTests/DynamicInstructionsTests.swift +++ b/Tests/AnyLanguageModelTests/DynamicInstructionsTests.swift @@ -1,7 +1,7 @@ import Foundation import Testing -@testable import AnyLanguageModel +@_spi(Compatibility) @testable import AnyLanguageModel @Suite("Dynamic instructions") struct DynamicInstructionsTests { @@ -407,7 +407,7 @@ private struct DynamicContextModel: LanguageModel { var entries: [Transcript.Entry] = [] if continuesAfterTool { - let tool = try #require(first.tools.first) + let tool = try #require(first.toolsForExecution(in: session).first) state.select(.b) let call = Transcript.ToolCall( id: "request-a-call", @@ -475,7 +475,7 @@ private struct FailingAfterToolModel: LanguageModel { return !didFail } if shouldFail { - let tool = try #require(context.tools.first) + let tool = try #require(context.toolsForExecution(in: session).first) _ = try await tool.makeOutputSegments(from: GeneratedContent(properties: [:])) state.select(.b) throw DynamicFixtureError.failed @@ -565,7 +565,7 @@ private struct CancellingAfterToolModel: LanguageModel { return !didSuspend } if shouldSuspend { - let tool = try #require(context.tools.first) + let tool = try #require(context.toolsForExecution(in: session).first) _ = try await tool.makeOutputSegments(from: GeneratedContent(properties: [:])) state.select(.b) control.signalStarted() diff --git a/Tests/AnyLanguageModelTests/DynamicProfileTests.swift b/Tests/AnyLanguageModelTests/DynamicProfileTests.swift new file mode 100644 index 00000000..45431adc --- /dev/null +++ b/Tests/AnyLanguageModelTests/DynamicProfileTests.swift @@ -0,0 +1,132 @@ +import Testing + +@testable import AnyLanguageModel + +@Suite("Dynamic Profile Composition") +struct DynamicProfileBuilderTests { + @Test func resolvesInstructionsToolsAndModifiers() { + let profile = LanguageModelSession.AnyDynamicProfile( + LanguageModelSession.Profile { + Instructions("Use this profile") + ProfileBuilderTool() + } + .model(MockLanguageModel()) + .temperature(0.2) + .samplingMode(.greedy) + .maximumResponseTokens(32) + .reasoningLevel(.light) + .toolCallingMode(.allowed) + .transcriptErrorHandlingPolicy(nil) + ).resolveForRequest() + #expect(profile.model is MockLanguageModel) + #expect(profile.instructions?.description == "Use this profile") + #expect(profile.tools.map(\.name) == ["builder"]) + #expect(profile.temperature == 0.2) + #expect(profile.hasTemperature) + #expect(profile.samplingMode == .greedy) + #expect(profile.maximumResponseTokens == 32) + #expect(profile.reasoningLevel == .light) + #expect(profile.toolCallingMode == .allowed) + #expect(profile.hasTranscriptErrorHandlingPolicy) + } + + @Test func conditionalAndErasedProfilesRetainBranchIdentity() { + let first = LanguageModelSession.AnyDynamicProfile(ConditionalBuilderProfile(first: true)) + .resolveForRequest() + let second = LanguageModelSession.AnyDynamicProfile( + LanguageModelSession.AnyDynamicProfile(ConditionalBuilderProfile(first: false)) + ).resolveForRequest() + #expect(first.instructions?.description == "First") + #expect(second.instructions?.description == "Second") + #expect(first.identity != second.identity) + } + + @Test func customModifierComposesWithBuiltinModifiers() { + let profile = LanguageModelSession.AnyDynamicProfile( + LanguageModelSession.Profile { Instructions("Custom") } + .modifier(BuilderModifier()) + .maximumResponseTokens(64) + ).resolveForRequest() + #expect(profile.temperature == 0.4) + #expect(profile.maximumResponseTokens == 64) + } + + @Test func historyTransformCanCaptureNonSendableState() { + let recorder = HistoryTransformRecorder() + let profile = LanguageModelSession.AnyDynamicProfile( + LanguageModelSession.Profile { Instructions("Projection") } + .historyTransform { entries in + recorder.entryIDs = entries.map(\.id) + return entries + } + ).resolveForRequest() + let history: [Transcript.Entry] = [.prompt(.init(segments: [.text(.init(content: "Hello"))]))] + #expect(profile.historyTransforms[0](history) == history) + #expect(recorder.entryIDs == history.map(\.id)) + } + + @Test func historyTransformsAndCallbacksKeepCompositionOrder() async throws { + let events = Locked<[String]>([]) + let profile = LanguageModelSession.AnyDynamicProfile( + LanguageModelSession.Profile { Instructions("Callbacks") } + .historyTransform { entries in + events.withLock { $0.append("inner transform") } + return entries + } + .onPrompt { _ in events.withLock { $0.append("inner prompt") } } + .onActivate { events.withLock { $0.append("inner activate") } } + .onDeactivate { events.withLock { $0.append("inner deactivate") } } + .historyTransform { entries in + events.withLock { $0.append("outer transform") } + return entries + } + .onPrompt { _ in events.withLock { $0.append("outer prompt") } } + .onActivate { events.withLock { $0.append("outer activate") } } + .onDeactivate { events.withLock { $0.append("outer deactivate") } } + ).resolveForRequest() + var history: [Transcript.Entry] = [] + for transform in profile.historyTransforms { history = transform(history) } + for action in profile.onActivate { await action() } + let prompt = Transcript.Prompt(segments: [.text(.init(content: "Hello"))]) + for action in profile.onPrompt { try await action(prompt) } + for action in profile.onDeactivate { await action() } + #expect( + events.withLock { $0 } == [ + "outer transform", "inner transform", + "outer activate", "inner activate", + "outer prompt", "inner prompt", + "inner deactivate", "outer deactivate", + ] + ) + } +} + +private final class HistoryTransformRecorder { + var entryIDs: [String] = [] +} + +private struct ConditionalBuilderProfile: LanguageModelSession.DynamicProfile { + let first: Bool + + var body: some LanguageModelSession.DynamicProfile { + if first { + Profile { Instructions("First") } + } else { + Profile { Instructions("Second") } + } + } +} + +private struct BuilderModifier: LanguageModelSession.DynamicProfileModifier { + func body(content: Content) -> some DynamicProfile { + content.temperature(0.4) + } +} + +private struct ProfileBuilderTool: Tool { + let name = "builder" + let description = "Builder fixture" + @Generable struct Arguments {} + + func call(arguments: Arguments) async throws -> String { "done" } +} diff --git a/Tests/AnyLanguageModelTests/GeminiProfileTests.swift b/Tests/AnyLanguageModelTests/GeminiProfileTests.swift new file mode 100644 index 00000000..d06746c9 --- /dev/null +++ b/Tests/AnyLanguageModelTests/GeminiProfileTests.swift @@ -0,0 +1,42 @@ +import Testing + +@_spi(Compatibility) @testable import AnyLanguageModel + +#if canImport(Darwin) && !canImport(AsyncHTTPClient) + extension ProviderUsageTests { + @Test(arguments: [Provider.gemini], [false, true]) + func geminiProfileReevaluation(_ provider: Provider, _ streaming: Bool) async throws { + try await providerProfileReevaluatesBeforeEveryToolContinuation(provider, streaming) + } + + @Test(arguments: [Provider.gemini], [false, true]) + func geminiModelHandoff(_ provider: Provider, _ streaming: Bool) async throws { + try await providerProfileCanSwitchModelForToolContinuation(provider, streaming) + } + + @Test(arguments: [Provider.gemini], [false, true]) + func geminiSchemaPromptHandoff(_ provider: Provider, _ streaming: Bool) async throws { + try await modelHandoffsRestoreOriginalSchemaPromptOption(provider, streaming) + } + + @Test(arguments: [Provider.gemini], [false, true]) + func geminiHistoryAcrossHandoffs(_ provider: Provider, _ streaming: Bool) async throws { + try await modelHandoffsPreserveEarlierToolResultsAndHistoryWrites(provider, streaming) + } + + @Test(arguments: [Provider.gemini], [false, true]) + func geminiSharedToolRoundBudget(_ provider: Provider, _ streaming: Bool) async throws { + try await providerModelHandoffsShareToolRoundLimit(provider, streaming) + } + + @Test(arguments: [Provider.gemini], [false, true]) + func geminiFailedHandoffHistory(_ provider: Provider, _ streaming: Bool) async throws { + try await failedHandoffDoesNotPersistPendingHistory(provider, streaming) + } + + @Test(arguments: [Provider.gemini], [false, true]) + func geminiOrdinaryToolLoop(_ provider: Provider, _ streaming: Bool) async throws { + try await ordinaryWrapperKeepsTheProviderToolLoop(provider, streaming) + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift b/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift index 5265718f..6ee026b6 100644 --- a/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift +++ b/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift @@ -4,11 +4,27 @@ import Testing @testable import AnyLanguageModel #if canImport(Darwin) && !canImport(AsyncHTTPClient) + import CoreGraphics @Suite("GeminiLanguageModel thought signatures", .serialized) struct GeminiThoughtSignatureTests { private static let signature = "CvsBAdHtim8n5xQK1pVX2H0lPQeXAMPLEsignature==" + @Generable + struct ImageArguments { + var request: String + } + + private struct ImageOutputTool: Tool { + let name = "image" + let description = "Return an image" + let output: Prompt + + func call(arguments: ImageArguments) async throws -> Prompt { + output + } + } + private func makeModel() -> GeminiLanguageModel { GeminiLanguageModel( apiKey: "test-key", @@ -51,6 +67,21 @@ import Testing """ } + private func image() throws -> CGImage { + let context = try #require( + CGContext( + data: nil, + width: 2, + height: 2, + bitsPerComponent: 8, + bytesPerRow: 8, + space: CGColorSpaceCreateDeviceRGB(), + bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue + ) + ) + return try #require(context.makeImage()) + } + /// Signatures of every `functionCall` part in a request body, in order. /// A part without a signature contributes `nil`. private func functionCallSignatures(in body: Data) throws -> [String?] { @@ -108,6 +139,161 @@ import Testing #expect(signatures.allSatisfy { $0 == Self.signature }) } + @Test("preserves all text and image data in function responses") + func preservesSegmentedToolOutput() async throws { + StubURLProtocol.reset() + StubURLProtocol.enqueue( + json: + #"{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"image","args":{"request":"read"}}}]},"finishReason":"STOP"}]}"# + ) + StubURLProtocol.enqueue(json: textResponse("OK")) + + let output = try Prompt { + "before" + Attachment(try image()) + "after" + } + let session = LanguageModelSession( + model: makeModel(), + tools: [ImageOutputTool(output: output)] + ) + _ = try await session.respond(to: "read image") + + let bodies = StubURLProtocol.recordedBodies + try #require(bodies.count == 2) + let json = try #require( + try JSONSerialization.jsonObject(with: bodies[1]) as? [String: Any] + ) + let contents = try #require(json["contents"] as? [[String: Any]]) + let outputContent = try #require(contents.last) + let outputParts = try #require(outputContent["parts"] as? [[String: Any]]) + let functionResponse = try #require( + outputParts.first?["functionResponse"] as? [String: Any] + ) + let response = try #require(functionResponse["response"] as? [String: Any]) + #expect(response["result"] as? String == "before\nafter") + + let mediaParts = try #require(functionResponse["parts"] as? [[String: Any]]) + let inlineData = try #require(mediaParts.first?["inlineData"] as? [String: Any]) + #expect(inlineData["mime_type"] as? String == "image/png") + #expect(Data(base64Encoded: try #require(inlineData["data"] as? String)) != nil) + } + + @Test("inlines URL-backed images in function responses") + func inlinesURLBackedToolOutput() async throws { + StubURLProtocol.reset() + StubURLProtocol.enqueue( + json: + #"{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"image","args":{"request":"read"}}}]},"finishReason":"STOP"}]}"# + ) + StubURLProtocol.enqueue(json: textResponse("OK")) + + let imageData = Data([0x89, 0x50, 0x4E, 0x47]) + let url = FileManager.default.temporaryDirectory + .appending(path: UUID().uuidString) + .appendingPathExtension("png") + try imageData.write(to: url) + defer { try? FileManager.default.removeItem(at: url) } + + let output = Prompt { Attachment(imageURL: url) } + let session = LanguageModelSession( + model: makeModel(), + tools: [ImageOutputTool(output: output)] + ) + _ = try await session.respond(to: "read image") + + let body = try #require(StubURLProtocol.recordedBodies.last) + let json = try #require( + try JSONSerialization.jsonObject(with: body) as? [String: Any] + ) + let contents = try #require(json["contents"] as? [[String: Any]]) + let outputParts = try #require(contents.last?["parts"] as? [[String: Any]]) + let functionResponse = try #require( + outputParts.first?["functionResponse"] as? [String: Any] + ) + let mediaParts = try #require(functionResponse["parts"] as? [[String: Any]]) + let inlineData = try #require(mediaParts.first?["inlineData"] as? [String: Any]) + #expect(inlineData["mime_type"] as? String == "image/png") + #expect(Data(base64Encoded: try #require(inlineData["data"] as? String)) == imageData) + #expect(mediaParts.first?["fileData"] == nil) + } + + @Test("sends remote Tool image URLs as ordered text without downloading", arguments: [false, true]) + func describesRemoteURLBackedToolOutput(streaming: Bool) async throws { + StubURLProtocol.reset() + let toolResponse = + #"{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"name":"image","args":{"request":"read"}}}]},"finishReason":"STOP"}]}"# + if streaming { + StubURLProtocol.enqueue(eventStream: [toolResponse]) + StubURLProtocol.enqueue(eventStream: [try self.response(parts: [["text": "OK"]])]) + } else { + StubURLProtocol.enqueue(json: toolResponse) + StubURLProtocol.enqueue(json: textResponse("OK")) + } + + let url = try #require(URL(string: "https://example.invalid/image.png")) + let output = Prompt { + "before" + Attachment(imageURL: url) + "after" + } + let session = LanguageModelSession( + model: makeModel(), + tools: [ImageOutputTool(output: output)] + ) + + let response = + if streaming { + try await session.streamResponse(to: "read image").collect() + } else { + try await session.respond(to: "read image") + } + #expect(response.content == "OK") + let bodies = StubURLProtocol.recordedBodies + try #require(bodies.count == 2) + + func expectURLText(in body: Data) throws { + let responses = try contents(in: body).flatMap { content in + (content["parts"] as? [[String: Any]] ?? []).compactMap { + $0["functionResponse"] as? [String: Any] + } + } + try #require(responses.count == 1) + let functionResponse = try #require(responses.first) + let result = try #require(functionResponse["response"] as? [String: Any]) + #expect(result["result"] as? String == "before\n\(url.absoluteString)\nafter") + #expect(functionResponse["parts"] == nil) + } + try expectURLText(in: bodies[1]) + + let outputs = session.transcript.compactMap { entry -> Transcript.ToolOutput? in + if case .toolOutput(let output) = entry { return output } + return nil + } + let segments = try #require(outputs.first).segments + try #require(segments.count == 3) + guard case .image(let image) = segments[1] else { + Issue.record("The transcript must retain the remote image segment") + return + } + #expect(image.source == .url(url)) + + let data = try JSONEncoder().encode(session.transcript) + let transcript = try JSONDecoder().decode(Transcript.self, from: data) + #expect(transcript == session.transcript) + let restored = LanguageModelSession(model: makeModel(), transcript: transcript) + if streaming { + StubURLProtocol.enqueue(eventStream: [try self.response(parts: [["text": "Replayed"]])]) + #expect(try await restored.streamResponse(to: "continue").collect().content == "Replayed") + } else { + StubURLProtocol.enqueue(json: textResponse("Replayed")) + #expect(try await restored.respond(to: "continue").content == "Replayed") + } + let replayBodies = StubURLProtocol.recordedBodies + try #require(replayBodies.count == 3) + try expectURLText(in: replayBodies[2]) + } + @Test("does not replay the conversation history on later turns") func doesNotReplayHistoryOnLaterTurns() async throws { StubURLProtocol.reset() diff --git a/Tests/AnyLanguageModelTests/HistoryViewTests.swift b/Tests/AnyLanguageModelTests/HistoryViewTests.swift new file mode 100644 index 00000000..57ee4107 --- /dev/null +++ b/Tests/AnyLanguageModelTests/HistoryViewTests.swift @@ -0,0 +1,47 @@ +import Testing + +import AnyLanguageModel + +@Suite("HistoryView collection contract") +struct HistoryViewTests { + private func entry(_ text: String) -> Transcript.Entry { + .prompt(.init(segments: [.text(.init(content: text))])) + } + + @Test func mutableSlicesPreserveIndicesAndValueSemantics() { + let first = entry("first") + let second = entry("second") + let third = entry("third") + var history: Transcript.HistoryView = [first, second, third] + let original = history + let middle = history.index(after: history.startIndex) + var slice: Transcript.HistoryView = history[middle ..< history.endIndex] + + #expect(slice.startIndex == middle) + #expect(slice.endIndex == history.endIndex) + #expect(slice.distance(from: slice.startIndex, to: slice.endIndex) == 2) + #expect(slice.index(before: slice.endIndex) == history.index(history.startIndex, offsetBy: 2)) + #expect(Array(slice.indices) == [middle, middle.advanced(by: 1)]) + #expect(slice.startIndex.distance(to: slice.endIndex) == 2) + slice[slice.startIndex] = first + #expect(Array(history) == [first, second, third]) + history[middle ..< history.endIndex] = slice + #expect(Array(history) == [first, first, third]) + #expect(Array(original) == [first, second, third]) + } + + @Test func rangeReplacementAndCollectionInitializers() { + let first = entry("first") + let second = entry("second") + let third = entry("third") + var history = Transcript.HistoryView([first, second]) + history.replaceSubrange(history.startIndex ..< history.index(after: history.startIndex), with: [third, first]) + #expect(Array(history) == [third, first, second]) + history.append(third) + history.append(contentsOf: [second]) + history.removeFirst() + #expect(Array(history) == [first, second, third, second]) + history.removeAll() + #expect(history.isEmpty) + } +} diff --git a/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift index 18841c72..3e931d8f 100644 --- a/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift @@ -1,7 +1,7 @@ import Foundation import Testing -@testable import AnyLanguageModel +@_spi(Compatibility) @testable import AnyLanguageModel #if MLX import class MLX.MLXArray diff --git a/Tests/AnyLanguageModelTests/NativeSessionPropertyTests.swift b/Tests/AnyLanguageModelTests/NativeSessionPropertyTests.swift new file mode 100644 index 00000000..cc76f281 --- /dev/null +++ b/Tests/AnyLanguageModelTests/NativeSessionPropertyTests.swift @@ -0,0 +1,54 @@ +import Testing + +#if canImport(FoundationModels) && !os(tvOS) && !os(watchOS) + import FoundationModels + @_spi(Compatibility) @testable import AnyLanguageModel + + @Suite("Native Tool session-property binding") + struct NativeSessionPropertyTests { + @Test func nativeWrappersBindPropertiesWithoutTaskLocalInheritance() async throws { + guard #available(macOS 26.0, iOS 26.0, visionOS 26.0, *) else { return } + let first = AnyLanguageModel.LanguageModelSession(model: MockLanguageModel(), tools: [NativePropertyTool()]) + let second = AnyLanguageModel.LanguageModelSession( + model: MockLanguageModel(), + tools: [NativePropertyTool()] + ) + first.properties.revision = 1 + second.properties.revision = 20 + let firstTool = try #require( + first.resolvedRequestContext().toolsForExecution(in: first).toFoundationModels().first + ) + let secondTool = try #require( + second.resolvedRequestContext().toolsForExecution(in: second).toFoundationModels().first + ) + + let firstCall = Task.detached { try await callNativeTool(firstTool) } + let secondCall = Task.detached { try await callNativeTool(secondTool) } + try await firstCall.value + try await secondCall.value + + #expect(first.properties.revision == 2) + #expect(second.properties.revision == 21) + #expect(AnyLanguageModel.SessionPropertyBinding.values == nil) + } + } + + @available(macOS 26.0, iOS 26.0, visionOS 26.0, *) + private func callNativeTool(_ tool: T) async throws { + #expect(AnyLanguageModel.SessionPropertyBinding.values == nil) + let arguments = try T.Arguments(FoundationModels.GeneratedContent(json: "\"read\"")) + _ = try await tool.call(arguments: arguments) + } + + private struct NativePropertyTool: AnyLanguageModel.Tool { + @AnyLanguageModel.LanguageModelSession.SessionProperty(\.revision) private var revision + let name = "property" + let description = "Update a session property" + var parameters: AnyLanguageModel.GenerationSchema { String.generationSchema } + + func call(arguments: String) async throws -> String { + revision += 1 + return "done" + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/ProfileRuntimeTests.swift b/Tests/AnyLanguageModelTests/ProfileRuntimeTests.swift new file mode 100644 index 00000000..9d593d28 --- /dev/null +++ b/Tests/AnyLanguageModelTests/ProfileRuntimeTests.swift @@ -0,0 +1,1422 @@ +import Foundation +import Observation +import Testing + +#if canImport(FoundationModels) && compiler(>=6.4) && !os(tvOS) && !os(watchOS) + import class FoundationModels.SystemLanguageModel +#endif + +@_spi(Compatibility) @testable import AnyLanguageModel + +extension SessionPropertyValues { + @SessionPropertyEntry var profileRevision = 0 +} + +@Suite("Dynamic profiles") +struct DynamicProfileTests { + @Test func callbacksAndToolExecutionRetainRequestScopedHistoryEdits() async throws { + let earlier = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Remove me"))])) + let promptMarker = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Prompt marker"))])) + let callMarker = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Call marker"))])) + let outputMarker = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Output marker"))])) + let responseMarker = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Response marker"))])) + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Evolving history") } + .model(LifecycleFixtureModel()) + .onPrompt { _ in + let properties = try #require(SessionPropertyBinding.values) + properties.history.removeAll { $0.id == earlier.id } + properties.history.append(promptMarker) + } + .onToolCall { _ in + let properties = try #require(SessionPropertyBinding.values) + #expect(!properties.history.contains { $0.id == earlier.id }) + #expect(properties.history.contains(promptMarker)) + properties.history.append(callMarker) + } + .onToolOutput { _, _ in + let properties = try #require(SessionPropertyBinding.values) + #expect(!properties.history.contains { $0.id == earlier.id }) + #expect(properties.history.contains(callMarker)) + properties.history.removeAll { $0.id == callMarker.id } + properties.history.append(outputMarker) + } + .onReasoning { _ in + let properties = try #require(SessionPropertyBinding.values) + #expect(!properties.history.contains { $0.id == earlier.id || $0.id == callMarker.id }) + #expect(properties.history.contains(outputMarker)) + properties.history = properties.history + } + .onResponse { _ in + let properties = try #require(SessionPropertyBinding.values) + #expect(!properties.history.contains { $0.id == earlier.id || $0.id == callMarker.id }) + properties.history.append(responseMarker) + }, + history: [earlier] + ) + let context = try await session.resolvedRequestContext(including: []) + let call = Transcript.ToolCall(id: "call", toolName: "fixture", arguments: .init(properties: [:])) + let calls = Transcript.Entry.toolCalls(.init([call])) + let output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: [.text(.init(content: "Done"))] + ) + let reasoning = Transcript.Entry.reasoning(.init(segments: [.text(.init(content: "Thinking"))])) + let response = Transcript.Entry.response(.init(assetIDs: [], segments: [.text(.init(content: "Answer"))])) + try await session.profileWillExecuteToolCall(call, requestContext: context, currentRoundEntries: [calls]) + let toolHistory = await session.withProfileToolExecution(requestContext: context, currentRoundEntries: [calls]) + { + session.properties.history + } + #expect(!toolHistory.contains { $0.id == earlier.id }) + #expect(toolHistory.contains(callMarker)) + try await session.profileDidProduceToolOutput( + for: call, + output: output, + requestContext: context, + currentRoundEntries: [calls, .toolOutput(output)] + ) + try await session.profileDidProduceEntries( + [reasoning, response], + requestContext: context, + currentRoundEntries: [calls, .toolOutput(output), reasoning, response] + ) + + #expect(Array(session.properties.history) == [promptMarker, outputMarker, responseMarker]) + } + + @Test func erasedConcreteProfilesKeepDistinctLifecycleIdentities() async throws { + let state = LifecycleFixtureState() + let session = LanguageModelSession(profile: ErasedRoutingProfile(state: state)) + _ = try await session.respond(to: "First") + _ = try await session.respond(to: "Same profile") + state.useSecond.withLock { $0 = true } + _ = try await session.respond(to: "Switch") + + #expect(state.events.withLock { $0 } == ["activate:A", "deactivate:A", "activate:B"]) + } + + @Test func lifecycleHistoryEditsAccumulateAcrossPreparationStages() async throws { + let state = LifecycleFixtureState() + let first = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "First old entry"))])) + let second = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Second old entry"))])) + let seen = Locked<[[String]]>([]) + let session = LanguageModelSession( + profile: HistoryEditingLifecycleProfile(state: state, first: first.id, second: second.id, seen: seen), + history: [first, second] + ) + _ = try await session.respond(to: "Activate A") + state.useSecond.withLock { $0 = true } + _ = try await session.respond(to: "Activate B") + + let stages = seen.withLock { $0 } + #expect(stages.count == 2) + #expect(stages.allSatisfy { !$0.contains(first.id) }) + #expect(stages.last?.contains(second.id) == false) + #expect(!session.properties.history.contains { $0.id == first.id || $0.id == second.id }) + } + + @Test(arguments: [false, true]) + func completionEntryCallbacksShareHistoryEdits(streaming: Bool) async throws { + let earlier = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Remove me"))])) + let marker = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Marker"))])) + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Evolving callback history") } + .model(LifecycleFixtureModel()) + .onReasoning { _ in + let properties = try #require(SessionPropertyBinding.values) + properties.history.removeAll { $0.id == earlier.id } + } + .onResponse { _ in + let properties = try #require(SessionPropertyBinding.values) + #expect(!properties.history.contains { $0.id == earlier.id }) + properties.history.append(marker) + }, + history: [earlier] + ) + if streaming { + _ = try await session.streamResponse(to: "Complete").collect() + } else { + _ = try await session.respond(to: "Complete") + } + + #expect(!session.properties.history.contains { $0.id == earlier.id }) + #expect(session.properties.history.last == marker) + } + + @Test(arguments: [false, true]) + func prewarmUsesActiveModelAndFeedbackUsesProducingModel(streaming: Bool) async throws { + let state = ProfileFixtureState() + let session = LanguageModelSession( + profile: FixtureProfile( + state: state, + first: ProfileFixtureModel(id: "a", state: state), + second: ProfileFixtureModel(id: "b", state: state) + ) + ) + session.prewarm(promptPrefix: Prompt("Prefix A")) + if streaming { + _ = try await session.streamResponse(to: "First").collect() + } else { + _ = try await session.respond(to: "First") + } + state.selectSecond() + session.prewarm(promptPrefix: Prompt("Prefix B")) + #expect(session.logFeedbackAttachment(sentiment: .positive) == Data("a".utf8)) + if streaming { + _ = try await session.streamResponse(to: "Second").collect() + } else { + _ = try await session.respond(to: "Second") + } + #expect(session.logFeedbackAttachment(sentiment: .positive) == Data("b".utf8)) + #expect(state.prewarms.withLock { $0 } == ["a:Prefix A", "b:Prefix B"]) + #expect(state.feedback.withLock { $0 } == ["a", "b"]) + } + + @Test func sameSessionReevaluatesModelInstructionsOptionsAndProperties() async throws { + let state = ProfileFixtureState() + let first = ProfileFixtureModel(id: "model-a", state: state) + let second = ProfileFixtureModel(id: "model-b", state: state) + let session = LanguageModelSession( + profile: FixtureProfile(state: state, first: first, second: second) + ) + + _ = try await session.respond(to: "First") + state.selectSecond() + session.properties.profileRevision = 7 + _ = try await session.respond(to: "Second") + + let requests = state.requests.withLock { $0 } + #expect(requests.map(\.modelID) == ["model-a", "model-b"]) + #expect(requests.map(\.instructions) == ["Profile A revision 0", "Profile B revision 7"]) + #expect(requests.map(\.temperature) == [0.2, 0.8]) + #expect(requests.map(\.maximumResponseTokens) == [64, 128]) + #expect(requests.map(\.reasoningLevel) == [.light, .deep]) + #expect(requests.map(\.includeSchemaInPrompt) == [false, true]) + #expect(session.transcript.count == 4) + #expect(!session.transcript.contains { if case .instructions = $0 { true } else { false } }) + } + + @Test func callSiteOptionsOverrideProfileDefaults() async throws { + let state = ProfileFixtureState() + let model = ProfileFixtureModel(id: "model", state: state) + let session = LanguageModelSession( + profile: FixtureProfile(state: state, first: model, second: model) + ) + + _ = try await session.respond( + to: "Override", + options: .init( + temperature: 0.45, + maximumResponseTokens: 17, + toolCallingMode: .required + ) + ) + + let request = try #require(state.requests.withLock { $0.last }) + #expect(request.temperature == 0.45) + #expect(request.maximumResponseTokens == 17) + #expect(request.toolCallingMode == .required) + #expect(request.reasoningLevel == .light) + #expect(request.includeSchemaInPrompt == false) + } + + @Test func lifecycleCallbacksAccumulateInOrderForResponseAndStream() async throws { + let state = LifecycleFixtureState() + let model = LifecycleFixtureModel() + let session = LanguageModelSession(profile: LifecycleFixtureProfile(state: state, model: model)) + + _ = try await session.respond(to: "One") + _ = try await session.streamResponse(to: "Two").collect() + state.useSecond.withLock { $0 = true } + _ = try await session.respond(to: "Three") + + #expect( + state.events.withLock { $0 } == [ + "outer-activate:A", "inner-activate:A", "outer-prompt:A", "inner-prompt:A", + "outer-reasoning:A", "inner-reasoning:A", "outer-response:A", "inner-response:A", + "outer-prompt:A", "inner-prompt:A", "outer-reasoning:A", "inner-reasoning:A", + "outer-response:A", "inner-response:A", "inner-deactivate:A", "outer-deactivate:A", + "outer-activate:B", "inner-activate:B", "outer-prompt:B", "inner-prompt:B", + "outer-reasoning:B", "inner-reasoning:B", "outer-response:B", "inner-response:B", + ] + ) + } + + @MainActor + @Test func appleConcurrencyAndObservationSurfaceCompiles() { + var callbackCount = 0 + _ = LanguageModelSession.Profile { Instructions("Actor callbacks") } + .model(LifecycleFixtureModel()) + .onActivate { callbackCount += 1 } + .onDeactivate { callbackCount += 1 } + .onPrompt { (_: Transcript.Prompt) in callbackCount += 1 } + .onResponse { (_: Transcript.Response) in callbackCount += 1 } + .onReasoning { (_: Transcript.Reasoning) in callbackCount += 1 } + .onToolCall { (_: Transcript.ToolCall) in callbackCount += 1 } + .onToolOutput { (_: Transcript.ToolCall, _: Transcript.ToolOutput) in + callbackCount += 1 + } + + let session = LanguageModelSession(model: LifecycleFixtureModel()) + let _: any Observable = session.properties + } + + @Test func concurrentSnapshotsDispatchEachCallbackOnce() async throws { + let probe = ConcurrentCallbackProbe() + let model = LifecycleFixtureModel() + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Concurrent callback") } + .model(model) + .onResponse { _ in await probe.recordInvocation() } + ) + let context = try await session.resolvedRequestContext(including: [], options: .init()) + let entry = Transcript.Entry.response( + .init(assetIDs: [], segments: [.text(.init(content: "response"))]) + ) + + let first = Task { + try await session.profileDidProduceEntries( + [entry], + requestContext: context, + currentRoundEntries: [entry] + ) + } + await probe.waitUntilFirstInvocation() + let second = Task { + try await session.profileDidProduceEntries( + [entry], + requestContext: context, + currentRoundEntries: [entry] + ) + } + + try await second.value + let invocationCount = await probe.invocationCount + await probe.releaseFirstInvocation() + try await first.value + + #expect(invocationCount == 1) + } + + @Test(arguments: [false, true]) + func completionCallbacksSeeAndPreserveCommittedHistory(streaming: Bool) async throws { + let histories = Locked<[[Transcript.Entry]]>([]) + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Completion history") } + .model(LifecycleFixtureModel()) + .onResponse { _ in + let properties = try #require(SessionPropertyBinding.values) + let history = properties.history + histories.withLock { $0.append(Array(history)) } + properties.history = history + } + ) + + if streaming { + _ = try await session.streamResponse(to: "Complete").collect() + } else { + _ = try await session.respond(to: "Complete") + } + + let history = try #require(histories.withLock { $0.first }) + #expect(history == Array(session.transcript)) + #expect(history.contains { if case .response = $0 { true } else { false } }) + #expect(history.contains { if case .reasoning = $0 { true } else { false } }) + } + + @Test(arguments: [false, true]) + func historyObserversSeeCompletionAndCallbackRollback(streaming: Bool) async throws { + let completionObserved = Locked(false) + let rollbackObserved = Locked(false) + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Observe completion") } + .model(LifecycleFixtureModel()) + .onPrompt { _ in + let properties = try #require(SessionPropertyBinding.values) + withObservationTracking { + _ = properties.history + } onChange: { + completionObserved.withLock { $0 = true } + } + } + .onResponse { _ in + let properties = try #require(SessionPropertyBinding.values) + withObservationTracking { + _ = properties.history + } onChange: { + rollbackObserved.withLock { $0 = true } + } + throw LifecycleFixtureError.rejected + } + ) + + await #expect(throws: LifecycleFixtureError.self) { + if streaming { + _ = try await session.streamResponse(to: "Roll back").collect() + } else { + _ = try await session.respond(to: "Roll back") + } + } + + #expect(completionObserved.withLock { $0 }) + #expect(rollbackObserved.withLock { $0 }) + #expect(session.transcript.count == 1) + } + + @Test func nestedInstructionsReceiveReadOnlyHistory() { + let writability = Locked<[Bool]>([]) + let session = LanguageModelSession(profile: HistoryBindingProfile(writability: writability)) + + _ = session.resolvedRequestContext() + + #expect(writability.withLock { $0 } == [true, false]) + } + + @Test func toolHistoryExcludesStaticInstructions() async { + let session = LanguageModelSession(model: LifecycleFixtureModel(), instructions: "Static") + let context = session.resolvedRequestContext() + let history = await session.withProfileToolExecution(requestContext: context, currentRoundEntries: []) { + session.properties.history + } + + #expect(history.isEmpty) + #expect(session.transcript.count == 1) + } + + @Test(arguments: [false, true]) + func promptCallbackUsesTheActualRequestAfterHistoryMutation(streaming: Bool) async throws { + let prompts = Locked<[String]>([]) + let session = LanguageModelSession(profile: PromptRemovingProfile(prompts: prompts)) + + if streaming { + _ = try await session.streamResponse(to: "Actual request").collect() + } else { + _ = try await session.respond(to: "Actual request") + } + + #expect(prompts.withLock { $0 } == ["Actual request"]) + } + + @Test func concurrentStreamsKeepTheirOwnPromptCallbacks() async throws { + let prompts = Locked<[String]>([]) + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Concurrent prompts") } + .model(LifecycleFixtureModel()) + .onPrompt { (prompt: Transcript.Prompt) in + let text = prompt.segments.compactMap { segment in + if case .text(let text) = segment { return text.content } + return nil + }.joined() + prompts.withLock { $0.append(text) } + } + ) + let first = session.streamResponse(to: "First") + let second = session.streamResponse(to: "Second") + async let firstResponse = first.collect() + async let secondResponse = second.collect() + _ = try await (firstResponse, secondResponse) + + #expect(prompts.withLock { $0.sorted() } == ["First", "Second"]) + } + + @Test(arguments: [false, true]) + func callbackErrorsPropagateWithoutCommittingAResponse(streaming: Bool) async throws { + let model = LifecycleFixtureModel() + let session = LanguageModelSession(profile: ThrowingResponseProfile(model: model)) + + if streaming { + await #expect(throws: LifecycleFixtureError.self) { + _ = try await session.streamResponse(to: "Fail").collect() + } + } else { + await #expect(throws: LifecycleFixtureError.self) { + _ = try await session.respond(to: "Fail") + } + } + + #expect(session.transcript.count == 1) + #expect(session.transcript.allSatisfy { if case .prompt = $0 { true } else { false } }) + } + + @Test(arguments: [false, true]) + func promptFailureUsesTheProducingRequestPolicy(streaming: Bool) async throws { + let state = PromptFailureProfileState() + let session = LanguageModelSession( + profile: PromptFailureProfile(state: state, model: LifecycleFixtureModel()) + ) + + if streaming { + await #expect(throws: LifecycleFixtureError.self) { + _ = try await session.streamResponse(to: "Revert").collect() + } + } else { + await #expect(throws: LifecycleFixtureError.self) { + _ = try await session.respond(to: "Revert") + } + } + #expect(session.transcript.isEmpty) + #expect(session.transcriptErrorHandlingPolicy == nil) + + state.useRevert.withLock { $0 = false } + if streaming { + await #expect(throws: LifecycleFixtureError.self) { + _ = try await session.streamResponse(to: "Keep").collect() + } + } else { + await #expect(throws: LifecycleFixtureError.self) { + _ = try await session.respond(to: "Keep") + } + } + #expect(session.transcript.count == 1) + guard case .prompt = session.transcript[0] else { + Issue.record("Expected the second request prompt to be retained") + return + } + #expect(session.transcriptErrorHandlingPolicy == nil) + } + + @Test func dynamicInstructionsAndToolsShareSessionProperties() async throws { + let model = ProfileFixtureModel(id: "model", state: ProfileFixtureState()) + let session = LanguageModelSession( + model: model, + dynamicInstructions: PropertyInstructions() + ) + session.properties.profileRevision = 4 + + let context = session.resolvedRequestContext() + #expect(context.instructions?.description == "Revision 4") + + let tool = PropertyFixtureTool() + _ = try await session.withProfileToolExecution( + requestContext: context, + currentRoundEntries: [] + ) { + try await tool.makeOutputSegments(from: GeneratedContent(properties: [:])) + } + #expect(session.properties.profileRevision == 5) + } + + @Test func rehydrationKeepsCanonicalHistoryAndUsesTheCurrentProfile() async throws { + let oldState = ProfileFixtureState() + let model = ProfileFixtureModel(id: "model", state: oldState) + let oldSession = LanguageModelSession( + profile: FixtureProfile(state: oldState, first: model, second: model) + ) + _ = try await oldSession.respond(to: "Old") + let saved = oldSession.transcript + + let restoredState = ProfileFixtureState() + restoredState.selectSecond() + let restoredModel = ProfileFixtureModel(id: "restored", state: restoredState) + let restored = LanguageModelSession( + profile: FixtureProfile( + state: restoredState, + first: restoredModel, + second: restoredModel + ), + history: saved + ) + _ = try await restored.respond(to: "New") + + #expect(restoredState.requests.withLock { $0.last?.instructions } == "Profile B revision 0") + #expect(restored.transcript.count == saved.count + 2) + #expect(!restored.transcript.contains { if case .instructions = $0 { true } else { false } }) + } + + @Test func transformedCurrentRoundDoesNotReplaceCanonicalTranscript() async throws { + let state = ProfileFixtureState() + let model = ProfileFixtureModel(id: "model-a", state: state) + let session = LanguageModelSession( + profile: FixtureProfile(state: state, first: model, second: model) + ) + _ = try await session.respond(to: "First") + let durableBefore = session.transcript + let call = Transcript.ToolCall( + id: "call-1", + toolName: "fixture", + arguments: GeneratedContent(properties: [:]) + ) + let output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: [.text(.init(content: "output"))] + ) + let inFlight: [Transcript.Entry] = [ + .toolCalls(.init([call])), + .toolOutput(output), + ] + + let context = try await session.resolvedRequestContext( + including: inFlight, + options: .init() + ) + + #expect(context.transcript.contains { $0.id == output.id }) + #expect(state.transformedHistory.withLock { $0.last?.contains(output.id) } == true) + #expect(session.transcript == durableBefore) + } + + @Test func toolCallbacksUseProducingSnapshotAndSessionPropertiesStayIsolated() async throws { + let state = ProfileFixtureState() + let model = ProfileFixtureModel(id: "model-a", state: state) + let first = LanguageModelSession( + profile: FixtureProfile(state: state, first: model, second: model) + ) + let second = LanguageModelSession( + profile: FixtureProfile(state: state, first: model, second: model) + ) + let call = Transcript.ToolCall( + id: "call", + toolName: "fixture", + arguments: GeneratedContent(properties: [:]) + ) + let output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: [.text(.init(content: "done"))] + ) + let request = try await first.resolvedRequestContext(including: [], options: .init()) + let unrelated = Transcript.Entry.prompt( + .init(id: "unrelated", segments: [.text(.init(content: "Other request"))]) + ) + first.properties.history = [unrelated] + let callEntry = Transcript.Entry.toolCalls(.init([call])) + let outputEntry = Transcript.Entry.toolOutput(output) + + try await first.profileWillExecuteToolCall( + call, + requestContext: request, + currentRoundEntries: [callEntry] + ) + try await first.profileDidProduceToolOutput( + for: call, + output: output, + requestContext: request, + currentRoundEntries: [callEntry, outputEntry] + ) + let toolHistory = await first.withProfileToolExecution( + requestContext: request, + currentRoundEntries: [callEntry, outputEntry] + ) { + first.properties.history.map(\.id) + } + + #expect(first.properties.profileRevision == 1) + #expect(second.properties.profileRevision == 0) + #expect(state.callbackEvents.withLock { $0 } == ["call:call", "output:call"]) + #expect( + state.callbackHistories.withLock { $0 } + == [[callEntry.id], [callEntry.id, outputEntry.id]] + ) + #expect(toolHistory == [callEntry.id, outputEntry.id]) + } + + @Test func callbackHistoryDoesNotDuplicateCanonicalEntries() async throws { + let state = ProfileFixtureState() + let model = ProfileFixtureModel(id: "model-a", state: state) + let call = Transcript.ToolCall( + id: "call", + toolName: "fixture", + arguments: GeneratedContent(properties: [:]) + ) + let output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: [.text(.init(content: "done"))] + ) + let callEntry = Transcript.Entry.toolCalls(.init([call])) + let outputEntry = Transcript.Entry.toolOutput(output) + let session = LanguageModelSession( + profile: FixtureProfile(state: state, first: model, second: model), + history: [callEntry] + ) + let request = try await session.resolvedRequestContext(including: [], options: .init()) + + try await session.profileDidProduceToolOutput( + for: call, + output: output, + requestContext: request, + currentRoundEntries: [callEntry, outputEntry] + ) + let toolHistory = await session.withProfileToolExecution( + requestContext: request, + currentRoundEntries: [callEntry, outputEntry] + ) { + session.properties.history.map(\.id) + } + + #expect(state.callbackHistories.withLock { $0 } == [[callEntry.id, outputEntry.id]]) + #expect(toolHistory == [callEntry.id, outputEntry.id]) + } + + @Test func continuationHistoryMutationDoesNotCommitPendingEntries() async throws { + let state = ProfileFixtureState() + let model = ProfileFixtureModel(id: "model", state: state) + let session = LanguageModelSession(profile: HistoryWritingProfile(model: model)) + _ = try await session.respond(to: "Durable") + let durableBefore = session.transcript + let call = Transcript.ToolCall( + id: "pending-call", + toolName: "fixture", + arguments: GeneratedContent(properties: [:]) + ) + let output = Transcript.ToolOutput( + id: call.id, + toolName: call.toolName, + segments: [.text(.init(content: "pending"))] + ) + let pending: [Transcript.Entry] = [ + .toolCalls(.init([call])), + .toolOutput(output), + ] + + let context = try await session.resolvedRequestContext( + including: pending, + options: .init() + ) + + #expect(context.transcript.contains { $0.id == output.id }) + #expect(session.transcript == durableBefore) + } + + enum UsageEntryPoint: CaseIterable { + case prompt, schema, images + } + + private enum UsageFailure: Error { case failed } + + @Test(arguments: UsageEntryPoint.allCases, [false, true]) + func profileCountsEachRequestBeforeResponseCompletes(entryPoint: UsageEntryPoint, fails: Bool) async throws { + let usage = LanguageModelSession.Usage( + input: .init(totalTokenCount: 13, cachedTokenCount: 0), + output: .init(totalTokenCount: 5, reasoningTokenCount: 0) + ) + let activeSession = Locked(nil) + let observed = Locked(nil) + var model = MockLanguageModel { _, _ in + var requests = RequestsUsage() + requests.add(usage) + observed.withLock { $0 = activeSession.withLock { $0?.usage } } + if fails { throw UsageFailure.failed } + return #""Done""# + } + model.usage = usage + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Request accounting") }.model(model) + ) + activeSession.withLock { $0 = session } + func respond() async throws { + switch entryPoint { + case .prompt: + _ = try await session.respond(to: Prompt("Hi")) + case .schema: + _ = try await session.respond(to: Prompt("Hi"), schema: String.generationSchema) + case .images: + _ = try await session.respond(to: "Hi", images: []) + } + } + if fails { + await #expect(throws: UsageFailure.self) { try await respond() } + } else { + try await respond() + } + #expect(observed.withLock { $0 } == usage) + #expect(session.usage == usage) + #expect(!session.isResponding) + } + + @Test(arguments: UsageEntryPoint.allCases, [false, true]) + func aStreamDoesNotReportUsageToItsSurroundingResponse(entryPoint: UsageEntryPoint, profile: Bool) async throws { + let usage = LanguageModelSession.Usage( + input: .init(totalTokenCount: 7, cachedTokenCount: 0), + output: .init(totalTokenCount: 3, reasoningTokenCount: 0) + ) + var model = MockLanguageModel { _, _ in + var requests = RequestsUsage() + requests.add(usage) + return #""Done""# + } + model.usage = usage + let session = + profile + ? LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Stream accounting") }.model(model) + ) : LanguageModelSession(model: model) + let channel = RequestUsageChannel { _ in } + try await RequestUsageChannel.$current.withValue(channel) { + switch entryPoint { + case .prompt: + _ = try await session.streamResponse(to: Prompt("Hi")).collect() + case .schema: + _ = try await session.streamResponse(to: Prompt("Hi"), schema: String.generationSchema).collect() + case .images: + _ = try await session.streamResponse(to: "Hi", images: []).collect() + } + } + #expect(channel.reportedUsage == .zero) + #expect(session.usage == usage) + #expect(!session.isResponding) + } + + @Test(arguments: [false, true], [false, true]) + func modelHandoffPreservesFinalAndEarlierRequestUsage(streaming: Bool, finalReportsUsage: Bool) async throws { + let prefix = LanguageModelSession.Usage( + input: .init(totalTokenCount: 11, cachedTokenCount: 0), + output: .init(totalTokenCount: 3, reasoningTokenCount: 0) + ) + let earlier = LanguageModelSession.Usage( + input: .init(totalTokenCount: 7, cachedTokenCount: 0), + output: .init(totalTokenCount: 2, reasoningTokenCount: 0) + ) + let final = + finalReportsUsage + ? LanguageModelSession.Usage( + input: .init(totalTokenCount: 5, cachedTokenCount: 0), + output: .init(totalTokenCount: 1, reasoningTokenCount: 0) + ) : .zero + let model = HandoffUsageModel(requests: .init(earlier: earlier, final: final)) + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Usage handoff") }.model(model) + ) + let context = session.resolvedRequestContext() + let call = Transcript.ToolCall(id: "earlier-call", toolName: "fixture", arguments: .init(properties: [:])) + let entries: [Transcript.Entry] = [ + .toolCalls(.init([call])), + .toolOutput(.init(id: call.id, toolName: call.toolName, segments: [.text(.init(content: "Earlier"))])), + ] + let response: LanguageModelSession.Response + if streaming { + let stream: LanguageModelSession.ResponseStream = try #require( + session.profileContinuationStream( + ifSelectedModelDiffersFrom: UUID(), + requestContext: context, + prompt: Prompt("Continue"), + generating: String.self, + schema: String.generationSchema, + includeSchemaInPrompt: true, + prefixEntries: entries, + prefixUsage: prefix, + prefixText: "" + ) + ) + response = try await stream.collect() + } else { + response = try #require( + try await session.profileContinuationResponse( + ifSelectedModelDiffersFrom: UUID(), + requestContext: context, + prompt: Prompt("Continue"), + generating: String.self, + schema: String.generationSchema, + includeSchemaInPrompt: true, + prefixEntries: entries, + prefixUsage: prefix, + prefixText: "" + ) + ) + } + var total = prefix + total.add(earlier) + let expectedEarlier = total + total.add(final) + #expect(response.usage == final) + #expect(response.requests.earlier == expectedEarlier) + #expect(response.requests.total == total) + } + + @Test func modelHandoffIsSkippedForTheInitialRequest() async throws { + let session = LanguageModelSession( + profile: LifecycleFixtureProfile( + state: LifecycleFixtureState(), + model: LifecycleFixtureModel() + ) + ) + let context = try await session.resolvedRequestContext(including: [], options: .init()) + let response: LanguageModelSession.Response? = try await session.profileContinuationResponse( + ifSelectedModelDiffersFrom: UUID(), + requestContext: context, + prompt: Prompt("Initial"), + generating: String.self, + schema: String.generationSchema, + includeSchemaInPrompt: true, + prefixEntries: [], + prefixUsage: .zero, + prefixText: "" + ) + let stream: LanguageModelSession.ResponseStream? = session.profileContinuationStream( + ifSelectedModelDiffersFrom: UUID(), + requestContext: context, + prompt: Prompt("Initial"), + generating: String.self, + schema: String.generationSchema, + includeSchemaInPrompt: true, + prefixEntries: [], + prefixUsage: .zero, + prefixText: "" + ) + + #expect(response == nil) + #expect(stream == nil) + } + + @Test func profileErrorPolicyDoesNotLeakIntoTheNextProfile() async throws { + let state = LifecycleFixtureState() + let session = LanguageModelSession( + profile: ErrorPolicyProfile(state: state, model: FailingProfileModel()) + ) + session.transcriptErrorHandlingPolicy = .preserveTranscript + + await #expect(throws: LifecycleFixtureError.self) { + _ = try await session.respond(to: "Revert") + } + #expect(session.transcript.isEmpty) + + state.useSecond.withLock { $0 = true } + await #expect(throws: LifecycleFixtureError.self) { + _ = try await session.streamResponse(to: "Preserve").collect() + } + #expect(session.transcript.count == 1) + #expect(session.transcript.allSatisfy { if case .prompt = $0 { true } else { false } }) + #expect(session.transcriptErrorHandlingPolicy == .preserveTranscript) + } +} + +private struct FixtureProfile: LanguageModelSession.DynamicProfile, @unchecked Sendable { + let state: ProfileFixtureState + let first: ProfileFixtureModel + let second: ProfileFixtureModel + @SessionProperty(\.profileRevision) private var revision + @SessionProperty(\.history) private var history + + var body: some LanguageModelSession.DynamicProfile { + if state.usesSecond { + LanguageModelSession.Profile { + Instructions("Profile B revision \(revision)") + FixtureProfileTool() + } + .model(second) + .temperature(0.8) + .maximumResponseTokens(128) + .reasoningLevel(.deep) + .includeSchemaInPrompt(true) + .historyTransform(recordingTransform) + .onToolCall(perform: recordCall) + .onToolOutput(perform: recordOutput) + } else { + LanguageModelSession.Profile { + Instructions("Profile A revision \(revision)") + FixtureProfileTool() + } + .model(first) + .temperature(0.2) + .maximumResponseTokens(64) + .reasoningLevel(.light) + .includeSchemaInPrompt(false) + .historyTransform(recordingTransform) + .onToolCall(perform: recordCall) + .onToolOutput(perform: recordOutput) + } + } + + private func recordingTransform(_ entries: [Transcript.Entry]) -> [Transcript.Entry] { + state.transformedHistory.withLock { $0.append(entries.map(\.id)) } + return entries + } + + private func recordCall(_ call: Transcript.ToolCall) async throws { + state.callbackEvents.withLock { $0.append("call:\(call.id)") } + state.callbackHistories.withLock { $0.append(history.map(\.id)) } + } + + private func recordOutput( + _ call: Transcript.ToolCall, + _ output: Transcript.ToolOutput + ) async throws { + revision += 1 + state.callbackEvents.withLock { $0.append("output:\(output.id)") } + state.callbackHistories.withLock { $0.append(history.map(\.id)) } + } +} + +private struct FixtureProfileTool: Tool { + let name = "fixture" + let description = "Fixture" + + @Generable + struct Arguments {} + + func call(arguments: Arguments) async throws -> String { "done" } +} + +private struct HandoffUsageModel: LanguageModel { + typealias UnavailableReason = Never + let requests: RequestsUsage + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response { + let raw = GeneratedContent("Final") + return .init(content: try Content(raw), rawContent: raw, transcriptEntries: [], requests: requests) + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream { + .init( + stream: AsyncThrowingStream { continuation in + do { + let raw = GeneratedContent("Final") + continuation.yield( + .init(content: try Content(raw).asPartiallyGenerated(), rawContent: raw, requests: requests) + ) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + ) + } +} + +private struct ProfileFixtureModel: LanguageModel { + typealias UnavailableReason = Never + let id: String + let state: ProfileFixtureState + + func prewarm(for session: LanguageModelSession, promptPrefix: Prompt?) { + state.prewarms.withLock { $0.append("\(id):\(promptPrefix?.description ?? "")") } + } + + func logFeedbackAttachment( + within session: LanguageModelSession, + sentiment: LanguageModelFeedback.Sentiment?, + issues: [LanguageModelFeedback.Issue], + desiredOutput: Transcript.Entry? + ) -> Data { + state.feedback.withLock { $0.append(id) } + return Data(id.utf8) + } + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + let context = try await session.resolvedRequestContext(including: [], options: options) + state.requests.withLock { + $0.append( + .init( + modelID: id, + instructions: context.instructions?.description, + temperature: context.options.temperature, + maximumResponseTokens: context.options.maximumResponseTokens, + toolCallingMode: context.options.toolCallingMode, + reasoningLevel: context.contextOptions.reasoningLevel, + includeSchemaInPrompt: includeSchemaInPrompt + ) + ) + } + let raw = GeneratedContent("response") + return .init(content: try Content(raw), rawContent: raw, transcriptEntries: []) + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + let stream = AsyncThrowingStream.Snapshot, any Error> { + continuation in + let task = Task { + do { + let response: LanguageModelSession.Response = try await respond( + within: session, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) + continuation.yield( + .init( + content: response.content.asPartiallyGenerated(), + rawContent: response.rawContent + ) + ) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in task.cancel() } + } + return .init(stream: stream) + } +} + +private final class ProfileFixtureState: @unchecked Sendable { + struct Request: Equatable { + let modelID: String + let instructions: String? + let temperature: Double? + let maximumResponseTokens: Int? + let toolCallingMode: GenerationOptions.ToolCallingMode? + let reasoningLevel: ContextOptions.ReasoningLevel? + let includeSchemaInPrompt: Bool + } + + private let selection = Locked(false) + let requests = Locked<[Request]>([]) + let transformedHistory = Locked<[[String]]>([]) + let callbackEvents = Locked<[String]>([]) + let callbackHistories = Locked<[[String]]>([]) + let prewarms = Locked<[String]>([]) + let feedback = Locked<[String]>([]) + + var usesSecond: Bool { selection.withLock { $0 } } + func selectSecond() { selection.withLock { $0 = true } } +} + +private struct PropertyInstructions: DynamicInstructions, @unchecked Sendable { + @SessionProperty(\.profileRevision) private var revision + + var body: some DynamicInstructions { + Instructions("Revision \(revision)") + } +} + +private struct PropertyFixtureTool: Tool, @unchecked Sendable { + let name = "property" + let description = "Updates a session property" + @SessionProperty(\.profileRevision) private var revision + + @Generable struct Arguments {} + + func call(arguments: Arguments) async throws -> String { + revision += 1 + return "updated" + } +} + +private enum LifecycleFixtureError: Error { + case rejected +} + +private struct HistoryBindingProfile: LanguageModelSession.DynamicProfile, @unchecked Sendable { + let writability: Locked<[Bool]> + + var body: some LanguageModelSession.DynamicProfile { + writability.withLock { $0.append(SessionPropertyBinding.history?.isWritable ?? false) } + return LanguageModelSession.Profile { HistoryBindingInstructions(writability: writability) } + .model(LifecycleFixtureModel()) + } +} + +private struct HistoryBindingInstructions: DynamicInstructions, @unchecked Sendable { + let writability: Locked<[Bool]> + + var body: some DynamicInstructions { + writability.withLock { $0.append(SessionPropertyBinding.history?.isWritable ?? true) } + return Instructions("Read-only instructions") + } +} + +private struct PromptRemovingProfile: LanguageModelSession.DynamicProfile, @unchecked Sendable { + let prompts: Locked<[String]> + @SessionProperty(\.history) private var history + + var body: some LanguageModelSession.DynamicProfile { + history = [] + return LanguageModelSession.Profile { Instructions("Removed prompt") } + .model(LifecycleFixtureModel()) + .onPrompt { (prompt: Transcript.Prompt) in + let text = prompt.segments.compactMap { segment in + if case .text(let text) = segment { return text.content } + return nil + }.joined() + prompts.withLock { $0.append(text) } + } + } +} + +private actor ConcurrentCallbackProbe { + private(set) var invocationCount = 0 + private var firstInvocationWaiter: CheckedContinuation? + private var firstInvocationRelease: CheckedContinuation? + + func recordInvocation() async { + invocationCount += 1 + guard invocationCount == 1 else { return } + firstInvocationWaiter?.resume() + firstInvocationWaiter = nil + await withCheckedContinuation { firstInvocationRelease = $0 } + } + + func waitUntilFirstInvocation() async { + guard invocationCount == 0 else { return } + await withCheckedContinuation { firstInvocationWaiter = $0 } + } + + func releaseFirstInvocation() { + firstInvocationRelease?.resume() + firstInvocationRelease = nil + } +} + +private final class LifecycleFixtureState: @unchecked Sendable { + let useSecond = Locked(false) + let events = Locked<[String]>([]) + + func record(_ value: String) { + events.withLock { $0.append(value) } + } +} + +private enum FirstErasedProfile {} +private enum SecondErasedProfile {} + +private struct ErasedLeafProfile: LanguageModelSession.DynamicProfile { + var body: some LanguageModelSession.DynamicProfile { + LanguageModelSession.Profile { Instructions("Same leaf type") } + } +} + +private struct ErasedRoutingProfile: LanguageModelSession.DynamicProfile { + let state: LifecycleFixtureState + + var body: some LanguageModelSession.DynamicProfile { + let second = state.useSecond.withLock { $0 } + let profile = + second + ? LanguageModelSession.AnyDynamicProfile(ErasedLeafProfile()) + : LanguageModelSession.AnyDynamicProfile(ErasedLeafProfile()) + let label = second ? "B" : "A" + return profile.model(LifecycleFixtureModel()) + .onActivate { state.record("activate:\(label)") } + .onDeactivate { state.record("deactivate:\(label)") } + } +} + +private struct HistoryEditingLifecycleProfile: LanguageModelSession.DynamicProfile { + let state: LifecycleFixtureState + let first: String + let second: String + let seen: Locked<[[String]]> + + var body: some LanguageModelSession.DynamicProfile { + if state.useSecond.withLock({ $0 }) { + LanguageModelSession.Profile { Instructions("B") } + .model(LifecycleFixtureModel()) + .onActivate { + guard let properties = SessionPropertyBinding.values else { return } + seen.withLock { $0.append(properties.history.map(\.id)) } + properties.history.removeAll { $0.id == second } + } + .onPrompt { _ in + let properties = try #require(SessionPropertyBinding.values) + seen.withLock { $0.append(properties.history.map(\.id)) } + } + } else { + LanguageModelSession.Profile { Instructions("A") } + .model(LifecycleFixtureModel()) + .onDeactivate { + guard let properties = SessionPropertyBinding.values else { return } + properties.history.removeAll { $0.id == first } + } + } + } +} + +private struct CountingNativeInstructions: DynamicInstructions { + @SessionProperty(\.profileRevision) private var revision + + var body: some DynamicInstructions { + revision += 1 + return Instructions("Revision \(revision)") + } +} + +private struct LifecycleFixtureProfile: LanguageModelSession.DynamicProfile, @unchecked Sendable { + let state: LifecycleFixtureState + let model: LifecycleFixtureModel + + var body: some LanguageModelSession.DynamicProfile { + if state.useSecond.withLock({ $0 }) { + profile(label: "B") + } else { + profile(label: "A") + } + } + + private func profile(label: String) -> some LanguageModelSession.DynamicProfile { + LanguageModelSession.Profile { Instructions("Lifecycle \(label)") } + .model(model) + .onActivate { state.record("inner-activate:\(label)") } + .onDeactivate { state.record("inner-deactivate:\(label)") } + .onPrompt { state.record("inner-prompt:\(label)") } + .onReasoning { state.record("inner-reasoning:\(label)") } + .onResponse { state.record("inner-response:\(label)") } + .onActivate { state.record("outer-activate:\(label)") } + .onDeactivate { state.record("outer-deactivate:\(label)") } + .onPrompt { state.record("outer-prompt:\(label)") } + .onReasoning { state.record("outer-reasoning:\(label)") } + .onResponse { state.record("outer-response:\(label)") } + } +} + +private struct ThrowingResponseProfile: LanguageModelSession.DynamicProfile { + let model: LifecycleFixtureModel + + var body: some LanguageModelSession.DynamicProfile { + LanguageModelSession.Profile { Instructions("Throw after response") } + .model(model) + .onResponse { throw LifecycleFixtureError.rejected } + } +} + +private struct HistoryWritingProfile: LanguageModelSession.DynamicProfile { + let model: ProfileFixtureModel + @SessionProperty(\.history) private var history + + var body: some LanguageModelSession.DynamicProfile { + history = history.filter { entry in + if case .toolCalls = entry { return false } + if case .toolOutput = entry { return false } + return true + } + return LanguageModelSession.Profile { Instructions("History") } + .model(model) + } +} + +private struct ErrorPolicyProfile: LanguageModelSession.DynamicProfile, @unchecked Sendable { + let state: LifecycleFixtureState + let model: FailingProfileModel + + var body: some LanguageModelSession.DynamicProfile { + if state.useSecond.withLock({ $0 }) { + LanguageModelSession.Profile { Instructions("Caller policy") } + .model(model) + } else { + LanguageModelSession.Profile { Instructions("Profile policy") } + .model(model) + .transcriptErrorHandlingPolicy(.revertTranscript) + } + } +} + +private final class PromptFailureProfileState: @unchecked Sendable { + let useRevert = Locked(true) +} + +private struct PromptFailureProfile: LanguageModelSession.DynamicProfile, @unchecked Sendable { + let state: PromptFailureProfileState + let model: LifecycleFixtureModel + + var body: some LanguageModelSession.DynamicProfile { + if state.useRevert.withLock({ $0 }) { + LanguageModelSession.Profile { Instructions("Revert prompt failures") } + .model(model) + .transcriptErrorHandlingPolicy(.revertTranscript) + .onPrompt { throw LifecycleFixtureError.rejected } + } else { + LanguageModelSession.Profile { Instructions("Keep prompt failures") } + .model(model) + .onPrompt { throw LifecycleFixtureError.rejected } + } + } +} + +private struct FailingProfileModel: LanguageModel { + typealias UnavailableReason = Never + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + throw LifecycleFixtureError.rejected + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + .init(stream: AsyncThrowingStream { $0.finish(throwing: LifecycleFixtureError.rejected) }) + } +} + +private struct LifecycleFixtureModel: LanguageModel { + typealias UnavailableReason = Never + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + let raw = GeneratedContent("response") + let reasoning = Transcript.Entry.reasoning( + .init(segments: [.text(.init(content: "reasoning"))]) + ) + return .init( + content: try Content(raw), + rawContent: raw, + transcriptEntries: [reasoning] + ) + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + let stream = AsyncThrowingStream.Snapshot, any Error> { + continuation in + let task = Task { + do { + let response: LanguageModelSession.Response = try await respond( + within: session, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) + continuation.yield( + .init( + content: response.content.asPartiallyGenerated(), + rawContent: response.rawContent, + transcriptEntries: response.transcriptEntries + ) + ) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in task.cancel() } + } + return .init(stream: stream) + } +} diff --git a/Tests/AnyLanguageModelTests/ProfileToolRoundLimitTests.swift b/Tests/AnyLanguageModelTests/ProfileToolRoundLimitTests.swift new file mode 100644 index 00000000..1aeec4a9 --- /dev/null +++ b/Tests/AnyLanguageModelTests/ProfileToolRoundLimitTests.swift @@ -0,0 +1,65 @@ +import Foundation +import Testing + +@_spi(Compatibility) @testable import AnyLanguageModel + +@Test func ordinarySessionToolRoundLimitUsesRequestOptions() throws { + let session = LanguageModelSession(model: MockLanguageModel()) + var options = GenerationOptions() + options.maximumToolRounds = 2 + let budget = session.toolRoundLimit(for: "fixture", options: options) + try budget.record(completed: []) + try budget.record(completed: []) + #expect(throws: LanguageModelSession.ToolRoundLimitExceeded.self) { + try budget.record(completed: []) + } +} + +@Test func profileToolRoundLimitSharesOriginalCallSiteBudget() async throws { + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("Budget") } + .model(ToolBudgetProbeModel()) + ) + var options = GenerationOptions() + options.maximumToolRounds = 2 + do { + _ = try await session.respond(to: "Use the configured budget", options: options) + Issue.record("Expected the third round to exceed the shared budget") + } catch let error as LanguageModelSession.ToolRoundLimitExceeded { + #expect(error.limit == 2) + #expect(error.provider == "first") + #expect(error.transcriptEntries.count == 1) + } +} + +private struct ToolBudgetProbeModel: LanguageModel { + typealias UnavailableReason = Never + var availability: Availability { .available } + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + let first = session.toolRoundLimit(for: "first", options: .init()) + let second = session.toolRoundLimit(for: "second", options: .init()) + #expect(first === second) + try first.record(completed: []) + try second.record(completed: []) + let completed: [Transcript.Entry] = [.response(.init(assetIDs: [], segments: []))] + try second.record(completed: completed) + throw CancellationError() + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + fatalError("ToolBudgetProbeModel only supports nonstreaming tests") + } +} diff --git a/Tests/AnyLanguageModelTests/PromptAttachmentTests.swift b/Tests/AnyLanguageModelTests/PromptAttachmentTests.swift index 34d6cdf4..348c93e3 100644 --- a/Tests/AnyLanguageModelTests/PromptAttachmentTests.swift +++ b/Tests/AnyLanguageModelTests/PromptAttachmentTests.swift @@ -1,7 +1,7 @@ import Foundation import Testing -@testable import AnyLanguageModel +@_spi(Compatibility) @testable import AnyLanguageModel #if canImport(CoreGraphics) && canImport(ImageIO) import CoreGraphics @@ -122,6 +122,64 @@ struct PromptAttachmentTests { } } + @Test func ordinaryToolReturnsOrderedImageSegments() async throws { + let url = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString + ".png") + try testImageData.write(to: url) + defer { try? FileManager.default.removeItem(at: url) } + let output = AnyLanguageModel.Prompt { + "before" + AnyLanguageModel.Attachment(imageURL: url) + "after" + } + let segments = try await ImageTool(output: output).makeOutputSegments( + from: AnyLanguageModel.GeneratedContent("read") + ) + #expect(segments.count == 3) + guard case .text(let before) = segments[0], + case .image(let image) = segments[1], + case .text(let after) = segments[2] + else { + Issue.record("Tool output lost its ordered image") + return + } + #expect(before.content == "before") + #expect(after.content == "after") + #expect(image.source == .data(testImageData, mimeType: "image/png")) + try FileManager.default.removeItem(at: url) + + let transcript = AnyLanguageModel.Transcript(entries: [ + .toolOutput(.init(id: "image-call", toolName: "image", segments: segments)) + ]) + let encoded = try JSONEncoder().encode(transcript) + #expect(try JSONDecoder().decode(AnyLanguageModel.Transcript.self, from: encoded) == transcript) + } + + @Test(arguments: ["pdf", "txt"]) + func ordinaryToolRejectsNonImageFiles(fileExtension: String) async throws { + let url = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + .appendingPathExtension(fileExtension) + try Data("Not an image".utf8).write(to: url) + defer { try? FileManager.default.removeItem(at: url) } + await #expect(throws: (any Error).self) { + _ = try await ImageTool(output: AnyLanguageModel.Prompt { AnyLanguageModel.Attachment(imageURL: url) }) + .makeOutputSegments(from: AnyLanguageModel.GeneratedContent("read")) + } + } + + @Test func ordinaryToolRetainsRemoteImageURLs() async throws { + let url = try #require(URL(string: "https://example.com/image.png")) + let segments = try await ImageTool( + output: AnyLanguageModel.Prompt { + AnyLanguageModel.Attachment(imageURL: url) + } + ).makeOutputSegments(from: AnyLanguageModel.GeneratedContent("read")) + guard case .image(let image) = try #require(segments.first) else { + Issue.record("Missing image") + return + } + #expect(image.source == .url(url)) + } + #if canImport(CoreGraphics) && canImport(ImageIO) private func image() throws -> CGImage { let context = try #require( @@ -267,11 +325,41 @@ struct PromptAttachmentTests { #expect(second.content == "prompts:2 images:1 right:1") } + @available(macOS 27, iOS 27, visionOS 27, *) + @Test func nativeToolWrapperRetainsOrdinaryImageOutput() async throws { + let output = try AnyLanguageModel.Prompt { + "tool note" + AnyLanguageModel.Attachment(try image(), orientation: .right) + } + let tools: [any AnyLanguageModel.Tool] = [ImageTool(output: output)] + let session = FoundationModels.LanguageModelSession( + model: BridgeModel(callsTools: true), + tools: tools.toFoundationModels() + ) + _ = try await session.respond(to: "read image") + let recorded = try #require( + session.transcript.compactMap { entry -> FoundationModels.Transcript.ToolOutput? in + guard case .toolOutput(let output) = entry else { return nil } + return output + }.first + ) + #expect(recorded.id == "image-call") + #expect(recorded.segments.count == 2) + guard case .attachment(let segment) = recorded.segments[1], + case .image(let image) = segment.content + else { + Issue.record("Native Tool wrapper flattened image output") + return + } + #expect(image.orientation == .right) + } + @available(macOS 27, iOS 27, visionOS 27, *) private struct BridgeModel: FoundationModels.LanguageModel { typealias Executor = BridgeExecutor + var callsTools = false var capabilities: FoundationModels.LanguageModelCapabilities { - .init([.vision]) + .init([.vision, .toolCalling]) } var executorConfiguration: String { "bridge-fixture" @@ -287,6 +375,27 @@ struct PromptAttachmentTests { model: BridgeModel, streamingInto channel: FoundationModels.LanguageModelExecutorGenerationChannel ) async throws { + if model.callsTools, + !request.transcript.contains(where: { + if case .toolOutput = $0 { return true } + return false + }) + { + await channel.send( + .toolCalls( + action: .toolCall( + id: "image-call", + name: "image", + action: .appendArguments(#""read""#, tokenCount: 1) + ) + ) + ) + return + } + if model.callsTools { + await channel.send(.response(action: .appendText("OK", tokenCount: 1))) + return + } let prompts = request.transcript.compactMap { entry -> FoundationModels.Transcript.Prompt? in guard case .prompt(let prompt) = entry else { return nil } return prompt @@ -313,3 +422,17 @@ struct PromptAttachmentTests { } #endif } + +private struct ImageTool: AnyLanguageModel.Tool { + let name = "image" + let description = "Return an image" + let output: AnyLanguageModel.Prompt + + var parameters: AnyLanguageModel.GenerationSchema { + String.generationSchema + } + + func call(arguments: String) async throws -> AnyLanguageModel.Prompt { + output + } +} diff --git a/Tests/AnyLanguageModelTests/ProviderUsageTests.swift b/Tests/AnyLanguageModelTests/ProviderUsageTests.swift index 92fb58f1..98736157 100644 --- a/Tests/AnyLanguageModelTests/ProviderUsageTests.swift +++ b/Tests/AnyLanguageModelTests/ProviderUsageTests.swift @@ -9,14 +9,11 @@ import Testing enum Provider: CaseIterable, Equatable, Sendable { case chat, responses, openResponses, anthropic, gemini, ollama - func makeSession(tools: [any Tool] = [], responseStartTimeout: Duration = .seconds(60)) - -> LanguageModelSession - { + func makeModel(responseStartTimeout: Duration = .seconds(60)) -> any LanguageModel { let http = UsageURLProtocol.makeSession() - let model: any LanguageModel switch self { case .chat, .responses: - model = OpenAILanguageModel( + return OpenAILanguageModel( apiKey: "test", model: "test", apiVariant: self == .chat ? .chatCompletions : .responses, @@ -24,7 +21,7 @@ import Testing transport: http ) case .openResponses: - model = OpenResponsesLanguageModel( + return OpenResponsesLanguageModel( baseURL: URL(string: "https://example.com/v1")!, apiKey: "test", model: "test", @@ -32,27 +29,38 @@ import Testing transport: http ) case .anthropic: - model = AnthropicLanguageModel( + return AnthropicLanguageModel( apiKey: "test", model: "test", responseStartTimeout: responseStartTimeout, transport: http ) case .gemini: - model = GeminiLanguageModel( + return GeminiLanguageModel( apiKey: "test", model: "test", responseStartTimeout: responseStartTimeout, transport: http ) case .ollama: - model = OllamaLanguageModel( + return OllamaLanguageModel( model: "test", responseStartTimeout: responseStartTimeout, transport: http ) } - return LanguageModelSession(model: model, tools: tools) + } + + func makeSession( + tools: [any Tool] = [], + transcript: Transcript = .init(), + responseStartTimeout: Duration = .seconds(60) + ) -> LanguageModelSession { + LanguageModelSession( + model: makeModel(responseStartTimeout: responseStartTimeout), + tools: tools, + transcript: transcript + ) } var counts: [String: Any] { @@ -521,6 +529,335 @@ import Testing #expect(UsageURLProtocol.recordedBodies.count == 2) } + @Test(arguments: [Provider.anthropic, .gemini], [false, true]) + func localToolImagesSurviveDeletionAndTranscriptRestoration(_ provider: Provider, _ streaming: Bool) + async throws + { + UsageURLProtocol.reset() + UsageURLProtocol.enqueue( + json: try streaming ? provider.toolStream() : Self.json(provider.response(text: "", tool: true)) + ) + if streaming || provider == .gemini { + UsageURLProtocol.enqueue(json: try streaming ? provider.stream() : Self.json(provider.response())) + } + UsageURLProtocol.enqueue(json: try Self.json(provider.response())) + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString) + .appendingPathExtension("png") + try testImageData.write(to: url) + defer { try? FileManager.default.removeItem(at: url) } + let session = provider.makeSession(tools: [LocalImageWeatherTool(url: url)]) + session.toolExecutionDelegate = DeleteImageAfterToolOutput(url: url) + + if streaming { + _ = try await session.streamResponse(to: "Weather?").collect() + } else { + _ = try await session.respond(to: "Weather?") + } + #expect(!FileManager.default.fileExists(atPath: url.path)) + let archived = try JSONEncoder().encode(session.transcript) + let restored = provider.makeSession(transcript: try JSONDecoder().decode(Transcript.self, from: archived)) + _ = try await restored.respond(to: "Review the image") + + #expect(UsageURLProtocol.recordedBodies.count == (streaming || provider == .gemini ? 3 : 2)) + for body in UsageURLProtocol.recordedBodies.dropFirst() { + let request = try JSONDecoder().decode([String: JSONValue].self, from: body) + if provider == .gemini { + let contents = try #require(request["contents"]?.arrayValue) + let results = contents.flatMap { $0.objectValue?["parts"]?.arrayValue ?? [] } + .compactMap { $0.objectValue?["functionResponse"]?.objectValue } + let parts = try #require(results.first?["parts"]?.arrayValue) + let inlineData = try #require(parts.first?.objectValue?["inlineData"]?.objectValue) + #expect(inlineData["mime_type"] == .string("image/png")) + #expect(inlineData["data"] == .string(testImageData.base64EncodedString())) + continue + } + let messages = try #require(request["messages"]?.arrayValue) + let results = messages.flatMap { $0.objectValue?["content"]?.arrayValue ?? [] } + .filter { $0.objectValue?["type"] == .string("tool_result") } + let content = try #require(results.first?.objectValue?["content"]?.arrayValue) + let image = try #require(content.first { $0.objectValue?["type"] == .string("image") }) + let source = try #require(image.objectValue?["source"]?.objectValue) + #expect(source["type"] == .string("base64")) + #expect(source["media_type"] == .string("image/png")) + #expect(source["data"] == .string(testImageData.base64EncodedString())) + #expect(source["url"] == nil) + } + } + + @Test(arguments: [Provider.anthropic, .gemini], [false, true]) + func toolOutputImageMIMEsPreserveOnlySupportedWireParts(_ provider: Provider, _ streaming: Bool) async throws { + for mimeType in [ + "image/png", "image/jpeg", "image/gif", "image/webp", "image/tiff", "image/svg+xml", "image/avif", + "application/octet-stream", + ] { + let supported = + ["image/png", "image/jpeg", "image/webp"].contains(mimeType) + || (provider == .anthropic && mimeType == "image/gif") + let segments: [Transcript.Segment] = [ + .text(.init(content: "before")), + .image(.init(data: testImageData, mimeType: mimeType)), + .text(.init(content: "after")), + ] + let session = provider.makeSession(tools: [WeatherTool()]) + session.toolExecutionDelegate = ImageSegmentsDelegate(segments: segments) + try await checkImageToolContinuationAndReplay( + provider, + streaming: streaming, + session: session, + expectedTexts: supported ? ["before", "after"] : ["before", "", "after"], + expectedMIME: supported ? mimeType : nil + ) + let output = try #require( + session.transcript.compactMap { entry -> Transcript.ToolOutput? in + if case .toolOutput(let output) = entry { return output } + return nil + }.first + ) + #expect(output.segments == segments) + } + } + + @Test(arguments: [Provider.anthropic, .gemini], [false, true]) + func localToolImageFormatsSnapshotBeforePlaceholder(_ provider: Provider, _ streaming: Bool) async throws { + for (fileExtension, mimeType) in [("tiff", "image/tiff"), ("svg", "image/svg+xml"), ("gif", "image/gif")] { + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString).appendingPathExtension(fileExtension) + try testImageData.write(to: url) + defer { try? FileManager.default.removeItem(at: url) } + let session = provider.makeSession(tools: [LocalImageWeatherTool(url: url)]) + session.toolExecutionDelegate = DeleteImageAfterToolOutput(url: url, mimeType: mimeType) + let supported = provider == .anthropic && fileExtension == "gif" + try await checkImageToolContinuationAndReplay( + provider, + streaming: streaming, + session: session, + expectedTexts: supported ? [] : [""], + expectedMIME: supported ? mimeType : nil + ) + #expect(!FileManager.default.fileExists(atPath: url.path)) + let output = try #require( + session.transcript.compactMap { entry -> Transcript.ToolOutput? in + if case .toolOutput(let output) = entry { return output } + return nil + }.first + ) + guard case .image(let image)? = output.segments.first else { + Issue.record("Wire placeholder must not replace the canonical image") + continue + } + #expect(image.source == .data(testImageData, mimeType: mimeType)) + } + } + + private func checkImageToolContinuationAndReplay( + _ provider: Provider, + streaming: Bool, + session: LanguageModelSession, + expectedTexts: [String], + expectedMIME: String? + ) async throws { + UsageURLProtocol.reset() + UsageURLProtocol.enqueue( + json: try streaming ? provider.toolStream() : Self.json(provider.response(text: "", tool: true)) + ) + if streaming || provider == .gemini { + UsageURLProtocol.enqueue(json: try streaming ? provider.stream() : Self.json(provider.response())) + } + if streaming { + _ = try await session.streamResponse(to: "Weather?").collect() + } else { + _ = try await session.respond(to: "Weather?") + } + let archived = try JSONEncoder().encode(session.transcript) + let transcript = try JSONDecoder().decode(Transcript.self, from: archived) + #expect(transcript == session.transcript) + let restored = provider.makeSession(transcript: transcript) + UsageURLProtocol.enqueue(json: try streaming ? provider.stream() : Self.json(provider.response())) + if streaming { + _ = try await restored.streamResponse(to: "Review the image").collect() + } else { + _ = try await restored.respond(to: "Review the image") + } + #expect(UsageURLProtocol.recordedBodies.count == (streaming || provider == .gemini ? 3 : 2)) + for body in UsageURLProtocol.recordedBodies.dropFirst() { + try expectToolImage(in: body, provider: provider, texts: expectedTexts, mimeType: expectedMIME) + } + } + + private func expectToolImage(in body: Data, provider: Provider, texts: [String], mimeType: String?) throws { + let request = try JSONDecoder().decode([String: JSONValue].self, from: body) + if provider == .gemini { + let contents = try #require(request["contents"]?.arrayValue) + let results = contents.flatMap { $0.objectValue?["parts"]?.arrayValue ?? [] } + .compactMap { $0.objectValue?["functionResponse"]?.objectValue } + try #require(results.count == 1) + let result = try #require(results.first) + #expect( + result["response"]?.objectValue?["result"] + == (texts.isEmpty ? nil : .string(texts.joined(separator: "\n"))) + ) + let parts = result["parts"]?.arrayValue ?? [] + #expect(parts.count == (mimeType == nil ? 0 : 1)) + if let mimeType { + let image = try #require(parts.first?.objectValue?["inlineData"]?.objectValue) + #expect(image["mime_type"] == .string(mimeType)) + #expect(image["data"] == .string(testImageData.base64EncodedString())) + } + } else { + let messages = try #require(request["messages"]?.arrayValue) + let results = messages.flatMap { $0.objectValue?["content"]?.arrayValue ?? [] } + .filter { $0.objectValue?["type"] == .string("tool_result") } + try #require(results.count == 1) + let content = try #require(results.first?.objectValue?["content"]?.arrayValue) + #expect(content.compactMap { $0.objectValue?["text"]?.stringValue } == texts) + let images = content.filter { $0.objectValue?["type"] == .string("image") } + #expect(images.count == (mimeType == nil ? 0 : 1)) + if let mimeType { + let source = try #require(images.first?.objectValue?["source"]?.objectValue) + #expect(source["media_type"] == .string(mimeType)) + #expect(source["data"] == .string(testImageData.base64EncodedString())) + #expect(source["type"] == .string("base64")) + #expect( + content.map { $0.objectValue?["type"]?.stringValue } + == (texts.isEmpty ? ["image"] : ["text", "image", "text"]) + ) + } else { + #expect(content.allSatisfy { $0.objectValue?["type"] == .string("text") }) + } + } + } + + private struct ImageSegmentsDelegate: ToolExecutionDelegate { + let segments: [Transcript.Segment] + + func toolCallDecision(for toolCall: Transcript.ToolCall, in session: LanguageModelSession) async + -> ToolExecutionDecision + { .provideOutput(segments) } + } + + @Test(arguments: [Provider.anthropic, .gemini]) + func legacyUnsupportedToolImageURLsUsePlaceholder(_ provider: Provider) async throws { + for fileExtension in ["tiff", "svg"] { + UsageURLProtocol.reset() + UsageURLProtocol.enqueue(json: try Self.json(provider.response())) + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString).appendingPathExtension(fileExtension) + let transcript = Transcript(entries: [ + .toolCalls( + .init([.init(id: "call_1", toolName: "getWeather", arguments: GeneratedContent("Paris"))]) + ), + .toolOutput( + .init(id: "call_1", toolName: "getWeather", segments: [.image(.init(source: .url(url)))]) + ), + ]) + let archived = try JSONEncoder().encode(transcript) + let decoded = try JSONDecoder().decode(Transcript.self, from: archived) + #expect(decoded == transcript) + _ = try await provider.makeSession(transcript: decoded).respond(to: "Review the image") + try #require(UsageURLProtocol.recordedBodies.count == 1) + try expectToolImage( + in: try #require(UsageURLProtocol.recordedBodies.first), + provider: provider, + texts: [""], + mimeType: nil + ) + } + } + + @Test(arguments: [Provider.anthropic, .gemini], [false, true]) + func missingToolImageFilesStillFail(_ provider: Provider, _ streaming: Bool) async throws { + UsageURLProtocol.reset() + UsageURLProtocol.enqueue( + json: try streaming ? provider.toolStream() : Self.json(provider.response(text: "", tool: true)) + ) + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString).appendingPathExtension("png") + let session = provider.makeSession(tools: [LocalImageWeatherTool(url: url)]) + do { + if streaming { + _ = try await session.streamResponse(to: "Weather?").collect() + } else { + _ = try await session.respond(to: "Weather?") + } + Issue.record("A read failure must not become an image placeholder") + } catch let error as LanguageModelSession.ToolCallError { + #expect((error.underlyingError as NSError).domain == NSCocoaErrorDomain) + #expect((error.underlyingError as NSError).code == NSFileReadNoSuchFileError) + } + #expect(UsageURLProtocol.recordedBodies.count == 1) + } + + @Test(arguments: [Provider.anthropic, .gemini]) + func legacySupportedToolImageReadFailureIsNotAPlaceholder(_ provider: Provider) async throws { + UsageURLProtocol.reset() + let url = FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString).appendingPathExtension("png") + let transcript = Transcript(entries: [ + .toolCalls(.init([.init(id: "call_1", toolName: "getWeather", arguments: GeneratedContent("Paris"))])), + .toolOutput(.init(id: "call_1", toolName: "getWeather", segments: [.image(.init(source: .url(url)))])), + ]) + do { + _ = try await provider.makeSession(transcript: transcript).respond(to: "Review the image") + Issue.record("A supported image read failure must remain explicit") + } catch { + #expect((error as NSError).domain == NSCocoaErrorDomain) + #expect((error as NSError).code == NSFileReadNoSuchFileError) + } + #expect(UsageURLProtocol.recordedBodies.isEmpty) + } + + @Test(arguments: [Provider.anthropic, .gemini]) + func unsupportedUserImageMIMEsRemainImages(_ provider: Provider) async throws { + UsageURLProtocol.reset() + UsageURLProtocol.enqueue(json: try Self.json(provider.response())) + let transcript = Transcript(entries: [ + .prompt(.init(segments: [.image(.init(data: testImageData, mimeType: "image/tiff"))])) + ]) + _ = try await provider.makeSession(transcript: transcript).respond(to: "Review the image") + let body = try #require(UsageURLProtocol.recordedBodies.first) + let request = try JSONDecoder().decode([String: JSONValue].self, from: body) + if provider == .gemini { + let images = + request["contents"]?.arrayValue?.flatMap { $0.objectValue?["parts"]?.arrayValue ?? [] } + .compactMap { $0.objectValue?["inlineData"]?.objectValue } ?? [] + #expect(images.first?["mime_type"] == .string("image/tiff")) + } else { + let images = + request["messages"]?.arrayValue?.flatMap { $0.objectValue?["content"]?.arrayValue ?? [] } + .filter { $0.objectValue?["type"] == .string("image") } ?? [] + #expect(images.first?.objectValue?["source"]?.objectValue?["media_type"] == .string("image/tiff")) + } + } + + private struct DeleteImageAfterToolOutput: ToolExecutionDelegate { + let url: URL + var mimeType = "image/png" + + func didExecuteToolCall( + _ toolCall: Transcript.ToolCall, + output: Transcript.ToolOutput, + in session: LanguageModelSession + ) async { + guard case .image(let image)? = output.segments.first else { + Issue.record("Missing image in Tool output") + return + } + #expect(image.source == .data(testImageData, mimeType: mimeType)) + do { try FileManager.default.removeItem(at: url) } catch { Issue.record(error) } + } + } + + private struct LocalImageWeatherTool: Tool { + let name = "getWeather" + let description = "Get weather as an image" + let url: URL + + func call(arguments: WeatherTool.Arguments) async throws -> Prompt { + Prompt { Attachment(imageURL: url) } + } + } + @Test(arguments: Provider.allCases) func stoppedStreamedTools(_ provider: Provider) async throws { UsageURLProtocol.reset() @@ -881,7 +1218,7 @@ import Testing } /// Returns what `body` throws when it goes past its tool-round limit, if it does. - private static func toolRoundLimitExceeded( + static func toolRoundLimitExceeded( _ body: () async throws -> Void ) async -> LanguageModelSession.ToolRoundLimitExceeded? { do { @@ -1058,9 +1395,6 @@ import Testing #expect(UsageURLProtocol.recordedBodies.count == 1) } } - /// A `URLProtocol` that answers requests from a queue of canned responses - /// and records every request body it sees, - /// so request/response round trips can be asserted offline. extension ProviderUsageTests { // MARK: Responses that throw @@ -1143,7 +1477,10 @@ import Testing } } - private final class UsageURLProtocol: URLProtocol { + /// A `URLProtocol` that answers requests from a queue of canned responses + /// and records every request body it sees, + /// so request/response round trips can be asserted offline. + final class UsageURLProtocol: URLProtocol { struct Exchange: Sendable { var statusCode: Int = 200 var body: Data diff --git a/Tests/AnyLanguageModelTests/RequestContextPublicAPITests.swift b/Tests/AnyLanguageModelTests/RequestContextPublicAPITests.swift new file mode 100644 index 00000000..d5b313b2 --- /dev/null +++ b/Tests/AnyLanguageModelTests/RequestContextPublicAPITests.swift @@ -0,0 +1,44 @@ +import AnyLanguageModel +import Testing + +@Test func requestContextRemainsPublicWithoutSPI() { + let session = LanguageModelSession( + model: OpenAILanguageModel(apiKey: "unused", model: "test"), + tools: [PublicSessionPropertyTool()], + instructions: "Public request context" + ) + let context: LanguageModelSession.RequestContext = session.resolvedRequestContext() + #expect(context.transcript.count == session.transcript.count) + #expect(context.instructions?.description == "Public request context") + #expect(context.tools.map(\.name) == ["publicSessionProperty"]) + #expect(context.tools.first is PublicSessionPropertyTool) + let _: PublicSessionPropertyTool.SessionProperty = .init(\.history) +} + +@Test func requestContextRetainsConcreteToolIdentity() { + let original = PublicReferenceTool() + let session = LanguageModelSession( + model: OpenAILanguageModel(apiKey: "unused", model: "test"), + tools: [original] + ) + let context = session.resolvedRequestContext() + #expect(context.tools.first as? PublicReferenceTool === original) +} + +private final class PublicReferenceTool: Tool { + let name = "reference" + let description = "Preserve the supplied Tool instance" + + func call(arguments: String) async throws -> String { arguments } +} + +private struct PublicSessionPropertyTool: Tool { + @SessionProperty(\.history) private var history + + let name = "publicSessionProperty" + let description = "Read session history" + + func call(arguments: String) async throws -> String { + String(history.count) + } +} diff --git a/Tests/AnyLanguageModelTests/RequestContextTests.swift b/Tests/AnyLanguageModelTests/RequestContextTests.swift index a39cbf7a..5503f183 100644 --- a/Tests/AnyLanguageModelTests/RequestContextTests.swift +++ b/Tests/AnyLanguageModelTests/RequestContextTests.swift @@ -1,6 +1,6 @@ import Testing -@testable import AnyLanguageModel +@_spi(Compatibility) @testable import AnyLanguageModel @Suite("Request context") struct RequestContextTests { diff --git a/Tests/AnyLanguageModelTests/RequestUsageCountingTests.swift b/Tests/AnyLanguageModelTests/RequestUsageCountingTests.swift index 15836a02..f616b6a3 100644 --- a/Tests/AnyLanguageModelTests/RequestUsageCountingTests.swift +++ b/Tests/AnyLanguageModelTests/RequestUsageCountingTests.swift @@ -1,7 +1,7 @@ import Foundation import Testing -@testable import AnyLanguageModel +@_spi(Compatibility) @testable import AnyLanguageModel /// How a session counts the requests of a response it doesn't stream: each as the model reports it, /// and what the model didn't report from the response. diff --git a/Tests/AnyLanguageModelTests/SessionPropertyTests.swift b/Tests/AnyLanguageModelTests/SessionPropertyTests.swift new file mode 100644 index 00000000..6ac6d154 --- /dev/null +++ b/Tests/AnyLanguageModelTests/SessionPropertyTests.swift @@ -0,0 +1,496 @@ +import Foundation +import Observation +import Testing + +@_spi(Compatibility) @testable import AnyLanguageModel + +extension SessionPropertyValues { + @SessionPropertyEntry var revision = 0 + @SessionPropertyEntry var optionalSelection: String? = nil + @SessionPropertyEntry var items: [Int] = [] + @SessionPropertyEntry var mode: SessionPropertyMode = .first +} + +enum SessionPropertyMode { + case first + case second +} + +@Suite("Session properties") +struct SessionPropertyTests { + @Test(arguments: [false, true]) + func executionToolsCarryPropertiesIntoDetachedExecution(dynamic: Bool) async throws { + let session = + dynamic + ? LanguageModelSession( + model: SessionPropertyFixtureModel(), + dynamicInstructions: PropertyToolInstructions() + ) + : LanguageModelSession( + model: SessionPropertyFixtureModel(), + tools: [PropertyBoundTool()], + instructions: "Static" + ) + let earlier = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Earlier"))])) + session.properties.history = [earlier] + session.properties.revision = 7 + let context = session.resolvedRequestContext() + #expect(context.tools.first is PropertyBoundTool) + let tool = try #require(context.toolsForExecution(in: session).first) + let output = try await Task.detached { + #expect(SessionPropertyBinding.values == nil) + return try await tool.makeOutputSegments(from: GeneratedContent("read")) + }.value + + switch try #require(output.first) { + case .text(let text): + #expect(text.content == "revision:7 history:1") + case .structure(let structured): + #expect(try String(structured.content) == "revision:7 history:1") + case .image: + Issue.record("Unexpected image output") + } + #expect(session.properties.revision == 8) + #expect(Array(session.properties.history) == [earlier]) + #expect(SessionPropertyBinding.values == nil) + } + + @Test func propertyDefaultsCanReadAnotherProperty() { + let session = LanguageModelSession(model: SessionPropertyFixtureModel()) + SessionPropertyBinding.$values.withValue(session.properties) { + #expect(session.properties[DerivedPropertyKey.self] == 42) + session.properties[BasePropertyKey.self] = 7 + #expect(session.properties[DerivedPropertyKey.self] == 42) + } + } + + @Test(arguments: [false, true], [false, true]) + func replacingHistoryDiscardsStaleGenerationCommits(streaming: Bool, images: Bool) async throws { + let (started, start) = AsyncStream.makeStream() + let (resume, resumeContinuation) = AsyncStream.makeStream() + let session = LanguageModelSession( + model: PausedHistoryModel(started: start, resume: resume), + instructions: "Keep instructions" + ) + let request = Task { + if streaming { + let stream = + images + ? session.streamResponse(to: "Old request", images: []) + : session.streamResponse(to: "Old request") + return try await stream.collect().content + } + let response = + try await images + ? session.respond(to: "Old request", images: []) + : session.respond(to: "Old request") + return response.content + } + for await _ in started.prefix(1) {} + let replacement = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "New conversation"))])) + session.properties.history = [replacement] + resumeContinuation.yield() + + #expect(try await request.value == "Done") + #expect(Array(session.properties.history) == [replacement]) + #expect(session.transcript.count == 2) + #expect(!session.isResponding) + } + + @Test func replacingHistoryDiscardsStaleStreamingFailureCheckpoint() async { + let (started, start) = AsyncStream.makeStream() + let (resume, resumeContinuation) = AsyncStream.makeStream() + let session = LanguageModelSession(model: PausedHistoryModel(started: start, resume: resume, fail: true)) + session.transcriptErrorHandlingPolicy = .preserveTranscript + let request = Task { try await session.streamResponse(to: "Old request").collect() } + for await _ in started.prefix(1) {} + let replacement = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "New conversation"))])) + session.properties.history = [replacement] + resumeContinuation.yield() + + await #expect(throws: PausedHistoryError.self) { try await request.value } + #expect(Array(session.properties.history) == [replacement]) + #expect(!session.isResponding) + } + + @Test func dynamicInstructionsReadAndUpdateSessionProperties() { + let session = LanguageModelSession( + model: SessionPropertyFixtureModel(), + dynamicInstructions: SessionPropertyFixtureInstructions() + ) + + #expect(session.properties.revision == 0) + let first = session.resolvedRequestContext() + let second = session.resolvedRequestContext() + + #expect(first.instructions?.description == "Revision 0") + #expect(second.instructions?.description == "Revision 1") + #expect(session.properties.revision == 2) + } + + @Test func historyIsInstructionFreeAndUsesTheCurrentRequestSnapshot() { + let prompt = Transcript.Entry.prompt( + Transcript.Prompt(segments: [.text(.init(content: "Earlier"))]) + ) + let session = LanguageModelSession( + model: SessionPropertyFixtureModel(), + dynamicInstructions: SessionPropertyHistoryInstructions(), + history: [prompt] + ) + + let context = session.resolvedRequestContext() + + #expect(context.instructions?.description == "History count 1") + #expect(context.transcript.count == 2) + #expect(Array(session.properties.history) == [prompt]) + } + + @Test func requestUsesOneHistorySnapshotDuringDynamicResolution() async { + let initial = Transcript.Entry.prompt( + Transcript.Prompt(segments: [.text(.init(content: "Initial"))]) + ) + let replacement = Transcript.Entry.prompt( + Transcript.Prompt(segments: [.text(.init(content: "Replacement"))]) + ) + let (started, startedContinuation) = AsyncStream.makeStream() + let resume = DispatchSemaphore(value: 0) + let session = LanguageModelSession( + model: SessionPropertyFixtureModel(), + dynamicInstructions: BlockingHistoryInstructions( + started: startedContinuation, + resume: resume + ), + history: [initial] + ) + + let request = Task.detached { session.resolvedRequestContext() } + for await _ in started.prefix(1) {} + session.properties.history = [replacement] + resume.signal() + let context = await request.value + + #expect(context.instructions?.description == "History count 1") + #expect( + context.transcript.filter { entry in + if case .instructions = entry { return false } + return true + } == [initial] + ) + #expect(Array(session.properties.history) == [replacement]) + } + + @Test func replacingHistoryNotifiesTranscriptObservers() { + let session = LanguageModelSession(model: SessionPropertyFixtureModel()) + let changed = Locked(false) + let prompt = Transcript.Entry.prompt( + Transcript.Prompt(segments: [.text(.init(content: "Replacement"))]) + ) + + withObservationTracking { + _ = session.transcript + } onChange: { + changed.withLock { $0 = true } + } + session.properties.history = [prompt] + + #expect(changed.withLock { $0 }) + } + + @Test func generatedPropertiesParticipateInObservation() { + let session = LanguageModelSession(model: SessionPropertyFixtureModel()) + let changed = Locked(false) + + withObservationTracking { + _ = session.properties.revision + } onChange: { + changed.withLock { $0 = true } + } + session.properties.revision = 1 + + #expect(changed.withLock { $0 }) + } + + @Test func generatedPropertiesPreserveExplicitTypeContext() { + let session = LanguageModelSession(model: SessionPropertyFixtureModel()) + + #expect(session.properties.optionalSelection == nil) + #expect(session.properties.items.isEmpty) + #expect(session.properties.mode == .first) + session.properties.optionalSelection = "selected" + session.properties.items = [1, 2] + session.properties.mode = .second + #expect(session.properties.optionalSelection == "selected") + #expect(session.properties.items == [1, 2]) + #expect(session.properties.mode == .second) + } + + @Test func historyObservationTracksSessionTranscriptChanges() async throws { + let session = LanguageModelSession(model: CompletingSessionPropertyModel()) + let changed = Locked(false) + + withObservationTracking { + _ = session.properties.history + } onChange: { + changed.withLock { $0 = true } + } + _ = try await session.respond(to: "Observe history") + + #expect(changed.withLock { $0 }) + #expect(session.properties.history.count == 2) + } + + @Test func propertiesAndHistoryAreIsolatedBetweenSessions() { + let first = LanguageModelSession(model: SessionPropertyFixtureModel()) + let second = LanguageModelSession(model: SessionPropertyFixtureModel()) + let prompt = Transcript.Entry.prompt( + Transcript.Prompt(segments: [.text(.init(content: "Replacement"))]) + ) + + first.properties.revision = 7 + first.properties.history = [prompt] + + #expect(first.properties.revision == 7) + #expect(first.transcript == Transcript(entries: [prompt])) + #expect(second.properties.revision == 0) + #expect(second.transcript.isEmpty) + let _: any Observable = first.properties + } + + @Test func replacingHistoryPreservesStaticInstructions() { + let session = LanguageModelSession( + model: SessionPropertyFixtureModel(), + instructions: "Keep these instructions" + ) + let prompt = Transcript.Entry.prompt( + Transcript.Prompt(segments: [.text(.init(content: "Replacement"))]) + ) + + session.properties.history = [prompt] + + #expect(Array(session.properties.history) == [prompt]) + #expect(session.transcript.count == 2) + guard case .instructions = session.transcript[0] else { + Issue.record("Expected the static instructions entry to remain first") + return + } + #expect(session.transcript[1] == prompt) + } + + @Test func boundHistoryOnlyAppliesToItsOwningSession() { + let earlier = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Earlier"))])) + let replacement = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Replacement"))])) + let other = LanguageModelSession(model: SessionPropertyFixtureModel()) + other.properties.history = [earlier] + let session = LanguageModelSession( + model: SessionPropertyFixtureModel(), + dynamicInstructions: ForeignSessionHistoryInstructions(other: other, replacement: replacement) + ) + + let context = session.resolvedRequestContext() + + #expect(context.instructions?.description == "Other history count 1") + #expect(session.properties.history.isEmpty) + #expect(Array(other.properties.history) == [replacement]) + } + + @Test func replacingHistoryRejectsForeignInstructions() { + let other = LanguageModelSession(model: SessionPropertyFixtureModel(), instructions: "Foreign") + let prompt = Transcript.Entry.prompt(.init(segments: [.text(.init(content: "Earlier"))])) + other.properties.history = [prompt] + let session = LanguageModelSession( + model: SessionPropertyFixtureModel(), + dynamicInstructions: SessionPropertyHistoryInstructions() + ) + + session.properties.history = .init(other.transcript) + let context = session.resolvedRequestContext() + + #expect(session.transcript == Transcript(entries: [prompt])) + #expect(context.instructions?.description == "History count 1") + #expect(context.transcript.count == 2) + } +} + +private struct PropertyToolInstructions: DynamicInstructions { + var body: some DynamicInstructions { + Instructions("Property Tool") + PropertyBoundTool() + } +} + +private struct PropertyBoundTool: Tool { + @LanguageModelSession.SessionProperty(\.revision) private var revision + @LanguageModelSession.SessionProperty(\.history) private var history + + let name = "property" + let description = "Read and update session properties" + var parameters: GenerationSchema { String.generationSchema } + + func call(arguments: String) async throws -> String { + let result = "revision:\(revision) history:\(history.count)" + revision += 1 + return result + } +} + +private enum BasePropertyKey: SessionPropertyKey { + static let defaultValue = 42 +} + +private enum DerivedPropertyKey: SessionPropertyKey { + static var defaultValue: Int { + SessionPropertyBinding.values![BasePropertyKey.self] + } +} + +private enum PausedHistoryError: Error { case failed } + +private struct PausedHistoryModel: LanguageModel { + typealias UnavailableReason = Never + let started: AsyncStream.Continuation + let resume: AsyncStream + var fail = false + + private var entries: [Transcript.Entry] { + [.reasoning(.init(segments: [.text(.init(content: "Old reasoning"))]))] + } + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + started.yield() + for await _ in resume.prefix(1) {} + let raw = GeneratedContent("Done") + return .init(content: try Content(raw), rawContent: raw, transcriptEntries: entries[...]) + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + .init( + stream: AsyncThrowingStream { continuation in + Task { + do { + let raw = GeneratedContent("Done") + continuation.yield( + .init( + content: try Content.PartiallyGenerated(raw), + rawContent: raw, + transcriptEntries: entries[...] + ) + ) + started.yield() + for await _ in resume.prefix(1) {} + if fail { throw PausedHistoryError.failed } + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + } + ) + } +} + +private struct ForeignSessionHistoryInstructions: DynamicInstructions { + let other: LanguageModelSession + let replacement: Transcript.Entry + + var body: some DynamicInstructions { + let count = other.properties.history.count + other.properties.history = [replacement] + return Instructions("Other history count \(count)") + } +} + +private struct SessionPropertyFixtureInstructions: DynamicInstructions { + @SessionProperty(\.revision) private var revision + + var body: some DynamicInstructions { + let currentRevision = revision + revision += 1 + return Instructions("Revision \(currentRevision)") + } +} + +private struct SessionPropertyHistoryInstructions: DynamicInstructions { + @SessionProperty(\.history) private var history + + var body: some DynamicInstructions { + Instructions("History count \(history.count)") + } +} + +private struct BlockingHistoryInstructions: DynamicInstructions { + @SessionProperty(\.history) private var history + let started: AsyncStream.Continuation + let resume: DispatchSemaphore + + var body: some DynamicInstructions { + let count = history.count + started.yield() + resume.wait() + return Instructions("History count \(count)") + } +} + +private struct SessionPropertyFixtureModel: LanguageModel { + typealias UnavailableReason = Never + + var availability: Availability { .available } + + func supports(locale: Locale) -> Bool { true } + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + fatalError("SessionPropertyFixtureModel does not generate responses") + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + fatalError("SessionPropertyFixtureModel does not generate responses") + } +} + +private struct CompletingSessionPropertyModel: LanguageModel { + typealias UnavailableReason = Never + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + let raw = GeneratedContent("Done") + return .init(content: try Content(raw), rawContent: raw, transcriptEntries: []) + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + fatalError("This fixture only supports non-streaming responses") + } +} diff --git a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift index 7fa23bc1..e4e46c98 100644 --- a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift +++ b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift @@ -1,4 +1,4 @@ -@testable import AnyLanguageModel +@_spi(Compatibility) @testable import AnyLanguageModel struct MockLanguageModel: LanguageModel { enum UnavailableReason: Hashable, Sendable { diff --git a/Tests/AnyLanguageModelTests/Shared/ProfileProviderFixtures.swift b/Tests/AnyLanguageModelTests/Shared/ProfileProviderFixtures.swift new file mode 100644 index 00000000..fd4b1de2 --- /dev/null +++ b/Tests/AnyLanguageModelTests/Shared/ProfileProviderFixtures.swift @@ -0,0 +1,702 @@ +import Foundation +import Testing + +@_spi(Compatibility) @testable import AnyLanguageModel + +#if canImport(Darwin) && !canImport(AsyncHTTPClient) + extension ProviderUsageTests { + func providerReusedCallIDsPreserveEveryOutput( + _ provider: Provider, + streaming: Bool, + profile: Bool + ) async throws { + UsageURLProtocol.reset() + let legacyAnthropic = provider == .anthropic && !streaming && !profile + let roundsPerTurn = legacyAnthropic ? 1 : 2 + let observedCalls = Locked<[Transcript.Entry]>([]) + var expectedOutputs: [String] = [] + var expectedRequests: [[String]] = [] + for turn in 0 ..< 2 { + for round in 0 ..< roundsPerTurn { + expectedRequests.append(expectedOutputs) + let city = "Turn\(turn)Round\(round)" + var response = + try streaming + ? provider.toolStream(city: city, callID: "call_0") + : Self.json(provider.response(text: "", tool: true)) + .replacingOccurrences(of: "Paris", with: city) + .replacingOccurrences(of: "call_1", with: "call_0") + if provider == .ollama { + response = response.replacingOccurrences( + of: "\"function\":", + with: "\"id\":\"call_0\",\"function\":" + ) + } + UsageURLProtocol.enqueue(json: response) + expectedOutputs.append("OUTPUT:\(city)") + } + if !legacyAnthropic { + expectedRequests.append(expectedOutputs) + UsageURLProtocol.enqueue( + json: try streaming ? provider.stream() : Self.json(provider.response()) + ) + } + } + if legacyAnthropic { + expectedRequests.append(expectedOutputs) + UsageURLProtocol.enqueue(json: try Self.json(provider.response())) + } + let session: LanguageModelSession + if profile { + session = LanguageModelSession( + profile: LanguageModelSession.Profile { + Instructions("Repeated call IDs") + ReusedCallIDTool(observedCalls: observedCalls) + } + .model(provider.makeModel()) + .historyTransform { entries in + let properties = SessionPropertyBinding.values! + properties.history = properties.history + return entries + } + ) + } else { + session = provider.makeSession(tools: [ReusedCallIDTool(observedCalls: observedCalls)]) + } + for _ in 0 ..< (legacyAnthropic ? 3 : 2) { + if streaming { + _ = try await session.streamResponse(to: "Weather?").collect() + } else { + _ = try await session.respond(to: "Weather?") + } + } + let bodies = UsageURLProtocol.recordedBodies + let outputIDs = session.transcript.compactMap { entry -> String? in + if case .toolOutput(let output) = entry { return output.id } + return nil + } + #expect(outputIDs == Array(repeating: "call_0", count: expectedOutputs.count)) + try #require(bodies.count == expectedRequests.count) + let committedCalls = session.transcript.filter { if case .toolCalls = $0 { true } else { false } } + #expect(observedCalls.withLock { $0 } == committedCalls) + for (body, outputs) in zip(bodies, expectedRequests) { + let json = String(decoding: body, as: UTF8.self) + for output in outputs { + #expect(json.components(separatedBy: output).count - 1 == 1) + } + } + } + + struct ReusedCallIDTool: Tool { + @SessionProperty(\.history) private var history + let observedCalls: Locked<[Transcript.Entry]> + let name = "getWeather" + let description = "A deterministic Tool with distinguishable outputs" + @Generable struct Arguments { let city: String } + + func call(arguments: Arguments) async throws -> String { + if let calls = history.last(where: { if case .toolCalls = $0 { true } else { false } }) { + observedCalls.withLock { $0.append(calls) } + } + return "OUTPUT:\(arguments.city)" + } + } + + func providerProfileReevaluatesBeforeEveryToolContinuation( + _ provider: Provider, + _ streaming: Bool + ) async throws { + UsageURLProtocol.reset() + if streaming { + UsageURLProtocol.enqueue(json: try provider.toolStream()) + UsageURLProtocol.enqueue(json: try provider.stream(text: "Profile done")) + } else { + UsageURLProtocol.enqueue(json: try Self.json(provider.response(text: "", tool: true))) + UsageURLProtocol.enqueue(json: try Self.json(provider.response(text: "Profile done"))) + } + let model = provider.makeModel() + let state = ProviderProfileState() + let session = LanguageModelSession( + profile: ProviderProfile(state: state, model: model) + ) + + let response = + try await streaming + ? session.streamResponse(to: "Weather?").collect() + : session.respond(to: "Weather?") + + #expect(response.content == "Profile done") + #expect(state.executions.withLock { $0 } == ["Paris"]) + #expect(state.transformSawOutput.withLock { $0 } == [false, true]) + #expect( + state.events.withLock { $0 } == [ + "call:A", "execute:A", "output:A", + ] + ) + #expect(UsageURLProtocol.recordedBodies.count == 2) + let committedCalls = try #require( + session.transcript.first { if case .toolCalls = $0 { true } else { false } } + ) + #expect(state.observedCallEntries.withLock { $0 } == Array(repeating: committedCalls, count: 3)) + let first = String(decoding: UsageURLProtocol.recordedBodies[0], as: UTF8.self) + let second = String(decoding: UsageURLProtocol.recordedBodies[1], as: UTF8.self) + #expect(first.contains("Profile A")) + #expect(first.contains("getWeather")) + #expect(first.contains("0.2")) + #expect(second.contains("Profile B")) + #expect(second.contains("otherTool")) + #expect(second.contains("0.8")) + } + + func providerProfileCanSwitchModelForToolContinuation( + _ provider: Provider, + _ streaming: Bool + ) async throws { + UsageURLProtocol.reset() + if streaming { + UsageURLProtocol.enqueue(json: try provider.toolStream()) + } else { + UsageURLProtocol.enqueue( + json: try Self.json(provider.response(text: "", counts: provider.counts, tool: true)) + ) + } + let state = ProviderProfileState() + let session = LanguageModelSession( + profile: ProviderProfile( + state: state, + model: provider.makeModel(), + continuationModel: ProviderProfileTerminalModel(state: state) + ) + ) + + let response = + try await streaming + ? session.streamResponse(to: "Weather?").collect() + : session.respond(to: "Weather?") + + #expect(response.content == "Switched") + #expect(response.usage == .zero) + #expect(session.usage == provider.expected) + #expect(UsageURLProtocol.recordedBodies.count == 1) + #expect(response.transcriptEntries.contains { if case .toolCalls = $0 { true } else { false } }) + #expect(response.transcriptEntries.contains { if case .toolOutput = $0 { true } else { false } }) + let selected = try #require(state.selectedRequests.withLock { $0.last }) + #expect(selected.instructions == "Profile B") + #expect(selected.temperature == 0.8) + #expect(selected.tools == ["otherTool"]) + #expect(selected.sawToolOutput) + } + + func providerModelHandoffsShareToolRoundLimit(_ provider: Provider, _ streaming: Bool) async throws { + UsageURLProtocol.reset() + let limit = GenerationOptions().maximumToolRounds + for round in 0 ... limit { + if streaming { + UsageURLProtocol.enqueue( + json: try provider.toolStream( + city: "City \(round)", + callID: "call_\(round)" + ) + ) + } else { + let response = try Self.json(provider.response(text: "", tool: true)) + .replacingOccurrences(of: "Paris", with: "City \(round)") + .replacingOccurrences(of: "call_1", with: "call_\(round)") + UsageURLProtocol.enqueue(json: response) + } + } + let state = AlternatingProviderProfileState() + let session = LanguageModelSession( + profile: AlternatingProviderProfile( + state: state, + first: provider.makeModel(), + second: provider.makeModel() + ) + ) + + let exceeded = await Self.toolRoundLimitExceeded { + if streaming { + _ = try await session.streamResponse(to: "Weather?").collect() + } else { + _ = try await session.respond(to: "Weather?") + } + } + + #expect(exceeded?.limit == limit) + #expect( + exceeded?.transcriptEntries.filter { if case .toolCalls = $0 { true } else { false } }.count == limit + ) + #expect( + exceeded?.transcriptEntries.filter { if case .toolOutput = $0 { true } else { false } }.count == limit + ) + #expect(state.executions.withLock { $0.count } == limit) + #expect(UsageURLProtocol.recordedBodies.count == limit + 1) + } + + private enum HandoffResponseFormat: CaseIterable { + case text, typed, schema, images + } + + func modelHandoffsRestoreOriginalSchemaPromptOption(_ provider: Provider, _ streaming: Bool) async throws { + for format in HandoffResponseFormat.allCases { + for callSiteValue in [false, true] { + UsageURLProtocol.reset() + UsageURLProtocol.enqueue( + json: try streaming + ? provider.toolStream() + : Self.json(provider.response(text: "", tool: true)) + ) + let state = ProviderProfileState() + let raw = + format == .text + ? GeneratedContent("Switched") + : GeneratedContent(properties: ["answer": "Switched"]) + let session = LanguageModelSession( + profile: ProviderProfile( + state: state, + model: provider.makeModel(), + continuationModel: ProviderProfileTerminalModel(state: state, rawResponse: raw), + initialSchemaOverride: !callSiteValue + ) + ) + + switch format { + case .text: + let result = + try await streaming + ? session.streamResponse( + to: "Weather?", + generating: String.self, + includeSchemaInPrompt: callSiteValue + ).collect() + : session.respond( + to: "Weather?", + generating: String.self, + includeSchemaInPrompt: callSiteValue + ) + #expect(result.content == "Switched") + case .typed: + let result = + try await streaming + ? session.streamResponse( + to: "Weather?", + generating: Answer.self, + includeSchemaInPrompt: callSiteValue + ).collect() + : session.respond( + to: "Weather?", + generating: Answer.self, + includeSchemaInPrompt: callSiteValue + ) + #expect(result.content.answer == "Switched") + case .schema: + let result = + try await streaming + ? session.streamResponse( + to: "Weather?", + schema: Answer.generationSchema, + includeSchemaInPrompt: callSiteValue + ).collect() + : session.respond( + to: "Weather?", + schema: Answer.generationSchema, + includeSchemaInPrompt: callSiteValue + ) + #expect(try result.content.value(String.self, forProperty: "answer") == "Switched") + case .images: + let result = + try await streaming + ? session.streamResponse( + to: "Weather?", + images: [], + generating: Answer.self, + includeSchemaInPrompt: callSiteValue + ).collect() + : session.respond( + to: "Weather?", + images: [], + generating: Answer.self, + includeSchemaInPrompt: callSiteValue + ) + #expect(result.content.answer == "Switched") + } + let selected = try #require(state.selectedRequests.withLock { $0.last }) + #expect(selected.includeSchemaInPrompt == callSiteValue) + #expect(UsageURLProtocol.recordedBodies.count == 1) + } + } + } + + func modelHandoffsPreserveEarlierToolResultsAndHistoryWrites( + _ provider: Provider, + _ streaming: Bool + ) async throws { + UsageURLProtocol.reset() + for round in 0 ..< 2 { + let response = + try streaming + ? provider.toolStream(city: "City \(round)", callID: "call_\(round)") + : Self.json(provider.response(text: "", tool: true)) + .replacingOccurrences(of: "Paris", with: "City \(round)") + .replacingOccurrences(of: "call_1", with: "call_\(round)") + UsageURLProtocol.enqueue(json: response) + } + UsageURLProtocol.enqueue( + json: try streaming ? provider.stream() : Self.json(provider.response()) + ) + let state = AlternatingProviderProfileState(writesHistory: true) + let session = LanguageModelSession( + profile: AlternatingProviderProfile( + state: state, + first: provider.makeModel(), + second: provider.makeModel() + ) + ) + + if streaming { + _ = try await session.streamResponse(to: "Weather?").collect() + } else { + _ = try await session.respond(to: "Weather?") + } + + #expect(state.executions.withLock { $0 } == ["City 0", "City 1"]) + #expect(state.outputHistoryCounts.withLock { $0 } == [1, 2]) + #expect(state.promptIDs.withLock { $0.count } == 1) + #expect(session.transcript.count == 6) + #expect(Set(session.transcript.map(\.id)).count == session.transcript.count) + let lastRequest = try #require(UsageURLProtocol.recordedBodies.last) + let text = String(decoding: lastRequest, as: UTF8.self) + #expect(text.contains("Sunny for City 0")) + #expect(text.contains("Sunny for City 1")) + } + + func failedHandoffDoesNotPersistPendingHistory(_ provider: Provider, _ streaming: Bool) async throws { + UsageURLProtocol.reset() + for round in 0 ..< 2 { + UsageURLProtocol.enqueue( + json: try streaming + ? provider.toolStream(city: "City \(round)", callID: "call_\(round)") + : Self.json(provider.response(text: "", tool: true)) + .replacingOccurrences(of: "Paris", with: "City \(round)") + .replacingOccurrences(of: "call_1", with: "call_\(round)") + ) + } + let state = AlternatingProviderProfileState(writesHistory: true, failsAfterSecondOutput: true) + let session = LanguageModelSession( + profile: AlternatingProviderProfile( + state: state, + first: provider.makeModel(), + second: provider.makeModel() + ) + ) + + await #expect(throws: ProfileHistoryError.self) { + if streaming { + _ = try await session.streamResponse(to: "Weather?").collect() + } else { + _ = try await session.respond(to: "Weather?") + } + } + + #expect(state.outputHistoryCounts.withLock { $0 } == [1, 2]) + #expect(session.transcript.count == 1) + #expect(session.transcript.allSatisfy { if case .prompt = $0 { true } else { false } }) + } + + func ordinaryWrapperKeepsTheProviderToolLoop(_ provider: Provider, _ streaming: Bool) async throws { + UsageURLProtocol.reset() + UsageURLProtocol.enqueue( + json: try streaming ? provider.toolStream() : Self.json(provider.response(text: "", tool: true)) + ) + UsageURLProtocol.enqueue( + json: try streaming ? provider.stream() : Self.json(provider.response()) + ) + let invocations = Locked(0) + let session = LanguageModelSession( + model: ProviderWrapper(model: provider.makeModel(), invocations: invocations), + tools: [WeatherTool()] + ) + + let response = + try await streaming + ? session.streamResponse(to: "Weather?").collect() + : session.respond(to: "Weather?") + + #expect(response.content == "Hello") + #expect(invocations.withLock { $0 } == 1) + #expect(UsageURLProtocol.recordedBodies.count == 2) + #expect(session.transcript.count == 4) + } + + private final class ProviderProfileState: @unchecked Sendable { + struct SelectedRequest: Sendable { + let instructions: String? + let temperature: Double? + let tools: [String] + let sawToolOutput: Bool + let includeSchemaInPrompt: Bool + } + + let useSecond = Locked(false) + let executions = Locked<[String]>([]) + let transformSawOutput = Locked<[Bool]>([]) + let events = Locked<[String]>([]) + let observedCallEntries = Locked<[Transcript.Entry]>([]) + let selectedRequests = Locked<[SelectedRequest]>([]) + + func recordCalls(in history: Transcript.HistoryView) { + if let calls = history.last(where: { if case .toolCalls = $0 { true } else { false } }) { + observedCallEntries.withLock { $0.append(calls) } + } + } + } + + private struct ProviderProfile: LanguageModelSession.DynamicProfile, @unchecked Sendable { + let state: ProviderProfileState + let model: any LanguageModel + var continuationModel: (any LanguageModel)? = nil + var initialSchemaOverride: Bool? = nil + + var body: some LanguageModelSession.DynamicProfile { + if state.useSecond.withLock({ $0 }) { + LanguageModelSession.Profile { + Instructions("Profile B") + ProviderProfileOtherTool() + } + .model(continuationModel ?? model) + .temperature(0.8) + .historyTransform(transform) + .onToolCall { _ in state.events.withLock { $0.append("call:B") } } + .onToolOutput { _, _ in state.events.withLock { $0.append("output:B") } } + } else { + LanguageModelSession.Profile { + Instructions("Profile A") + ProviderProfileWeatherTool(state: state) + } + .model(model) + .includeSchemaInPrompt(initialSchemaOverride) + .temperature(0.2) + .historyTransform(transform) + .onToolCall { _ in + state.events.withLock { $0.append("call:A") } + state.recordCalls(in: try #require(SessionPropertyBinding.values).history) + } + .onToolOutput { _, _ in + state.events.withLock { $0.append("output:A") } + state.recordCalls(in: try #require(SessionPropertyBinding.values).history) + } + } + } + + private func transform(_ entries: [Transcript.Entry]) -> [Transcript.Entry] { + state.transformSawOutput.withLock { values in + values.append(entries.contains { if case .toolOutput = $0 { true } else { false } }) + } + return entries + } + } + + private final class AlternatingProviderProfileState: @unchecked Sendable { + let useSecond = Locked(false) + let executions = Locked<[String]>([]) + let outputHistoryCounts = Locked<[Int]>([]) + let promptIDs = Locked<[String]>([]) + let writesHistory: Bool + let failsAfterSecondOutput: Bool + + init(writesHistory: Bool = false, failsAfterSecondOutput: Bool = false) { + self.writesHistory = writesHistory + self.failsAfterSecondOutput = failsAfterSecondOutput + } + + func toggleModel() { + useSecond.withLock { $0.toggle() } + } + } + + private struct AlternatingProviderProfile: LanguageModelSession.DynamicProfile, @unchecked Sendable { + let state: AlternatingProviderProfileState + let first: any LanguageModel + let second: any LanguageModel + + var body: some LanguageModelSession.DynamicProfile { + if state.useSecond.withLock({ $0 }) { + profile(model: second) + } else { + profile(model: first) + } + } + + private func profile(model: any LanguageModel) -> some LanguageModelSession.DynamicProfile { + LanguageModelSession.Profile { AlternatingProviderProfileTool(state: state) } + .model(model) + .onPrompt { (prompt: Transcript.Prompt) in + state.promptIDs.withLock { $0.append(prompt.id) } + } + .onToolOutput { _, _ in + guard state.writesHistory else { return } + let properties = try #require(SessionPropertyBinding.values) + let history = properties.history + let outputs = history.filter { if case .toolOutput = $0 { true } else { false } } + state.outputHistoryCounts.withLock { $0.append(outputs.count) } + properties.history = history + if state.failsAfterSecondOutput && outputs.count == 2 { + throw ProfileHistoryError.rejected + } + } + .onResponse { _ in + guard state.writesHistory else { return } + let properties = try #require(SessionPropertyBinding.values) + properties.history = properties.history + } + } + } + + private struct AlternatingProviderProfileTool: Tool { + let name = "getWeather" + let description = "Get weather and switch models" + let state: AlternatingProviderProfileState + + func call(arguments: WeatherTool.Arguments) async throws -> String { + state.executions.withLock { $0.append(arguments.city) } + state.toggleModel() + return "Sunny for \(arguments.city)" + } + } + + private enum ProfileHistoryError: Error { + case rejected + } + + private struct ProviderWrapper: LanguageModel { + typealias UnavailableReason = Never + let model: any LanguageModel + let invocations: Locked + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + invocations.withLock { $0 += 1 } + return try await model.respond( + within: session, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + invocations.withLock { $0 += 1 } + return model.streamResponse( + within: session, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) + } + } + + private struct ProviderProfileWeatherTool: Tool { + @SessionProperty(\.history) private var history + let name = "getWeather" + let description = "Get weather" + let state: ProviderProfileState + + func call(arguments: WeatherTool.Arguments) async throws -> String { + state.recordCalls(in: history) + state.executions.withLock { $0.append(arguments.city) } + state.events.withLock { $0.append("execute:A") } + state.useSecond.withLock { $0 = true } + return "Sunny" + } + } + + struct ProviderProfileOtherTool: Tool { + @Generable struct Arguments {} + let name = "otherTool" + let description = "Only available in profile B" + func call(arguments: Arguments) async throws -> String { "unused" } + } + + private struct ProviderProfileTerminalModel: LanguageModel { + typealias UnavailableReason = Never + let state: ProviderProfileState + var rawResponse = GeneratedContent("Switched") + + func respond( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) async throws -> LanguageModelSession.Response where Content: Generable { + let context = session.resolvedRequestContext() + state.selectedRequests.withLock { + $0.append( + .init( + instructions: context.instructions?.description, + temperature: context.options.temperature, + tools: context.tools.map(\.name), + sawToolOutput: context.transcript.contains { + if case .toolOutput = $0 { true } else { false } + }, + includeSchemaInPrompt: includeSchemaInPrompt + ) + ) + } + let raw = rawResponse + return .init( + content: try Content(raw), + rawContent: raw, + transcriptEntries: [] + ) + } + + func streamResponse( + within session: LanguageModelSession, + to prompt: Prompt, + generating type: Content.Type, + includeSchemaInPrompt: Bool, + options: GenerationOptions + ) -> sending LanguageModelSession.ResponseStream where Content: Generable { + let stream = AsyncThrowingStream.Snapshot, any Error> { + continuation in + let task = Task { + do { + let response: LanguageModelSession.Response = try await respond( + within: session, + to: prompt, + generating: type, + includeSchemaInPrompt: includeSchemaInPrompt, + options: options + ) + continuation.yield( + .init( + content: response.content.asPartiallyGenerated(), + rawContent: response.rawContent + ) + ) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in task.cancel() } + } + return .init(stream: stream) + } + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift b/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift index 31b62a62..a7682b1f 100644 --- a/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift @@ -323,6 +323,50 @@ import Testing } #if compiler(>=6.4) + @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) + @Test(arguments: [false, true]) + func ordinaryNativeToolImageSurvivesDeletionAndReplay(useDelegate: Bool) async throws { + let data = try onePixelPNG() + let url = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + .appendingPathExtension("png") + try data.write(to: url) + defer { try? FileManager.default.removeItem(at: url) } + let tool = NativeImageOutputTool(imageURL: url) + let session = LanguageModelSession(model: SystemLanguageModel.default, tools: [tool]) + if useDelegate { session.toolExecutionDelegate = ScriptedDelegate([.execute]) } + let record = ToolCallRecord() + try await ToolCallRecord.$current.withValue(record) { + _ = try await callTool(tool, with: GeneratedContent("image"), in: session) + } + try FileManager.default.removeItem(at: url) + let transcript = Transcript(entries: await record.entries) + let replay = try JSONDecoder().decode(Transcript.self, from: JSONEncoder().encode(transcript)) + let outputs = replay.compactMap { entry -> Transcript.ToolOutput? in + if case .toolOutput(let output) = entry { return output } + return nil + } + let output = try #require(outputs.first) + #expect(output.segments.count == 2) + guard case .image(let image) = output.segments.last else { + Issue.record("Expected the native Tool's image, not its text description") + return + } + #expect(image.source == .data(data, mimeType: "image/png")) + _ = try output.segments.toFoundationModelsToolOutput() + } + + private struct NativeImageOutputTool: Tool { + let imageURL: URL + let description = "Return an image" + + func call(arguments: String) async throws -> Prompt { + Prompt { + "Image output" + Attachment(imageURL: imageURL) + } + } + } + /// An image the delegate provides goes to Foundation Models as an attachment. @available(macOS 27.0, iOS 27.0, visionOS 27.0, *) @Test func providedImageIsAttached() async throws { diff --git a/Tests/AnyLanguageModelTests/ToolContinuationTests.swift b/Tests/AnyLanguageModelTests/ToolContinuationTests.swift new file mode 100644 index 00000000..bfe7c378 --- /dev/null +++ b/Tests/AnyLanguageModelTests/ToolContinuationTests.swift @@ -0,0 +1,128 @@ +import Testing + +@_spi(Compatibility) @testable import AnyLanguageModel + +@Suite("Tool continuation history") +struct ToolContinuationTests { + @Test func explicitEmptyOfferedToolsDoesNotExecuteASessionTool() async throws { + let calls = Locked(0) + let tool = OfferedToolProbe(calls: calls) + let session = LanguageModelSession(model: MockLanguageModel(), tools: [tool]) + let call = Transcript.ToolCall(id: "call_0", toolName: tool.name, arguments: .init("input")) + let result = try await session.resolveProfileToolCalls( + [call], + requestContext: session.resolvedRequestContext(), + currentRoundEntries: [.toolCalls(.init([call]))], + offeredTools: [] + ) + guard case .invocations(let invocations) = result else { + Issue.record("Expected the existing missing-tool outcome") + return + } + #expect(calls.withLock { $0 } == 0) + #expect(invocations.count == 1) + } + + @Test func removingDurableParentDoesNotCommitItsPendingOutput() async throws { + let round = Self.round("pending") + let session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("History") } + .model(MockLanguageModel()) + .historyTransform { entries in + let properties = SessionPropertyBinding.values! + properties.history = [] + properties.history = properties.history + return entries + }, + history: [round[0]] + ) + _ = try await session.resolvedRequestContext(including: round) + #expect(session.properties.history.isEmpty) + } + + @Test func ordinaryContinuationKeepsTranscriptAndOpaqueMessages() { + let previous = Self.round("previous") + let current = Self.round("current") + let messages = toolContinuationMessages( + from: Transcript(entries: previous), + replacing: current, + with: current, + hasDynamicProfile: false, + convert: Array.init + ) + #expect(messages == previous + current) + } + + @Test func profileContinuationReplacesOnlyTheUnchangedCurrentRound() { + let previous = Self.round("previous") + let current = Self.round("current") + let opaque: [Transcript.Entry] = [.response(.init(assetIDs: [], segments: []))] + let messages = toolContinuationMessages( + from: Transcript(entries: previous + current), + replacing: current, + with: opaque, + hasDynamicProfile: true, + convert: Array.init + ) + #expect(messages == previous + opaque) + let projection = previous + current.reversed() + let reordered = toolContinuationMessages( + from: Transcript(entries: projection), + replacing: current, + with: opaque, + hasDynamicProfile: true, + convert: Array.init + ) + #expect(reordered == projection) + } + + @Test(arguments: [false, true]) + func reusedCallIDDoesNotDiscardANewTurnOutput(profile: Bool) async throws { + let previous = Self.round("previous") + let current = Self.round("current") + let session: LanguageModelSession + if profile { + session = LanguageModelSession( + profile: LanguageModelSession.Profile { Instructions("History") } + .model(MockLanguageModel()) + .historyTransform { entries in + let properties = SessionPropertyBinding.values! + properties.history = properties.history + return entries + }, + history: previous + ) + } else { + session = LanguageModelSession(model: MockLanguageModel(), transcript: Transcript(entries: previous)) + } + let context = try await session.resolvedRequestContext(including: current) + #expect(Self.outputs(context.transcript) == Self.outputs(Transcript(entries: previous + current))) + #expect(Array(session.properties.history) == previous) + } + + private static func round(_ result: String) -> [Transcript.Entry] { + let call = Transcript.ToolCall(id: "call_0", toolName: "fixture", arguments: .init(properties: [:])) + return [ + .toolCalls(.init([call])), + .toolOutput(.init(id: call.id, toolName: call.toolName, segments: [.text(.init(content: result))])), + ] + } + + private static func outputs(_ transcript: Transcript) -> [Transcript.ToolOutput] { + transcript.compactMap { entry in + if case .toolOutput(let output) = entry { return output } + return nil + } + } +} + +private struct OfferedToolProbe: Tool { + let name = "offeredProbe" + let description = "Record unexpected executions" + let calls: Locked + + func call(arguments: String) async throws -> String { + calls.withLock { $0 += 1 } + return "executed" + } +} diff --git a/Tests/AnyLanguageModelTests/ToolExecutionDelegateTests.swift b/Tests/AnyLanguageModelTests/ToolExecutionDelegateTests.swift index 1b258e32..383cca86 100644 --- a/Tests/AnyLanguageModelTests/ToolExecutionDelegateTests.swift +++ b/Tests/AnyLanguageModelTests/ToolExecutionDelegateTests.swift @@ -1,6 +1,6 @@ import Testing -@testable import AnyLanguageModel +@_spi(Compatibility) @testable import AnyLanguageModel private actor ToolExecutionDelegateSpy: ToolExecutionDelegate { private(set) var generatedToolCalls: [Transcript.ToolCall] = []