mirror of
https://github.com/firebase/firebase-ios-sdk
synced 2026-11-10 22:45:59 +05:45
786 lines
28 KiB
Swift
786 lines
28 KiB
Swift
|
// Copyright 2026 Google LLC
|
||
|
|
//
|
||
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
// you may not use this file except in compliance with the License.
|
||
|
|
// You may obtain a copy of the License at
|
||
|
|
//
|
||
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
//
|
||
|
|
// Unless required by applicable law or agreed to in writing, software
|
||
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
// See the License for the specific language governing permissions and
|
||
|
|
// limitations under the License.
|
||
|
|
|
||
|
|
#if compiler(>=6.4)
|
||
|
|
import Foundation
|
||
|
|
import Synchronization
|
||
|
|
import Testing
|
||
|
|
|
||
|
|
@testable import FoundationModelsFake
|
||
|
|
|
||
|
|
@Suite("LanguageModelSession Unit Tests")
|
||
|
|
struct LanguageModelSessionTests {
|
||
|
|
// MARK: - Test Support Types
|
||
|
|
|
||
|
|
private struct HiddenSchemaTool: Tool {
|
||
|
|
typealias Arguments = MacroFixtures.Address
|
||
|
|
let name = "hiddenTool"
|
||
|
|
let description = "A tool excluded from instructions schema."
|
||
|
|
let includesSchemaInInstructions = false
|
||
|
|
|
||
|
|
func call(arguments: MacroFixtures.Address) async throws -> String {
|
||
|
|
"hidden:\(arguments.city)"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private struct VisibleSchemaTool: Tool {
|
||
|
|
typealias Arguments = MacroFixtures.Address
|
||
|
|
let name = "visibleTool"
|
||
|
|
let description = "A tool included in instructions schema."
|
||
|
|
let includesSchemaInInstructions = true
|
||
|
|
|
||
|
|
func call(arguments: MacroFixtures.Address) async throws -> String {
|
||
|
|
"visible:\(arguments.city)"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private struct CountingModel: LanguageModel {
|
||
|
|
struct Executor: LanguageModelExecutor {
|
||
|
|
typealias Model = CountingModel
|
||
|
|
|
||
|
|
struct Configuration: Hashable, Sendable {
|
||
|
|
let key: String
|
||
|
|
}
|
||
|
|
|
||
|
|
static let initCounts = Mutex<[String: Int]>([:])
|
||
|
|
|
||
|
|
let configuration: Configuration
|
||
|
|
|
||
|
|
init(configuration: Configuration) throws {
|
||
|
|
self.configuration = configuration
|
||
|
|
Self.initCounts.withLock { counts in
|
||
|
|
counts[configuration.key, default: 0] += 1
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func prewarm(model: CountingModel, transcript: Transcript) {}
|
||
|
|
|
||
|
|
nonisolated(nonsending) func respond(
|
||
|
|
to request: LanguageModelExecutorGenerationRequest,
|
||
|
|
model: CountingModel,
|
||
|
|
streamingInto channel: LanguageModelExecutorGenerationChannel
|
||
|
|
) async throws {
|
||
|
|
await channel.send(
|
||
|
|
.response(
|
||
|
|
entryID: nil,
|
||
|
|
action: .appendText(model.replyText, tokenCount: 2)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let configKey: String
|
||
|
|
let replyText: String
|
||
|
|
|
||
|
|
var capabilities: LanguageModelCapabilities {
|
||
|
|
LanguageModelCapabilities([.guidedGeneration, .toolCalling, .reasoning, .vision])
|
||
|
|
}
|
||
|
|
|
||
|
|
var executorConfiguration: Executor.Configuration {
|
||
|
|
Executor.Configuration(key: configKey)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private struct SampleStreamError: Error, Equatable {}
|
||
|
|
|
||
|
|
private final class StreamStepGate: Sendable {
|
||
|
|
private struct State: Sendable {
|
||
|
|
var isOpen = false
|
||
|
|
var waiter: CheckedContinuation<Void, Never>?
|
||
|
|
}
|
||
|
|
|
||
|
|
private let state = Mutex(State())
|
||
|
|
|
||
|
|
func wait() async {
|
||
|
|
await withCheckedContinuation { continuation in
|
||
|
|
let resumeImmediately = state.withLock { state -> Bool in
|
||
|
|
if state.isOpen {
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
state.waiter = continuation
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
if resumeImmediately {
|
||
|
|
continuation.resume()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func open() {
|
||
|
|
let waiter = state.withLock { state -> CheckedContinuation<Void, Never>? in
|
||
|
|
state.isOpen = true
|
||
|
|
let current = state.waiter
|
||
|
|
state.waiter = nil
|
||
|
|
return current
|
||
|
|
}
|
||
|
|
waiter?.resume()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private struct CancellableStreamModel: LanguageModel {
|
||
|
|
struct Executor: LanguageModelExecutor {
|
||
|
|
typealias Model = CancellableStreamModel
|
||
|
|
|
||
|
|
struct Configuration: Hashable, Sendable {
|
||
|
|
let key: String
|
||
|
|
}
|
||
|
|
|
||
|
|
init(configuration: Configuration) throws {}
|
||
|
|
|
||
|
|
func prewarm(model: CancellableStreamModel, transcript: Transcript) {}
|
||
|
|
|
||
|
|
nonisolated(nonsending) func respond(
|
||
|
|
to request: LanguageModelExecutorGenerationRequest,
|
||
|
|
model: CancellableStreamModel,
|
||
|
|
streamingInto channel: LanguageModelExecutorGenerationChannel
|
||
|
|
) async throws {
|
||
|
|
let promptText: String =
|
||
|
|
if case .prompt(let promptEntry) = request.transcript.last,
|
||
|
|
case .text(let segment) = promptEntry.segments.first
|
||
|
|
{
|
||
|
|
segment.content
|
||
|
|
} else {
|
||
|
|
""
|
||
|
|
}
|
||
|
|
|
||
|
|
if promptText.hasPrefix("cancel:") {
|
||
|
|
await channel.send(
|
||
|
|
.response(entryID: nil, action: .appendText("partial", tokenCount: 1)))
|
||
|
|
try await Task.sleep(for: .seconds(10))
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
if promptText.hasPrefix("throw-cancel:") {
|
||
|
|
await channel.send(
|
||
|
|
.response(entryID: nil, action: .appendText("partial", tokenCount: 2)))
|
||
|
|
throw CancellationError()
|
||
|
|
}
|
||
|
|
|
||
|
|
if promptText.hasPrefix("gated:"), let gate = model.gate {
|
||
|
|
await channel.send(.response(entryID: nil, action: .appendText("part1", tokenCount: 1)))
|
||
|
|
await gate.wait()
|
||
|
|
await channel.send(.response(entryID: nil, action: .appendText("part2", tokenCount: 1)))
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
await channel.send(.response(entryID: nil, action: .appendText("ok", tokenCount: 1)))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let configKey: String
|
||
|
|
let gate: StreamStepGate?
|
||
|
|
|
||
|
|
var capabilities: LanguageModelCapabilities {
|
||
|
|
LanguageModelCapabilities([.guidedGeneration, .toolCalling, .reasoning, .vision])
|
||
|
|
}
|
||
|
|
|
||
|
|
var executorConfiguration: Executor.Configuration {
|
||
|
|
Executor.Configuration(key: configKey)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private struct ConditionalFailureModel: LanguageModel {
|
||
|
|
struct Executor: LanguageModelExecutor {
|
||
|
|
typealias Model = ConditionalFailureModel
|
||
|
|
|
||
|
|
struct Configuration: Hashable, Sendable {
|
||
|
|
let key: String
|
||
|
|
}
|
||
|
|
|
||
|
|
init(configuration: Configuration) throws {}
|
||
|
|
|
||
|
|
func prewarm(model: ConditionalFailureModel, transcript: Transcript) {}
|
||
|
|
|
||
|
|
nonisolated(nonsending) func respond(
|
||
|
|
to request: LanguageModelExecutorGenerationRequest,
|
||
|
|
model: ConditionalFailureModel,
|
||
|
|
streamingInto channel: LanguageModelExecutorGenerationChannel
|
||
|
|
) async throws {
|
||
|
|
let shouldFail =
|
||
|
|
if case .prompt(let promptEntry) = request.transcript.last,
|
||
|
|
case .text(let segment) = promptEntry.segments.first
|
||
|
|
{
|
||
|
|
segment.content.hasPrefix("fail:")
|
||
|
|
} else {
|
||
|
|
false
|
||
|
|
}
|
||
|
|
|
||
|
|
await channel.send(
|
||
|
|
.response(
|
||
|
|
entryID: nil,
|
||
|
|
action: .appendText(model.replyText, tokenCount: 2)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
if shouldFail {
|
||
|
|
throw SampleStreamError()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let configKey: String
|
||
|
|
let replyText: String
|
||
|
|
|
||
|
|
var capabilities: LanguageModelCapabilities {
|
||
|
|
LanguageModelCapabilities([.guidedGeneration, .toolCalling, .reasoning, .vision])
|
||
|
|
}
|
||
|
|
|
||
|
|
var executorConfiguration: Executor.Configuration {
|
||
|
|
Executor.Configuration(key: configKey)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// MARK: - Initializers & Tool Definition Filtering
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func sessionInitializersAndIncludesSchemaInInstructions() {
|
||
|
|
let model = CountingModel(configKey: "init-test", replyText: "ok")
|
||
|
|
let visible = VisibleSchemaTool()
|
||
|
|
let hidden = HiddenSchemaTool()
|
||
|
|
|
||
|
|
let session = LanguageModelSession(
|
||
|
|
model: model,
|
||
|
|
tools: [visible, hidden],
|
||
|
|
instructions: "Follow rules."
|
||
|
|
)
|
||
|
|
|
||
|
|
#expect(session.transcript.count == 1)
|
||
|
|
guard case .instructions(let instructionsEntry) = session.transcript.first else {
|
||
|
|
Issue.record("Expected instructions entry")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
#expect(instructionsEntry.toolDefinitions.count == 1)
|
||
|
|
#expect(instructionsEntry.toolDefinitions.first?.name == "visibleTool")
|
||
|
|
|
||
|
|
let builderSession = LanguageModelSession(model: model) {
|
||
|
|
"Line 1"
|
||
|
|
"Line 2"
|
||
|
|
}
|
||
|
|
#expect(builderSession.transcript.count == 1)
|
||
|
|
|
||
|
|
let emptySession = LanguageModelSession(model: model)
|
||
|
|
#expect(emptySession.transcript.isEmpty)
|
||
|
|
}
|
||
|
|
|
||
|
|
// MARK: - Per-Session Executor Caching
|
||
|
|
|
||
|
|
/// Executors are cached per session, as in Apple's session: sessions with equal
|
||
|
|
/// configurations get separate executors.
|
||
|
|
@Test
|
||
|
|
func executorCachingIsPerSession() async throws {
|
||
|
|
let keyA = "cache-key-A-\(UUID().uuidString)"
|
||
|
|
let keyB = "cache-key-B-\(UUID().uuidString)"
|
||
|
|
|
||
|
|
let sessionA1 = LanguageModelSession(
|
||
|
|
model: CountingModel(configKey: keyA, replyText: "first"))
|
||
|
|
let sessionA2 = LanguageModelSession(
|
||
|
|
model: CountingModel(configKey: keyA, replyText: "second"))
|
||
|
|
let sessionB = LanguageModelSession(model: CountingModel(configKey: keyB, replyText: "third"))
|
||
|
|
|
||
|
|
let r1 = try await sessionA1.respond(to: "Hi")
|
||
|
|
let r2 = try await sessionA2.respond(to: "Hi again")
|
||
|
|
let r3 = try await sessionB.respond(to: "Hi B")
|
||
|
|
|
||
|
|
#expect(r1.content == "first")
|
||
|
|
#expect(r2.content == "second")
|
||
|
|
#expect(r3.content == "third")
|
||
|
|
|
||
|
|
let countA = CountingModel.Executor.initCounts.withLock { $0[keyA] ?? 0 }
|
||
|
|
let countB = CountingModel.Executor.initCounts.withLock { $0[keyB] ?? 0 }
|
||
|
|
#expect(countA == 2)
|
||
|
|
#expect(countB == 1)
|
||
|
|
}
|
||
|
|
|
||
|
|
// MARK: - All 18 respond Overloads
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func allEighteenRespondOverloads() async throws {
|
||
|
|
let textModel = CountingModel(configKey: "respond-text", replyText: "hello")
|
||
|
|
let jsonModel = CountingModel(
|
||
|
|
configKey: "respond-json",
|
||
|
|
replyText: #"{"street":"1 Main","city":"London"}"#
|
||
|
|
)
|
||
|
|
let expectedAddress = MacroFixtures.Address(street: "1 Main", city: "London")
|
||
|
|
|
||
|
|
let sText = LanguageModelSession(model: textModel)
|
||
|
|
let sJSON = LanguageModelSession(model: jsonModel)
|
||
|
|
let schema = MacroFixtures.Address.generationSchema
|
||
|
|
|
||
|
|
// Legacy 1-3: String output
|
||
|
|
let r1 = try await sText.respond(to: Prompt("p1"))
|
||
|
|
let r2 = try await sText.respond(to: "p2")
|
||
|
|
let r3 = try await sText.respond { "p3" }
|
||
|
|
#expect(r1.content == "hello")
|
||
|
|
#expect(r2.content == "hello")
|
||
|
|
#expect(r3.content == "hello")
|
||
|
|
|
||
|
|
// Legacy 4-6: GenerationSchema output
|
||
|
|
let r4 = try await sJSON.respond(to: Prompt("p4"), schema: schema)
|
||
|
|
let r5 = try await sJSON.respond(to: "p5", schema: schema)
|
||
|
|
let r6 = try await sJSON.respond(schema: schema) { "p6" }
|
||
|
|
#expect(try MacroFixtures.Address(r4.content) == expectedAddress)
|
||
|
|
#expect(try MacroFixtures.Address(r5.content) == expectedAddress)
|
||
|
|
#expect(try MacroFixtures.Address(r6.content) == expectedAddress)
|
||
|
|
|
||
|
|
// Legacy 7-9: Generable type output
|
||
|
|
let r7 = try await sJSON.respond(to: Prompt("p7"), generating: MacroFixtures.Address.self)
|
||
|
|
let r8 = try await sJSON.respond(to: "p8", generating: MacroFixtures.Address.self)
|
||
|
|
let r9 = try await sJSON.respond(generating: MacroFixtures.Address.self) { "p9" }
|
||
|
|
#expect(r7.content == expectedAddress)
|
||
|
|
#expect(r8.content == expectedAddress)
|
||
|
|
#expect(r9.content == expectedAddress)
|
||
|
|
|
||
|
|
// New 10-12: String output with contextOptions & metadata
|
||
|
|
let r10 = try await sText.respond(
|
||
|
|
to: Prompt("p10"),
|
||
|
|
contextOptions: ContextOptions(),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
)
|
||
|
|
let r11 = try await sText.respond(
|
||
|
|
to: "p11",
|
||
|
|
contextOptions: ContextOptions(),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
)
|
||
|
|
let r12 = try await sText.respond(
|
||
|
|
contextOptions: ContextOptions(),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
) { "p12" }
|
||
|
|
#expect(r10.content == "hello")
|
||
|
|
#expect(r11.content == "hello")
|
||
|
|
#expect(r12.content == "hello")
|
||
|
|
|
||
|
|
// New 13-15: GenerationSchema output with contextOptions & metadata
|
||
|
|
let r13 = try await sJSON.respond(
|
||
|
|
to: Prompt("p13"),
|
||
|
|
schema: schema,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
)
|
||
|
|
let r14 = try await sJSON.respond(
|
||
|
|
to: "p14",
|
||
|
|
schema: schema,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
)
|
||
|
|
let r15 = try await sJSON.respond(
|
||
|
|
schema: schema,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
) { "p15" }
|
||
|
|
#expect(try MacroFixtures.Address(r13.content) == expectedAddress)
|
||
|
|
#expect(try MacroFixtures.Address(r14.content) == expectedAddress)
|
||
|
|
#expect(try MacroFixtures.Address(r15.content) == expectedAddress)
|
||
|
|
|
||
|
|
// New 16-18: Generable type output with contextOptions & metadata
|
||
|
|
let r16 = try await sJSON.respond(
|
||
|
|
to: Prompt("p16"),
|
||
|
|
generating: MacroFixtures.Address.self,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
)
|
||
|
|
let r17 = try await sJSON.respond(
|
||
|
|
to: "p17",
|
||
|
|
generating: MacroFixtures.Address.self,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
)
|
||
|
|
let r18 = try await sJSON.respond(
|
||
|
|
generating: MacroFixtures.Address.self,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
) { "p18" }
|
||
|
|
#expect(r16.content == expectedAddress)
|
||
|
|
#expect(r17.content == expectedAddress)
|
||
|
|
#expect(r18.content == expectedAddress)
|
||
|
|
|
||
|
|
// Generable where Content == String sets schema to nil
|
||
|
|
let rStringGenerable = try await sText.respond(to: "pString", generating: String.self)
|
||
|
|
#expect(rStringGenerable.content == "hello")
|
||
|
|
}
|
||
|
|
|
||
|
|
// MARK: - All 18 streamResponse Overloads
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func allEighteenStreamResponseOverloads() async throws {
|
||
|
|
let textModel = CountingModel(configKey: "stream-text", replyText: "streamed")
|
||
|
|
let jsonModel = CountingModel(
|
||
|
|
configKey: "stream-json",
|
||
|
|
replyText: #"{"street":"1 Main","city":"London"}"#
|
||
|
|
)
|
||
|
|
let expectedAddress = MacroFixtures.Address(street: "1 Main", city: "London")
|
||
|
|
|
||
|
|
let sText = LanguageModelSession(model: textModel)
|
||
|
|
let sJSON = LanguageModelSession(model: jsonModel)
|
||
|
|
let schema = MacroFixtures.Address.generationSchema
|
||
|
|
|
||
|
|
// Legacy 1-3: String output
|
||
|
|
let r1 = try await sText.streamResponse(to: Prompt("p1")).collect()
|
||
|
|
let r2 = try await sText.streamResponse(to: "p2").collect()
|
||
|
|
let r3 = try await sText.streamResponse { "p3" }.collect()
|
||
|
|
#expect(r1.content == "streamed")
|
||
|
|
#expect(r2.content == "streamed")
|
||
|
|
#expect(r3.content == "streamed")
|
||
|
|
|
||
|
|
// Legacy 4-6: GenerationSchema output
|
||
|
|
let r4 = try await sJSON.streamResponse(to: Prompt("p4"), schema: schema).collect()
|
||
|
|
let r5 = try await sJSON.streamResponse(to: "p5", schema: schema).collect()
|
||
|
|
let r6 = try await sJSON.streamResponse(schema: schema) { "p6" }.collect()
|
||
|
|
#expect(try MacroFixtures.Address(r4.content) == expectedAddress)
|
||
|
|
#expect(try MacroFixtures.Address(r5.content) == expectedAddress)
|
||
|
|
#expect(try MacroFixtures.Address(r6.content) == expectedAddress)
|
||
|
|
|
||
|
|
// Legacy 7-9: Generable type output
|
||
|
|
let r7 = try await sJSON.streamResponse(
|
||
|
|
to: Prompt("p7"),
|
||
|
|
generating: MacroFixtures.Address.self
|
||
|
|
).collect()
|
||
|
|
let r8 = try await sJSON.streamResponse(
|
||
|
|
to: "p8",
|
||
|
|
generating: MacroFixtures.Address.self
|
||
|
|
).collect()
|
||
|
|
let r9 = try await sJSON.streamResponse(generating: MacroFixtures.Address.self) {
|
||
|
|
"p9"
|
||
|
|
}.collect()
|
||
|
|
#expect(r7.content == expectedAddress)
|
||
|
|
#expect(r8.content == expectedAddress)
|
||
|
|
#expect(r9.content == expectedAddress)
|
||
|
|
|
||
|
|
// New 10-12: String output with contextOptions & metadata
|
||
|
|
let r10 = try await sText.streamResponse(
|
||
|
|
to: Prompt("p10"),
|
||
|
|
contextOptions: ContextOptions(),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
).collect()
|
||
|
|
let r11 = try await sText.streamResponse(
|
||
|
|
to: "p11",
|
||
|
|
contextOptions: ContextOptions(),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
).collect()
|
||
|
|
let r12 = try await sText.streamResponse(
|
||
|
|
contextOptions: ContextOptions(),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
) { "p12" }.collect()
|
||
|
|
#expect(r10.content == "streamed")
|
||
|
|
#expect(r11.content == "streamed")
|
||
|
|
#expect(r12.content == "streamed")
|
||
|
|
|
||
|
|
// New 13-15: GenerationSchema output with contextOptions & metadata
|
||
|
|
let r13 = try await sJSON.streamResponse(
|
||
|
|
to: Prompt("p13"),
|
||
|
|
schema: schema,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
).collect()
|
||
|
|
let r14 = try await sJSON.streamResponse(
|
||
|
|
to: "p14",
|
||
|
|
schema: schema,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
).collect()
|
||
|
|
let r15 = try await sJSON.streamResponse(
|
||
|
|
schema: schema,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
) { "p15" }.collect()
|
||
|
|
#expect(try MacroFixtures.Address(r13.content) == expectedAddress)
|
||
|
|
#expect(try MacroFixtures.Address(r14.content) == expectedAddress)
|
||
|
|
#expect(try MacroFixtures.Address(r15.content) == expectedAddress)
|
||
|
|
|
||
|
|
// New 16-18: Generable type output with contextOptions & metadata
|
||
|
|
let r16 = try await sJSON.streamResponse(
|
||
|
|
to: Prompt("p16"),
|
||
|
|
generating: MacroFixtures.Address.self,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
).collect()
|
||
|
|
let r17 = try await sJSON.streamResponse(
|
||
|
|
to: "p17",
|
||
|
|
generating: MacroFixtures.Address.self,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
).collect()
|
||
|
|
let r18 = try await sJSON.streamResponse(
|
||
|
|
generating: MacroFixtures.Address.self,
|
||
|
|
contextOptions: ContextOptions(includeSchemaInPrompt: false),
|
||
|
|
metadata: ["k": "v"]
|
||
|
|
) { "p18" }.collect()
|
||
|
|
#expect(r16.content == expectedAddress)
|
||
|
|
#expect(r17.content == expectedAddress)
|
||
|
|
#expect(r18.content == expectedAddress)
|
||
|
|
}
|
||
|
|
|
||
|
|
// MARK: - Refusal Explanation & ExplanationStream
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func refusalExplanationAndStream() async throws {
|
||
|
|
let refusal = LanguageModelError.Refusal(
|
||
|
|
explanation: "I cannot assist with that.",
|
||
|
|
debugDescription: "Refused"
|
||
|
|
)
|
||
|
|
|
||
|
|
let explanation = try await refusal.explanation
|
||
|
|
#expect(explanation.content == "I cannot assist with that.")
|
||
|
|
#expect(explanation.rawContent == GeneratedContent("I cannot assist with that."))
|
||
|
|
#expect(explanation.transcriptEntries.isEmpty)
|
||
|
|
#expect(explanation.usage.totalTokenCount == 0)
|
||
|
|
|
||
|
|
var snapshots: [String] = []
|
||
|
|
let stream = refusal.explanationStream
|
||
|
|
for try await snap in stream {
|
||
|
|
snapshots.append(snap.content)
|
||
|
|
}
|
||
|
|
let collected = try await stream.collect()
|
||
|
|
#expect(snapshots == ["I cannot assist with that."])
|
||
|
|
#expect(collected.content == "I cannot assist with that.")
|
||
|
|
#expect(collected.transcriptEntries.isEmpty)
|
||
|
|
}
|
||
|
|
|
||
|
|
// MARK: - Stream Final Conversion Failure & Repeated ToolCall ID Flat Lookup
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func streamResponseConvertFinalFailureFinishesSnapshotsBeforeFailingCollect() async throws {
|
||
|
|
// Missing required "city" field so PartiallyGenerated succeeds on snapshots,
|
||
|
|
// while final Address conversion throws in collect().
|
||
|
|
let incompleteModel = CountingModel(
|
||
|
|
configKey: "stream-incomplete-address",
|
||
|
|
replyText: #"{"street":"1 Main"}"#
|
||
|
|
)
|
||
|
|
let session = LanguageModelSession(model: incompleteModel)
|
||
|
|
let stream = session.streamResponse(
|
||
|
|
to: "Make an address",
|
||
|
|
generating: MacroFixtures.Address.self
|
||
|
|
)
|
||
|
|
|
||
|
|
var snapshots: [MacroFixtures.Address.PartiallyGenerated] = []
|
||
|
|
for try await snap in stream {
|
||
|
|
snapshots.append(snap.content)
|
||
|
|
}
|
||
|
|
|
||
|
|
#expect(!snapshots.isEmpty)
|
||
|
|
#expect(snapshots.last?.street == "1 Main")
|
||
|
|
#expect(snapshots.last?.city == nil)
|
||
|
|
await #expect(throws: GeneratedContent.ParsingError.self) {
|
||
|
|
_ = try await stream.collect()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func repeatedToolCallIDWithNilEntryIDUpdatesExistingEntryInPlace() throws {
|
||
|
|
var transcript = Transcript(entries: [
|
||
|
|
.toolCalls(
|
||
|
|
Transcript.ToolCalls(
|
||
|
|
id: "earlier-entry",
|
||
|
|
[
|
||
|
|
Transcript.ToolCall(
|
||
|
|
id: "call-1",
|
||
|
|
toolName: "visibleTool",
|
||
|
|
arguments: try GeneratedContent(json: #"{"street":"Old","city":"Rome"}"#)
|
||
|
|
)
|
||
|
|
]
|
||
|
|
)
|
||
|
|
)
|
||
|
|
])
|
||
|
|
|
||
|
|
var reducer = LanguageModelSession.TurnReducer(schema: nil)
|
||
|
|
let event = LanguageModelExecutorGenerationChannel.Event.toolCalls(
|
||
|
|
entryID: nil,
|
||
|
|
action: .toolCall(
|
||
|
|
id: "call-1",
|
||
|
|
name: "visibleTool",
|
||
|
|
action: .appendArguments(#"{"street":"New","city":"Oslo"}"#, tokenCount: 2)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
_ = try reducer.apply(event, to: &transcript)
|
||
|
|
|
||
|
|
#expect(transcript.count == 1)
|
||
|
|
#expect(reducer.createdToolCallsEntryIndices.isEmpty)
|
||
|
|
let entry = try #require(transcript.first)
|
||
|
|
guard case .toolCalls(let callsEntry) = entry else {
|
||
|
|
Issue.record("Expected toolCalls entry")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
#expect(callsEntry.id == "earlier-entry")
|
||
|
|
#expect(callsEntry.calls.count == 1)
|
||
|
|
let updatedCall = try #require(callsEntry.calls.first)
|
||
|
|
let decoded = try MacroFixtures.Address(updatedCall.arguments)
|
||
|
|
#expect(decoded == MacroFixtures.Address(street: "New", city: "Oslo"))
|
||
|
|
}
|
||
|
|
|
||
|
|
// MARK: - Stream Completion Ordering Regression Tests
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func backToBackStreamResponseAndRespondNeverRacesIsResponding() async throws {
|
||
|
|
let model = CountingModel(configKey: "race-regression-ok", replyText: "ok")
|
||
|
|
let session = LanguageModelSession(model: model)
|
||
|
|
let iterations = 200
|
||
|
|
|
||
|
|
for i in 0..<iterations {
|
||
|
|
let c1 = try await session.streamResponse(to: "collect-a-\(i)").collect()
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
#expect(c1.content == "ok")
|
||
|
|
|
||
|
|
let c2 = try await session.streamResponse(to: "collect-b-\(i)").collect()
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
#expect(c2.content == "ok")
|
||
|
|
|
||
|
|
var lastSnapshot: String?
|
||
|
|
for try await snapshot in session.streamResponse(to: "iter-\(i)") {
|
||
|
|
lastSnapshot = snapshot.content
|
||
|
|
}
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
#expect(lastSnapshot == "ok")
|
||
|
|
|
||
|
|
let r = try await session.respond(to: "respond-\(i)")
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
#expect(r.content == "ok")
|
||
|
|
}
|
||
|
|
|
||
|
|
#expect(session.transcript.count == iterations * 4 * 2)
|
||
|
|
#expect(session.usage.totalTokenCount == iterations * 4 * 2)
|
||
|
|
}
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func backToBackStreamResponseFailureImmediatelyLeavesSessionIdle() async throws {
|
||
|
|
let failingOrOkModel = ConditionalFailureModel(
|
||
|
|
configKey: "race-regression-fail",
|
||
|
|
replyText: #"{"street":"1 Main"}"#
|
||
|
|
)
|
||
|
|
let session = LanguageModelSession(model: failingOrOkModel)
|
||
|
|
let iterations = 200
|
||
|
|
|
||
|
|
for i in 0..<iterations {
|
||
|
|
await #expect(throws: SampleStreamError.self) {
|
||
|
|
_ = try await session.streamResponse(to: "fail:collect-\(i)").collect()
|
||
|
|
}
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
|
||
|
|
await #expect(throws: SampleStreamError.self) {
|
||
|
|
for try await _ in session.streamResponse(to: "fail:iter-\(i)") {}
|
||
|
|
}
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
|
||
|
|
await #expect(throws: GeneratedContent.ParsingError.self) {
|
||
|
|
_ = try await session.streamResponse(
|
||
|
|
to: "incomplete-address-\(i)",
|
||
|
|
generating: MacroFixtures.Address.self
|
||
|
|
).collect()
|
||
|
|
}
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
|
||
|
|
let recovered = try await session.respond(to: "ok-\(i)")
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
#expect(recovered.content == #"{"street":"1 Main"}"#)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Each iteration has 2 reverted mid-stream errors + 1 committed final-conversion failure
|
||
|
|
// (prompt + response) + 1 committed respond (prompt + response) = 4 entries per iteration.
|
||
|
|
#expect(session.transcript.count == iterations * 4)
|
||
|
|
}
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func cancelledStreamResponseRevertRemovesPromptAndLeavesSessionIdle() async throws {
|
||
|
|
let model = CancellableStreamModel(configKey: "cancel-stream-revert-200", gate: nil)
|
||
|
|
let session = LanguageModelSession(model: model)
|
||
|
|
let iterations = 200
|
||
|
|
|
||
|
|
for i in 0..<iterations {
|
||
|
|
let sawFirstSnapshot = StreamStepGate()
|
||
|
|
let task = Task {
|
||
|
|
for try await _ in session.streamResponse(to: "cancel:\(i)") {
|
||
|
|
sawFirstSnapshot.open()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
await sawFirstSnapshot.wait()
|
||
|
|
task.cancel()
|
||
|
|
let taskResult = await task.result
|
||
|
|
|
||
|
|
#expect(throws: Never.self) { try taskResult.get() }
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
#expect(session.transcript.count == i * 2)
|
||
|
|
let response = try await session.respond(to: "after:\(i)")
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
#expect(response.content == "ok")
|
||
|
|
#expect(session.transcript.count == (i + 1) * 2)
|
||
|
|
}
|
||
|
|
|
||
|
|
#expect(session.usage.totalTokenCount == iterations)
|
||
|
|
}
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func executorThrownCancellationErrorBehavesAsExecutorError() async throws {
|
||
|
|
let model = CancellableStreamModel(configKey: "executor-throws-cancellation", gate: nil)
|
||
|
|
let revertStreamSession = LanguageModelSession(model: model)
|
||
|
|
revertStreamSession.transcriptErrorHandlingPolicy = .revertTranscript
|
||
|
|
let preserveStreamSession = LanguageModelSession(model: model)
|
||
|
|
preserveStreamSession.transcriptErrorHandlingPolicy = .preserveTranscript
|
||
|
|
let preserveRespondSession = LanguageModelSession(model: model)
|
||
|
|
preserveRespondSession.transcriptErrorHandlingPolicy = .preserveTranscript
|
||
|
|
|
||
|
|
await #expect(throws: CancellationError.self) {
|
||
|
|
for try await _ in revertStreamSession.streamResponse(to: "throw-cancel:revert") {}
|
||
|
|
}
|
||
|
|
await #expect(throws: CancellationError.self) {
|
||
|
|
for try await _ in preserveStreamSession.streamResponse(to: "throw-cancel:preserve") {}
|
||
|
|
}
|
||
|
|
await #expect(throws: CancellationError.self) {
|
||
|
|
_ = try await preserveRespondSession.respond(to: "throw-cancel:respond")
|
||
|
|
}
|
||
|
|
|
||
|
|
#expect(revertStreamSession.transcript.isEmpty)
|
||
|
|
#expect(revertStreamSession.usage.output.totalTokenCount == 2)
|
||
|
|
#expect(preserveStreamSession.transcript.count == 2)
|
||
|
|
#expect(preserveStreamSession.usage.output.totalTokenCount == 2)
|
||
|
|
#expect(preserveRespondSession.transcript.count == 1)
|
||
|
|
#expect(preserveRespondSession.usage.output.totalTokenCount == 0)
|
||
|
|
}
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func shrinkingTranscriptMidStreamThrowsTranscriptMutationWhileResponding() async throws {
|
||
|
|
let gate = StreamStepGate()
|
||
|
|
let model = CancellableStreamModel(configKey: "shrink-transcript-mid-stream", gate: gate)
|
||
|
|
let session = LanguageModelSession(model: model)
|
||
|
|
_ = try await session.respond(to: "seed")
|
||
|
|
let seededCount = session.transcript.count
|
||
|
|
|
||
|
|
let stream = session.streamResponse(to: "gated:shrink")
|
||
|
|
var iterator = stream.makeAsyncIterator()
|
||
|
|
let firstSnapshot = try await iterator.next()
|
||
|
|
session.transcript = Transcript()
|
||
|
|
gate.open()
|
||
|
|
|
||
|
|
#expect(firstSnapshot?.content == "part1")
|
||
|
|
await #expect(throws: LanguageModelSession.Error.transcriptMutationWhileResponding) {
|
||
|
|
while try await iterator.next() != nil {}
|
||
|
|
}
|
||
|
|
await #expect(throws: LanguageModelSession.Error.transcriptMutationWhileResponding) {
|
||
|
|
_ = try await stream.collect()
|
||
|
|
}
|
||
|
|
#expect(!session.isResponding)
|
||
|
|
#expect(session.transcript.count == seededCount)
|
||
|
|
}
|
||
|
|
|
||
|
|
@Test
|
||
|
|
func anySendableHashableDistinguishesDynamicTypesWithCompatibleCasts() {
|
||
|
|
let directValue: Int = 7
|
||
|
|
let optionalValue: Int? = 7
|
||
|
|
|
||
|
|
let wrappedDirect = LanguageModelSession.AnySendableHashable(directValue)
|
||
|
|
let wrappedOptional = LanguageModelSession.AnySendableHashable(optionalValue)
|
||
|
|
|
||
|
|
#expect(((directValue as any Hashable) as? Int?) == optionalValue)
|
||
|
|
#expect(((optionalValue as any Hashable) as? Int) == directValue)
|
||
|
|
#expect(wrappedDirect != wrappedOptional)
|
||
|
|
#expect(wrappedOptional != wrappedDirect)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
#endif // compiler(>=6.4)
|