diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 1b569491..e25fa03a 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -77,7 +77,10 @@ public final class LanguageModelSession: @unchecked Sendable { let instructionsEntry = Transcript.Entry.instructions( Transcript.Instructions( segments: [.text(.init(content: instructions.description))], - toolDefinitions: tools.map { Transcript.ToolDefinition(tool: $0) } + toolDefinitions: + tools + .filter(\.includesSchemaInInstructions) + .map { Transcript.ToolDefinition(tool: $0) } ) ) finalTranscript.append(instructionsEntry) diff --git a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift index 52c1db90..97fddb0b 100644 --- a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift @@ -78,7 +78,12 @@ let fmSession = FoundationModels.LanguageModelSession( model: systemModel, tools: session.tools.toFoundationModels(), - transcript: session.transcript.toFoundationModels(instructions: session.instructions) + transcript: session.transcript.toFoundationModels( + instructions: session.instructions, + toolDefinitions: session.tools + .filter(\.includesSchemaInInstructions) + .map { Transcript.ToolDefinition(tool: $0) } + ) ) if type == String.self { @@ -154,7 +159,12 @@ let fmSession = FoundationModels.LanguageModelSession( model: systemModel, tools: session.tools.toFoundationModels(), - transcript: session.transcript.toFoundationModels(instructions: session.instructions) + transcript: session.transcript.toFoundationModels( + instructions: session.instructions, + toolDefinitions: session.tools + .filter(\.includesSchemaInInstructions) + .map { Transcript.ToolDefinition(tool: $0) } + ) ) let stream: AsyncThrowingStream.Snapshot, Error> = @@ -642,8 +652,10 @@ @available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *) extension Transcript { - fileprivate func toFoundationModels(instructions: AnyLanguageModel.Instructions?) -> FoundationModels.Transcript - { + fileprivate func toFoundationModels( + instructions: AnyLanguageModel.Instructions?, + toolDefinitions: [Transcript.ToolDefinition] + ) -> FoundationModels.Transcript { var fmEntries: [FoundationModels.Transcript.Entry] = [] // Add instructions entry if provided and not already in transcript @@ -656,7 +668,7 @@ if !hasInstructions { let fmInstructions = FoundationModels.Transcript.Instructions( segments: [.text(.init(content: instructions.description))], - toolDefinitions: [] + toolDefinitions: toolDefinitions.toFoundationModels() ) fmEntries.append(.instructions(fmInstructions)) } diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index baff85dd..ad8e1c06 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -2,6 +2,37 @@ import Testing @testable import AnyLanguageModel +private struct ToolSchemaOptOutTool: Tool { + let name = "schemaOptOutTool" + let description = "A tool that opts out of schema injection." + let includesSchemaInInstructions = false + + @Generable + struct Arguments { + @Guide(description: "A value to echo") + var value: String + } + + func call(arguments: Arguments) async throws -> String { + arguments.value + } +} + +private func waitUntil( + timeout: Duration = .milliseconds(500), + pollInterval: Duration = .milliseconds(5), + _ condition: @autoclosure () -> Bool +) async throws -> Bool { + let deadline = ContinuousClock.now + timeout + while !condition() { + if ContinuousClock.now >= deadline { + return false + } + try await Task.sleep(for: pollInterval) + } + return true +} + @Suite("MockLanguageModel") struct MockLanguageModelTests { @Test func fixedResponse() async throws { @@ -69,6 +100,92 @@ struct MockLanguageModelTests { } } + @Test func instructionsOnlyIncludeToolsThatOptInToSchemaInjection() { + let model = MockLanguageModel.fixed("ok") + let tools: [any Tool] = [WeatherTool(), ToolSchemaOptOutTool()] + let session = LanguageModelSession( + model: model, + tools: tools, + instructions: "Use tools when appropriate." + ) + + #expect(session.transcript.count == 1) + guard case let .instructions(instructionsEntry) = session.transcript[0] else { + Issue.record("First entry should be instructions") + return + } + + #expect(instructionsEntry.toolDefinitions.count == 1) + #expect(instructionsEntry.toolDefinitions.first?.name == "getWeather") + } + + @Test func transcriptAndInstructionsInitHaveEquivalentToolDefinitionBehavior() { + let model = MockLanguageModel.fixed("ok") + let tools: [any Tool] = [WeatherTool(), ToolSchemaOptOutTool()] + let expectedToolDefinitions = + tools + .filter(\.includesSchemaInInstructions) + .map { Transcript.ToolDefinition(tool: $0) } + + let instructionsSession = LanguageModelSession( + model: model, + tools: tools, + instructions: "Use tools when appropriate." + ) + + let explicitTranscript = Transcript(entries: [ + .instructions( + Transcript.Instructions( + segments: [.text(.init(content: "Use tools when appropriate."))], + toolDefinitions: expectedToolDefinitions + ) + ) + ]) + let transcriptSession = LanguageModelSession( + model: model, + tools: tools, + transcript: explicitTranscript + ) + + guard case let .instructions(instructionsEntry) = instructionsSession.transcript[0] else { + Issue.record("First instructions-based entry should be instructions") + return + } + guard case let .instructions(transcriptEntry) = transcriptSession.transcript[0] else { + Issue.record("First transcript-based entry should be instructions") + return + } + + #expect(instructionsEntry.toolDefinitions == transcriptEntry.toolDefinitions) + } + + @Test func explicitTranscriptInstructionsArePreservedWhenToolsOptOut() { + let model = MockLanguageModel.fixed("ok") + let tools: [any Tool] = [WeatherTool(), ToolSchemaOptOutTool()] + let explicitToolDefinitions = tools.map { Transcript.ToolDefinition(tool: $0) } + let transcript = Transcript(entries: [ + .instructions( + Transcript.Instructions( + segments: [.text(.init(content: "Custom instructions."))], + toolDefinitions: explicitToolDefinitions + ) + ) + ]) + + let session = LanguageModelSession( + model: model, + tools: tools, + transcript: transcript + ) + + guard case let .instructions(entry) = session.transcript[0] else { + Issue.record("First entry should be instructions") + return + } + + #expect(entry.toolDefinitions == explicitToolDefinitions) + } + @Test func unavailable() async throws { let model = MockLanguageModel.unavailable @@ -91,8 +208,7 @@ struct MockLanguageModelTests { try await asyncSession.respond(to: "Async test") } - try await Task.sleep(for: .milliseconds(50)) - #expect(asyncSession.isResponding == true) + #expect(try await waitUntil(asyncSession.isResponding == true)) _ = try await asyncTask.value try await Task.sleep(for: .milliseconds(10)) @@ -112,8 +228,7 @@ struct MockLanguageModelTests { for try await _ in stream {} } - try await Task.sleep(for: .milliseconds(50)) - #expect(streamSession.isResponding == true) + #expect(try await waitUntil(streamSession.isResponding == true)) _ = try await streamTask.value try await Task.sleep(for: .milliseconds(10))