diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 688a126..7dbb9fd 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -170,7 +170,10 @@ public final class LanguageModelSession: @unchecked Sendable { ) ) session.withMutation(keyPath: \.transcript) { - session.state.withLock { $0.transcript.append(responseEntry) } + session.state.withLock { + $0.transcript.append(contentsOf: lastSnapshot.transcriptEntries) + $0.transcript.append(responseEntry) + } } } } catch { @@ -834,13 +837,23 @@ extension LanguageModelSession { public var content: Content.PartiallyGenerated public var rawContent: GeneratedContent + /// Transcript entries (tool calls and outputs) produced so far while streaming. + /// Cumulative across tool rounds; empty for providers that don't stream tool activity. + public var transcriptEntries: ArraySlice + /// Creates a snapshot from partially generated content and raw content. /// - Parameters: /// - content: The partially generated content. /// - rawContent: The raw content produced by the model. - public init(content: Content.PartiallyGenerated, rawContent: GeneratedContent) { + /// - transcriptEntries: Transcript entries accumulated so far (tool calls/outputs). + public init( + content: Content.PartiallyGenerated, + rawContent: GeneratedContent, + transcriptEntries: ArraySlice = [] + ) { self.content = content self.rawContent = rawContent + self.transcriptEntries = transcriptEntries } } } @@ -903,7 +916,7 @@ extension LanguageModelSession.ResponseStream: AsyncSequence { return LanguageModelSession.Response( content: finalContent, rawContent: last.rawContent, - transcriptEntries: [] + transcriptEntries: last.transcriptEntries ) } } @@ -918,7 +931,7 @@ extension LanguageModelSession.ResponseStream: AsyncSequence { return LanguageModelSession.Response( content: finalContent, rawContent: fallbackSnapshot.rawContent, - transcriptEntries: [] + transcriptEntries: fallbackSnapshot.transcriptEntries ) } diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 8b8bb0b..9af2202 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -208,6 +208,11 @@ import Foundation /// A negative value offloads all layers, and `0` runs entirely on the CPU. public let gpuLayers: Int32 + /// The path to the multimodal projector GGUF file, when the model has one. + /// + /// Prompts may include image segments only when a projector is loaded. + public let mmprojPath: String? + /// The default GPU layer count for the current platform. /// /// All layers are offloaded by default: the prebuilt llama.cpp binaries @@ -453,18 +458,202 @@ import Foundation /// The model's vocabulary private var vocab: OpaquePointer? + /// The multimodal projector context, when a projector file was provided + private var mtmdContext: OpaquePointer? + /// Whether the model is currently loaded private var isModelLoaded: Bool = false + /// A context kept alive for one session so exchanges reuse its state. + private struct CachedSessionContext { + let sessionID: ObjectIdentifier + let context: OpaquePointer + var tokens: [llama_token] + let contextSize: UInt32 + let batchSize: UInt32 + } + + private var cachedSessionContext: CachedSessionContext? + + /// The number of prompt tokens reused from the cached context by the most + /// recent chat generation. + internal private(set) var lastReusedTokenCount: Int = 0 + + /// The number of prompt tokens decoded by the most recent chat generation. + internal private(set) var lastPrefillTokenCount: Int = 0 + + /// Frees the cached per-session context and the state it holds. + /// + /// The next chat generation prefills its full prompt again. Call this under + /// memory pressure or when a session is discarded. + public func clearCachedContext() { + discardCachedSessionContext() + } + + private func discardCachedSessionContext() { + if let cached = cachedSessionContext { + llama_free(cached.context) + } + cachedSessionContext = nil + } + + private func recordCachedTokens(_ tokens: [llama_token], context: OpaquePointer) { + guard var cached = cachedSessionContext, cached.context == context else { return } + cached.tokens = tokens + cachedSessionContext = cached + } + + /// Returns a context for the session along with the index of the first + /// prompt token that still needs to be decoded. + /// + /// A cached context whose recorded tokens share a prefix with the prompt + /// keeps that prefix: matching state past the divergence point is removed + /// with `llama_memory_seq_rm`, and backends whose state cannot be rewound + /// (recurrent models) fall back to clearing the memory and decoding the + /// full prompt. The final prompt token is always re-decoded so sampling + /// has fresh logits. + private func acquireSessionContext( + for session: LanguageModelSession, + promptTokens: [llama_token], + options: ResolvedGenerationOptions + ) throws -> (context: OpaquePointer, startIndex: Int) { + let sessionID = ObjectIdentifier(session) + + if var cached = cachedSessionContext, + cached.sessionID == sessionID, + cached.contextSize == options.contextSize, + cached.batchSize == options.batchSize + { + var common = 0 + while common < cached.tokens.count, common < promptTokens.count, + cached.tokens[common] == promptTokens[common] + { + common += 1 + } + if common == promptTokens.count { + common = max(0, promptTokens.count - 1) + } + if common < cached.tokens.count { + let memory = llama_get_memory(cached.context) + if !llama_memory_seq_rm(memory, 0, llama_pos(common), -1) { + llama_memory_clear(memory, true) + common = 0 + } + } + cached.tokens = Array(promptTokens.prefix(common)) + cachedSessionContext = cached + return (cached.context, common) + } + + discardCachedSessionContext() + let contextParams = createContextParams(from: options) + guard let context = llama_init_from_model(model!, contextParams) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + guard llama_get_memory(context) != nil else { + llama_free(context) + throw LlamaLanguageModelError.encoderOnlyModel + } + cachedSessionContext = CachedSessionContext( + sessionID: sessionID, + context: context, + tokens: [], + contextSize: options.contextSize, + batchSize: options.batchSize + ) + return (context, 0) + } + + /// Creates a single-use context for generations that do not reuse state. + private func makeFreshContext(options: ResolvedGenerationOptions) throws -> OpaquePointer { + let contextParams = createContextParams(from: options) + guard let context = llama_init_from_model(model!, contextParams) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + guard llama_get_memory(context) != nil else { + llama_free(context) + throw LlamaLanguageModelError.encoderOnlyModel + } + llama_set_causal_attn(context, true) + llama_set_n_threads(context, options.threads, options.threads) + return context + } + + /// Runs a chat text generation for the session, reusing the session's + /// cached context when its state matches a prefix of the prompt. + private func generateChatText( + session: LanguageModelSession, + prompt: String, + maxTokens: Int, + options: ResolvedGenerationOptions, + onToken: (String) -> Bool + ) throws { + guard let model = self.model, let vocab = llama_model_get_vocab(model) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + + let promptTokens = try tokenizeText(vocab: vocab, text: prompt) + guard !promptTokens.isEmpty else { + throw LlamaLanguageModelError.tokenizationFailed + } + + if llama_model_has_encoder(model) { + let context = try makeFreshContext(options: options) + defer { llama_free(context) } + try performTokenGeneration( + context: context, + vocab: vocab, + promptTokens: promptTokens, + startIndex: 0, + maxTokens: maxTokens, + options: options, + onToken: onToken + ) + return + } + + let (context, startIndex) = try acquireSessionContext( + for: session, + promptTokens: promptTokens, + options: options + ) + llama_set_causal_attn(context, true) + llama_set_n_threads(context, options.threads, options.threads) + + do { + try performTokenGeneration( + context: context, + vocab: vocab, + promptTokens: promptTokens, + startIndex: startIndex, + maxTokens: maxTokens, + options: options, + onToken: onToken + ) + } catch { + discardCachedSessionContext() + throw error + } + } + /// Creates a Llama language model. /// /// - Parameters: /// - modelPath: The path to the GGUF model file. /// - gpuLayers: The number of model layers to offload to the GPU. /// Defaults to ``defaultGPULayerCount``. - public init(modelPath: String, gpuLayers: Int32 = LlamaLanguageModel.defaultGPULayerCount) { + /// - mmprojPath: The path to a multimodal projector GGUF file matching + /// the model. When provided, prompts may include image segments, + /// which are encoded through the projector. Defaults to `nil` + /// (text only). + public init( + modelPath: String, + gpuLayers: Int32 = LlamaLanguageModel.defaultGPULayerCount, + mmprojPath: String? = nil + ) { self.modelPath = modelPath self.gpuLayers = gpuLayers + self.mmprojPath = mmprojPath self.legacyDefaults = ResolvedGenerationOptions() } @@ -505,11 +694,192 @@ import Foundation } deinit { + if let cached = cachedSessionContext { + llama_free(cached.context) + } + if let mtmdContext = mtmdContext { + mtmd_free(mtmdContext) + } if let model = model { llama_model_free(model) } } + // MARK: - Tool calling + + /// Prompt-side tool state for one exchange: the detected syntax, the + /// session's tool definitions, and the tool turns produced so far in + /// the current resolve-and-continue loop. + struct LlamaToolPromptContext { + let format: LlamaToolCallFormat + let definitions: [LlamaToolDefinition] + var pendingEntries: [Transcript.Entry] + } + + private struct ToolInvocationResult { + let call: Transcript.ToolCall + let output: Transcript.ToolOutput + } + + private enum ToolResolutionOutcome { + case stop(calls: [Transcript.ToolCall]) + case invocations([ToolInvocationResult]) + } + + private static func maxToolIterationsExceededError(limit: Int) -> LanguageModelSession.GenerationError { + .decodingFailure( + .init( + debugDescription: + "Exceeded maximum tool iterations (\(limit)) while processing Llama tool calls." + ) + ) + } + + private static func repeatedToolCallLoopError() -> LanguageModelSession.GenerationError { + .decodingFailure( + .init( + debugDescription: + "Detected repeated Llama tool-call signature and aborted to avoid an infinite tool loop." + ) + ) + } + + private func currentToolCallFormat() -> LlamaToolCallFormat { + guard let model = self.model else { return .hermesJSON } + let template = llama_model_chat_template(model, nil).map { String(cString: $0) } + return LlamaToolCallFormat.detect(template: template) + } + + private func makeToolPromptContext(for session: LanguageModelSession) throws -> LlamaToolPromptContext? { + guard !session.tools.isEmpty, self.model != nil else { return nil } + let format = currentToolCallFormat() + let definitions = try session.tools.map { tool -> LlamaToolDefinition in + let schema = tool.parameters.withResolvedRoot() ?? tool.parameters + let data = try JSONEncoder().encode(schema) + let parameters = try JSONSerialization.jsonObject(with: data) as? [String: Any] + return LlamaToolDefinition( + name: tool.name, + description: tool.description, + parameters: parameters + ) + } + return LlamaToolPromptContext(format: format, definitions: definitions, pendingEntries: []) + } + + private func toolOutputText(_ output: Transcript.ToolOutput) -> String { + var parts: [String] = [] + for segment in output.segments { + switch segment { + case .text(let text): + parts.append(text.content) + case .structure(let structure): + parts.append(structure.content.jsonString) + case .image: + break + } + } + return parts.joined(separator: "\n") + } + + private func makeTranscriptToolCalls( + from parsedCalls: [LlamaParsedToolCall] + ) throws -> [Transcript.ToolCall] { + try parsedCalls.map { parsed in + Transcript.ToolCall( + id: UUID().uuidString, + toolName: parsed.name, + arguments: try GeneratedContent(json: parsed.argumentsJSON) + ) + } + } + + private func resolveToolCalls( + _ parsedCalls: [LlamaParsedToolCall], + session: LanguageModelSession + ) async throws -> ToolResolutionOutcome { + if parsedCalls.isEmpty { return .invocations([]) } + + var toolsByName: [String: any Tool] = [:] + for tool in session.tools where toolsByName[tool.name] == nil { + toolsByName[tool.name] = tool + } + + let transcriptCalls = try makeTranscriptToolCalls(from: parsedCalls) + + if let delegate = session.toolExecutionDelegate { + await delegate.didGenerateToolCalls(transcriptCalls, in: session) + } + + 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: + 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) + } + } + } + + return .invocations(results) + } + public func respond( within session: LanguageModelSession, to prompt: Prompt, @@ -517,59 +887,157 @@ import Foundation includeSchemaInPrompt: Bool, options: GenerationOptions ) async throws -> LanguageModelSession.Response where Content: Generable { - // Validate that no image segments are present - try validateNoImageSegments(in: session) + if mmprojPath == nil { + try validateNoImageSegments(in: session) + } try await ensureModelLoaded() let runtimeOptions = resolvedOptions(from: options) let structuredOptions = resolvedStructuredOptions(from: options) - let contextParams = createContextParams(from: runtimeOptions) - // Try to create context with error handling - guard let context = llama_init_from_model(model!, contextParams) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - defer { llama_free(context) } - - // Check if this is an embedding model (no KV cache). - // This early check catches models configured for embeddings that lack a KV cache. - // A complementary architectural check in prepareInitialBatch catches encoder-only - // models (like BERT) by their architecture type. - if llama_get_memory(context) == nil { - throw LlamaLanguageModelError.encoderOnlyModel - } - - llama_set_causal_attn(context, true) - llama_set_warmup(context, false) - llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) - - let fullPrompt: String - if includeSchemaInPrompt, type != String.self { - fullPrompt = try formatPrompt( - for: session, - extraSystemMessage: schemaPrompt(for: type.generationSchema), - assistantPrefill: runtimeOptions.assistantPrefill - ) - } else { - fullPrompt = try formatPrompt(for: session, assistantPrefill: runtimeOptions.assistantPrefill) - } + let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil if type == String.self { let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - let text = try await generateText( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) + let outputFormat = currentToolCallFormat() + var toolContext = try makeToolPromptContext(for: session) + let maxToolIterations = 8 + var toolIteration = 0 + var previousToolCallSignature: String? + var allEntries: [Transcript.Entry] = [] + var text = "" + + generationLoop: while true { + var promptImages: [Data] = [] + let fullPrompt = try formatPrompt( + for: session, + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages, + toolContext: toolContext + ) + + var accumulated = "" + let terminator = toolContext?.format.callTerminator + let collectToken: (String) -> Bool = { tokenText in + accumulated += tokenText + if let terminator, + accumulated.suffix(terminator.count + 8).contains(terminator) + { + return false + } + return true + } + + if promptImages.isEmpty { + try generateChatText( + session: session, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: collectToken + ) + } else { + discardCachedSessionContext() + let context = try makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } + try performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: collectToken + ) + } + + guard let format = toolContext?.format else { + if outputFormat == .gemma { + text = LlamaToolCallFormat.stripGemmaThoughtChannels(from: accumulated) + .trimmingCharacters(in: .whitespacesAndNewlines) + } else { + text = accumulated + } + break generationLoop + } + let (visibleText, parsedCalls) = format.parseToolCalls(in: accumulated) + if parsedCalls.isEmpty { + text = visibleText + break generationLoop + } + + toolIteration += 1 + if toolIteration > maxToolIterations { + let unresolved = try makeTranscriptToolCalls(from: parsedCalls) + allEntries.append(.toolCalls(Transcript.ToolCalls(unresolved))) + throw Self.maxToolIterationsExceededError(limit: maxToolIterations) + } + let signature = + parsedCalls + .map { "\($0.name):\($0.argumentsJSON)" } + .joined(separator: "|") + if signature == previousToolCallSignature { + let unresolved = try makeTranscriptToolCalls(from: parsedCalls) + allEntries.append(.toolCalls(Transcript.ToolCalls(unresolved))) + throw Self.repeatedToolCallLoopError() + } + previousToolCallSignature = signature + + let resolution = try await resolveToolCalls(parsedCalls, session: session) + switch resolution { + case .stop(let calls): + if !calls.isEmpty { + allEntries.append(.toolCalls(Transcript.ToolCalls(calls))) + } + return LanguageModelSession.Response( + content: "" as! Content, + rawContent: GeneratedContent(""), + transcriptEntries: ArraySlice(allEntries) + ) + case .invocations(let invocations): + guard !invocations.isEmpty else { + text = visibleText + break generationLoop + } + let callsEntry = Transcript.Entry.toolCalls( + Transcript.ToolCalls(invocations.map(\.call)) + ) + allEntries.append(callsEntry) + toolContext?.pendingEntries.append(callsEntry) + for invocation in invocations { + let outputEntry = Transcript.Entry.toolOutput(invocation.output) + allEntries.append(outputEntry) + toolContext?.pendingEntries.append(outputEntry) + } + } + } return LanguageModelSession.Response( content: text as! Content, rawContent: GeneratedContent(text), - transcriptEntries: ArraySlice([]) + transcriptEntries: ArraySlice(allEntries) ) } else { + var promptImages: [Data] = [] + let fullPrompt: String + if includeSchemaInPrompt { + fullPrompt = try formatPrompt( + for: session, + extraSystemMessage: schemaPrompt(for: type.generationSchema), + assistantPrefill: runtimeOptions.assistantPrefill + ) + } else { + fullPrompt = try formatPrompt( + for: session, + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages + ) + } + let context = try makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } let maxTokens = structuredOptions.maximumResponseTokens ?? 512 let jsonString = try await generateStructuredJSON( context: context, @@ -600,15 +1068,16 @@ import Foundation fatalError("LlamaLanguageModel only supports generating String content") } - // Validate that no image segments are present - do { - try validateNoImageSegments(in: session) - } catch { - return LanguageModelSession.ResponseStream( - stream: AsyncThrowingStream { continuation in - continuation.finish(throwing: error) - } - ) + if mmprojPath == nil { + do { + try validateNoImageSegments(in: session) + } catch { + return LanguageModelSession.ResponseStream( + stream: AsyncThrowingStream { continuation in + continuation.finish(throwing: error) + } + ) + } } let stream: AsyncThrowingStream.Snapshot, any Error> = @@ -619,52 +1088,148 @@ import Foundation let runtimeOptions = resolvedOptions(from: options) let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - let contextParams = createContextParams(from: runtimeOptions) - guard let context = llama_init_from_model(model!, contextParams) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - defer { llama_free(context) } - - // Check if this is an embedding model (no KV cache). - // This early check catches models configured for embeddings that lack a KV cache. - // A complementary architectural check in prepareInitialBatch catches encoder-only - // models (like BERT) by their architecture type. - if llama_get_memory(context) == nil { - throw LlamaLanguageModelError.encoderOnlyModel + let outputFormat = self.currentToolCallFormat() + var toolContext = try self.makeToolPromptContext(for: session) + let maxToolIterations = 8 + var toolIteration = 0 + var previousToolCallSignature: String? + var accumulatedEntries: [Transcript.Entry] = [] + var emittedBase = "" + var lastYieldedText: String? + let imageMarker = + self.mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil + + func yieldSnapshot(_ text: String) { + lastYieldedText = text + let snapshot = LanguageModelSession.ResponseStream.Snapshot( + content: (text as! Content).asPartiallyGenerated(), + rawContent: GeneratedContent(text), + transcriptEntries: ArraySlice(accumulatedEntries) + ) + continuation.yield(snapshot) } - // Stabilize runtime behavior per-context - llama_set_causal_attn(context, true) - llama_set_warmup(context, false) - llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) - - var accumulatedText = "" - let fullPrompt = try self.formatPrompt( - for: session, - assistantPrefill: runtimeOptions.assistantPrefill - ) - - do { - for try await tokenText in generateTextStream( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) { - accumulatedText += tokenText - - let snapshot = LanguageModelSession.ResponseStream.Snapshot( - content: (accumulatedText as! Content).asPartiallyGenerated(), - rawContent: GeneratedContent(accumulatedText) + generationLoop: while true { + var promptImages: [Data] = [] + let fullPrompt = try self.formatPrompt( + for: session, + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages, + toolContext: toolContext + ) + + var roundRaw = "" + let terminator = toolContext?.format.callTerminator + let withholdToolCalls = toolContext != nil + let collectToken: (String) -> Bool = { tokenText in + roundRaw += tokenText + let visible = outputFormat.streamingVisibleText( + in: roundRaw, + withholdToolCalls: withholdToolCalls + ) + let total = emittedBase + visible + if !total.isEmpty, total != lastYieldedText { + yieldSnapshot(total) + } + if let terminator, + roundRaw.suffix(terminator.count + 8).contains(terminator) + { + return false + } + return true + } + + if promptImages.isEmpty { + try self.generateChatText( + session: session, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: collectToken + ) + } else { + self.discardCachedSessionContext() + let context = try self.makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } + try self.performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: collectToken + ) + } + + if Task.isCancelled { + break generationLoop + } + + let roundVisible = outputFormat.streamingVisibleText( + in: roundRaw, + withholdToolCalls: withholdToolCalls + ) + + guard let format = toolContext?.format else { + emittedBase += roundVisible + break generationLoop + } + let (_, parsedCalls) = format.parseToolCalls(in: roundRaw) + if parsedCalls.isEmpty { + emittedBase += roundVisible + break generationLoop + } + + toolIteration += 1 + if toolIteration > maxToolIterations { + let unresolved = try self.makeTranscriptToolCalls(from: parsedCalls) + accumulatedEntries.append(.toolCalls(Transcript.ToolCalls(unresolved))) + throw Self.maxToolIterationsExceededError(limit: maxToolIterations) + } + let signature = + parsedCalls + .map { "\($0.name):\($0.argumentsJSON)" } + .joined(separator: "|") + if signature == previousToolCallSignature { + let unresolved = try self.makeTranscriptToolCalls(from: parsedCalls) + accumulatedEntries.append(.toolCalls(Transcript.ToolCalls(unresolved))) + throw Self.repeatedToolCallLoopError() + } + previousToolCallSignature = signature + + let resolution = try await self.resolveToolCalls(parsedCalls, session: session) + switch resolution { + case .stop(let calls): + if !calls.isEmpty { + accumulatedEntries.append(.toolCalls(Transcript.ToolCalls(calls))) + yieldSnapshot(emittedBase + roundVisible) + } + break generationLoop + case .invocations(let invocations): + guard !invocations.isEmpty else { + emittedBase += roundVisible + break generationLoop + } + let callsEntry = Transcript.Entry.toolCalls( + Transcript.ToolCalls(invocations.map(\.call)) ) - continuation.yield(snapshot) + accumulatedEntries.append(callsEntry) + toolContext?.pendingEntries.append(callsEntry) + for invocation in invocations { + let outputEntry = Transcript.Entry.toolOutput(invocation.output) + accumulatedEntries.append(outputEntry) + toolContext?.pendingEntries.append(outputEntry) + } + emittedBase += roundVisible + yieldSnapshot(emittedBase) } - } catch { - continuation.finish(throwing: error) - return } + if emittedBase != lastYieldedText { + yieldSnapshot(emittedBase) + } continuation.finish() } catch { continuation.finish(throwing: error) @@ -693,6 +1258,11 @@ import Foundation llama_backend_init() // Free any existing model before loading a new one + discardCachedSessionContext() + if let existingContext = mtmdContext { + mtmd_free(existingContext) + self.mtmdContext = nil + } if let existingModel = model { llama_model_free(existingModel) self.model = nil @@ -703,6 +1273,22 @@ import Foundation throw LlamaLanguageModelError.modelLoadFailed } + if let mmprojPath { + guard FileManager.default.fileExists(atPath: mmprojPath) else { + llama_model_free(loadedModel) + throw LlamaLanguageModelError.invalidModelPath + } + var mtmdParams = mtmd_context_params_default() + mtmdParams.use_gpu = gpuLayers != 0 + mtmdParams.print_timings = false + mtmdParams.n_threads = Int32(ProcessInfo.processInfo.processorCount) + guard let projector = mtmd_init_from_file(mmprojPath, loadedModel, mtmdParams) else { + llama_model_free(loadedModel) + throw LlamaLanguageModelError.modelLoadFailed + } + self.mtmdContext = projector + } + self.model = loadedModel self.vocab = llama_model_get_vocab(loadedModel) self.isModelLoaded = true @@ -793,133 +1379,29 @@ import Foundation llama_sampler_chain_add(sampler, llama_sampler_init_top_p(1.0, 1)) llama_sampler_chain_add(sampler, llama_sampler_init_greedy()) case .topK(let k, let seed): - llama_sampler_chain_add(sampler, llama_sampler_init_top_k(Int32(k))) - llama_sampler_chain_add(sampler, llama_sampler_init_top_p(1.0, 1)) - llama_sampler_chain_add(sampler, llama_sampler_init_temp(effectiveTemperature)) - let samplingSeed = seed.map(UInt32.init) ?? options.seed - llama_sampler_chain_add(sampler, llama_sampler_init_dist(samplingSeed)) - case .nucleus(let threshold, let seed): - llama_sampler_chain_add(sampler, llama_sampler_init_top_k(0)) - llama_sampler_chain_add(sampler, llama_sampler_init_top_p(Float(threshold), 1)) - llama_sampler_chain_add(sampler, llama_sampler_init_temp(effectiveTemperature)) - let samplingSeed = seed.map(UInt32.init) ?? options.seed - llama_sampler_chain_add(sampler, llama_sampler_init_dist(samplingSeed)) - } - return - } - - if options.topK > 0 { - llama_sampler_chain_add(sampler, llama_sampler_init_top_k(options.topK)) - } - if options.topP < 1.0 { - llama_sampler_chain_add(sampler, llama_sampler_init_top_p(options.topP, 1)) - } - llama_sampler_chain_add(sampler, llama_sampler_init_temp(effectiveTemperature)) - llama_sampler_chain_add(sampler, llama_sampler_init_dist(options.seed)) - } - - private func generateText( - context: OpaquePointer, - model: OpaquePointer, - prompt: String, - maxTokens: Int, - options: ResolvedGenerationOptions - ) async throws - -> String - { - guard let vocab = llama_model_get_vocab(model) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - - // Tokenize the prompt - let promptTokens = try tokenizeText(vocab: vocab, text: prompt) - guard !promptTokens.isEmpty else { - throw LlamaLanguageModelError.tokenizationFailed - } - - var batch = llama_batch_init(Int32(options.batchSize), 0, 1) - defer { llama_batch_free(batch) } - - let hasEncoder = try prepareInitialBatch( - batch: &batch, - promptTokens: promptTokens, - model: model, - vocab: vocab, - context: context, - batchSize: options.batchSize, - contextSize: options.contextSize - ) - - // Initialize sampler chain with options - guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { - throw LlamaLanguageModelError.decodingFailed - } - defer { llama_sampler_free(sampler) } - let samplerPtr = UnsafeMutablePointer(sampler) - - let effectiveTemperature = Float(options.temperature) - - // Apply repeat/frequency/presence penalties from custom options - let effectiveRepeatPenalty = options.repeatPenalty - let effectiveRepeatLastN = options.repeatLastN - let effectiveFrequencyPenalty = options.frequencyPenalty - let effectivePresencePenalty = options.presencePenalty - - if effectiveRepeatPenalty != 1.0 || effectiveFrequencyPenalty != 0.0 || effectivePresencePenalty != 0.0 { - llama_sampler_chain_add( - samplerPtr, - llama_sampler_init_penalties( - llama_vocab_n_tokens(vocab), - effectiveRepeatLastN, - effectiveRepeatPenalty, - effectiveFrequencyPenalty, - effectivePresencePenalty - ) - ) - } - - applySampling(sampler: samplerPtr, effectiveTemperature: effectiveTemperature, options: options) - - // Generate tokens one by one - var generatedText = "" - // Track position - for encoder-decoder models, we start from position 1 (after decoder start token) - // For decoder-only models, we continue from the end of the prompt - var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) - - for _ in 0 ..< maxTokens { - // Sample next token from logits - llama_batch_get_one creates batch with single token at index 0 - let nextToken = llama_sampler_sample(sampler, context, batch.n_tokens - 1) - llama_sampler_accept(sampler, nextToken) - - // Check for end of sequence - if llama_vocab_is_eog(vocab, nextToken) { - break - } - - // Convert token to text - if let tokenText = tokenToText(vocab: vocab, token: nextToken) { - generatedText += tokenText - } - - // Prepare batch for next token - batch.n_tokens = 1 - batch.token[0] = nextToken - batch.pos[0] = n_cur - batch.n_seq_id[0] = 1 - if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { - seq_id[0] = 0 - } - batch.logits[0] = 1 - - n_cur += 1 - - let decodeResult = llama_decode(context, batch) - guard decodeResult == 0 else { - break + llama_sampler_chain_add(sampler, llama_sampler_init_top_k(Int32(k))) + llama_sampler_chain_add(sampler, llama_sampler_init_top_p(1.0, 1)) + llama_sampler_chain_add(sampler, llama_sampler_init_temp(effectiveTemperature)) + let samplingSeed = seed.map(UInt32.init) ?? options.seed + llama_sampler_chain_add(sampler, llama_sampler_init_dist(samplingSeed)) + case .nucleus(let threshold, let seed): + llama_sampler_chain_add(sampler, llama_sampler_init_top_k(0)) + llama_sampler_chain_add(sampler, llama_sampler_init_top_p(Float(threshold), 1)) + llama_sampler_chain_add(sampler, llama_sampler_init_temp(effectiveTemperature)) + let samplingSeed = seed.map(UInt32.init) ?? options.seed + llama_sampler_chain_add(sampler, llama_sampler_init_dist(samplingSeed)) } + return } - return generatedText + if options.topK > 0 { + llama_sampler_chain_add(sampler, llama_sampler_init_top_k(options.topK)) + } + if options.topP < 1.0 { + llama_sampler_chain_add(sampler, llama_sampler_init_top_p(options.topP, 1)) + } + llama_sampler_chain_add(sampler, llama_sampler_init_temp(effectiveTemperature)) + llama_sampler_chain_add(sampler, llama_sampler_init_dist(options.seed)) } /// Builds a JSONSchema-informed prompt for structured output. @@ -1186,133 +1668,254 @@ import Foundation } } - private func generateTextStream( + private func performTokenGeneration( context: OpaquePointer, - model: OpaquePointer, - prompt: String, + vocab: OpaquePointer, + promptTokens: [llama_token], + startIndex: Int, maxTokens: Int, - options: ResolvedGenerationOptions - ) -> AsyncThrowingStream { - return AsyncThrowingStream { continuation in - self.performTextGeneration( - context: context, - model: model, - prompt: prompt, - maxTokens: maxTokens, - options: options, - continuation: continuation + options: ResolvedGenerationOptions, + onToken: (String) -> Bool + ) throws { + guard let model = self.model else { + throw LlamaLanguageModelError.modelLoadFailed + } + + lastReusedTokenCount = startIndex + lastPrefillTokenCount = promptTokens.count - startIndex + + // Initialize batch + var batch = llama_batch_init(Int32(options.batchSize), 0, 1) + defer { llama_batch_free(batch) } + + let hasEncoder = try prepareInitialBatch( + batch: &batch, + promptTokens: promptTokens, + model: model, + vocab: vocab, + context: context, + batchSize: options.batchSize, + contextSize: options.contextSize, + startIndex: startIndex + ) + + // Initialize sampler chain with options + guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { + throw LlamaLanguageModelError.decodingFailed + } + defer { llama_sampler_free(sampler) } + let samplerPtr = UnsafeMutablePointer(sampler) + + let effectiveTemperature = Float(options.temperature) + + // Apply repeat/frequency/presence penalties from custom options + let effectiveRepeatPenalty = options.repeatPenalty + let effectiveRepeatLastN = options.repeatLastN + let effectiveFrequencyPenalty = options.frequencyPenalty + let effectivePresencePenalty = options.presencePenalty + + if effectiveRepeatPenalty != 1.0 || effectiveFrequencyPenalty != 0.0 || effectivePresencePenalty != 0.0 { + llama_sampler_chain_add( + samplerPtr, + llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), + effectiveRepeatLastN, + effectiveRepeatPenalty, + effectiveFrequencyPenalty, + effectivePresencePenalty + ) ) } - } - private func performTextGeneration( - context: OpaquePointer, - model: OpaquePointer, - prompt: String, - maxTokens: Int, - options: ResolvedGenerationOptions, - continuation: AsyncThrowingStream.Continuation - ) { - do { - guard let vocab = llama_model_get_vocab(model) else { - continuation.finish(throwing: LlamaLanguageModelError.contextInitializationFailed) - return + // Check for mirostat sampling (takes precedence over standard sampling) + applySampling(sampler: samplerPtr, effectiveTemperature: effectiveTemperature, options: options) + + // Generate tokens one by one + // Track position - for encoder-decoder models, we start from position 1 (after decoder start token) + // For decoder-only models, we continue from the end of the prompt + var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) + var decodedTokens = promptTokens + + for _ in 0 ..< maxTokens { + if Task.isCancelled { + break } - // Tokenize the prompt - let promptTokens = try tokenizeText(vocab: vocab, text: prompt) - guard !promptTokens.isEmpty else { - continuation.finish(throwing: LlamaLanguageModelError.tokenizationFailed) - return + // Sample next token from logits of the last token we just decoded + let nextToken = llama_sampler_sample(sampler, context, batch.n_tokens - 1) + llama_sampler_accept(sampler, nextToken) + + // Check for end of sequence + if llama_vocab_is_eog(vocab, nextToken) { + break } - // Initialize batch - var batch = llama_batch_init(Int32(options.batchSize), 0, 1) - defer { llama_batch_free(batch) } + // Convert token to text and yield it + if let tokenText = tokenToText(vocab: vocab, token: nextToken) { + guard onToken(tokenText) else { + break + } + } - let hasEncoder = try prepareInitialBatch( - batch: &batch, - promptTokens: promptTokens, - model: model, - vocab: vocab, - context: context, - batchSize: options.batchSize, - contextSize: options.contextSize - ) + // Prepare batch for next token + batch.n_tokens = 1 + batch.token[0] = nextToken + batch.pos[0] = n_cur + batch.n_seq_id[0] = 1 + if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { + seq_id[0] = 0 + } + batch.logits[0] = 1 - // Initialize sampler chain with options - guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { - throw LlamaLanguageModelError.decodingFailed + n_cur += 1 + + let decodeResult = llama_decode(context, batch) + guard decodeResult == 0 else { + break } - defer { llama_sampler_free(sampler) } - let samplerPtr = UnsafeMutablePointer(sampler) + decodedTokens.append(nextToken) + } - let effectiveTemperature = Float(options.temperature) + recordCachedTokens(decodedTokens, context: context) + } - // Apply repeat/frequency/presence penalties from custom options - let effectiveRepeatPenalty = options.repeatPenalty - let effectiveRepeatLastN = options.repeatLastN - let effectiveFrequencyPenalty = options.frequencyPenalty - let effectivePresencePenalty = options.presencePenalty + /// Evaluates a marker-annotated multimodal prompt through the projector, + /// then generates text tokens from the resulting state. + private func performMultimodalGeneration( + context: OpaquePointer, + prompt: String, + images: [Data], + maxTokens: Int, + options: ResolvedGenerationOptions, + onToken: (String) -> Bool + ) throws { + guard let mtmdContext, let model = self.model, + let vocab = llama_model_get_vocab(model) + else { + throw LlamaLanguageModelError.contextInitializationFailed + } - if effectiveRepeatPenalty != 1.0 || effectiveFrequencyPenalty != 0.0 || effectivePresencePenalty != 0.0 - { - llama_sampler_chain_add( - samplerPtr, - llama_sampler_init_penalties( - llama_vocab_n_tokens(vocab), - effectiveRepeatLastN, - effectiveRepeatPenalty, - effectiveFrequencyPenalty, - effectivePresencePenalty - ) + var bitmaps: [OpaquePointer?] = [] + defer { + for bitmap in bitmaps { + if let bitmap { + mtmd_bitmap_free(bitmap) + } + } + } + for imageData in images { + let wrapper = imageData.withUnsafeBytes { raw -> mtmd_helper_bitmap_wrapper in + mtmd_helper_bitmap_init_from_buf( + mtmdContext, + raw.bindMemory(to: UInt8.self).baseAddress, + imageData.count, + false ) } + if let videoContext = wrapper.video_ctx { + mtmd_helper_video_free(videoContext) + throw LlamaLanguageModelError.unsupportedFeature + } + guard let bitmap = wrapper.bitmap else { + throw LlamaLanguageModelError.encodingFailed + } + bitmaps.append(bitmap) + } - // Check for mirostat sampling (takes precedence over standard sampling) - applySampling(sampler: samplerPtr, effectiveTemperature: effectiveTemperature, options: options) + guard let chunks = mtmd_input_chunks_init() else { + throw LlamaLanguageModelError.encodingFailed + } + defer { mtmd_input_chunks_free(chunks) } + + let tokenizeResult = prompt.withCString { cPrompt -> Int32 in + var inputText = mtmd_input_text( + text: cPrompt, + text_len: strlen(cPrompt), + add_special: true, + parse_special: true + ) + return bitmaps.withUnsafeMutableBufferPointer { buffer in + mtmd_tokenize(mtmdContext, chunks, &inputText, buffer.baseAddress, buffer.count) + } + } + guard tokenizeResult == 0 else { + throw LlamaLanguageModelError.tokenizationFailed + } - // Generate tokens one by one - // Track position - for encoder-decoder models, we start from position 1 (after decoder start token) - // For decoder-only models, we continue from the end of the prompt - var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) + var pastPosition: llama_pos = 0 + let evalResult = mtmd_helper_eval_chunks( + mtmdContext, + context, + chunks, + 0, + 0, + Int32(options.batchSize), + true, + &pastPosition + ) + guard evalResult == 0 else { + throw LlamaLanguageModelError.decodingFailed + } - for _ in 0 ..< maxTokens { - // Sample next token from logits of the last token we just decoded - let nextToken = llama_sampler_sample(sampler, context, batch.n_tokens - 1) - llama_sampler_accept(sampler, nextToken) + guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { + throw LlamaLanguageModelError.decodingFailed + } + defer { llama_sampler_free(sampler) } + let samplerPtr = UnsafeMutablePointer(sampler) - // Check for end of sequence - if llama_vocab_is_eog(vocab, nextToken) { - break - } + if options.repeatPenalty != 1.0 || options.frequencyPenalty != 0.0 || options.presencePenalty != 0.0 { + llama_sampler_chain_add( + samplerPtr, + llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), + options.repeatLastN, + options.repeatPenalty, + options.frequencyPenalty, + options.presencePenalty + ) + ) + } + applySampling(sampler: samplerPtr, effectiveTemperature: options.temperature, options: options) - // Convert token to text and yield it - if let tokenText = tokenToText(vocab: vocab, token: nextToken) { - continuation.yield(tokenText) - } + var batch = llama_batch_init(1, 0, 1) + defer { llama_batch_free(batch) } - // Prepare batch for next token - batch.n_tokens = 1 - batch.token[0] = nextToken - batch.pos[0] = n_cur - batch.n_seq_id[0] = 1 - if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { - seq_id[0] = 0 - } - batch.logits[0] = 1 + var n_cur: Int32 = Int32(pastPosition) + var sampleIndex: Int32 = -1 + + for _ in 0 ..< maxTokens { + if Task.isCancelled { + break + } + + let nextToken = llama_sampler_sample(samplerPtr, context, sampleIndex) + llama_sampler_accept(samplerPtr, nextToken) - n_cur += 1 + if llama_vocab_is_eog(vocab, nextToken) { + break + } - let decodeResult = llama_decode(context, batch) - guard decodeResult == 0 else { + if let tokenText = tokenToText(vocab: vocab, token: nextToken) { + guard onToken(tokenText) else { break } } - continuation.finish() - } catch { - continuation.finish(throwing: error) + batch.n_tokens = 1 + batch.token[0] = nextToken + batch.pos[0] = n_cur + batch.n_seq_id[0] = 1 + if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { + seq_id[0] = 0 + } + batch.logits[0] = 1 + + n_cur += 1 + + guard llama_decode(context, batch) == 0 else { + break + } + sampleIndex = 0 } } @@ -1347,6 +1950,8 @@ import Foundation /// - context: The model context. /// - batchSize: The batch capacity per decode call. /// - contextSize: The context window the prompt must fit within. + /// - startIndex: The index of the first prompt token to decode. Earlier + /// tokens are already present in the context's state. Defaults to `0`. /// - Returns: `true` if the model has an encoder (for position tracking during generation). /// - Throws: `promptExceedsContextWindow` if the prompt cannot fit in the context window, /// `insufficientMemory` if an encoder prompt exceeds the batch capacity, `encoderOnlyModel` @@ -1358,7 +1963,8 @@ import Foundation vocab: OpaquePointer, context: OpaquePointer, batchSize: UInt32, - contextSize: UInt32 + contextSize: UInt32, + startIndex: Int = 0 ) throws -> Bool { // Leave at least one context cell free for generation. guard promptTokens.count < contextSize else { @@ -1369,7 +1975,7 @@ import Foundation let hasDecoder = llama_model_has_decoder(model) // Encoder models ingest the full prompt in a single llama_encode call. - guard !hasEncoder || promptTokens.count <= batchSize else { + guard !hasEncoder || (startIndex == 0 && promptTokens.count <= batchSize) else { throw LlamaLanguageModelError.insufficientMemory } @@ -1421,7 +2027,7 @@ import Foundation // batch-sized chunks with absolute positions, requesting logits // only for the final token. let capacity = Int(batchSize) - var start = 0 + var start = startIndex while start < promptTokens.count { let count = min(capacity, promptTokens.count - start) batch.n_tokens = Int32(count) @@ -1452,7 +2058,27 @@ import Foundation private func formatPrompt( for session: LanguageModelSession, extraSystemMessage: String? = nil, - assistantPrefill: String? = nil + assistantPrefill: String? = nil, + toolContext: LlamaToolPromptContext? = nil + ) throws -> String { + var images: [Data] = [] + return try formatPrompt( + for: session, + extraSystemMessage: extraSystemMessage, + assistantPrefill: assistantPrefill, + imageMarker: nil, + images: &images, + toolContext: toolContext + ) + } + + private func formatPrompt( + for session: LanguageModelSession, + extraSystemMessage: String?, + assistantPrefill: String?, + imageMarker: String?, + images: inout [Data], + toolContext: LlamaToolPromptContext? = nil ) throws -> String { guard let model = self.model else { throw LlamaLanguageModelError.modelLoadFailed @@ -1460,28 +2086,88 @@ import Foundation var messages: [(role: String, content: String)] = [] - for entry in session.transcript { + func appendEntry(_ entry: Transcript.Entry) throws { switch entry { case .instructions(let instructions): - let text = extractText(from: instructions.segments) + let text = try extractContent( + from: instructions.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("system", text)) } case .prompt(let prompt): - let text = extractText(from: prompt.segments) + let text = try extractContent( + from: prompt.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("user", text)) } case .response(let response): - let text = extractText(from: response.segments) + let text = try extractContent( + from: response.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("assistant", text)) } - default: - break + case .toolCalls(let toolCalls): + guard let toolContext else { break } + let parsed = toolCalls.map { + LlamaParsedToolCall(name: $0.toolName, argumentsJSON: $0.arguments.jsonString) + } + if let last = messages.last, last.role == "assistant" { + let markup = toolContext.format.assistantText(for: parsed, precededByContent: true) + messages[messages.count - 1].content += markup + } else { + let markup = toolContext.format.assistantText(for: parsed, precededByContent: false) + messages.append(("assistant", markup)) + } + + case .toolOutput(let output): + guard let toolContext else { break } + let message = toolContext.format.toolResponseMessage( + toolName: output.toolName, + content: toolOutputText(output) + ) + if let last = messages.last, last.role == message.role, last.role == "user", + last.content.hasSuffix("") + { + messages[messages.count - 1].content += "\n" + message.content + } else { + messages.append(message) + } + } + } + + for entry in session.transcript { + try appendEntry(entry) + } + if let toolContext { + for entry in toolContext.pendingEntries { + try appendEntry(entry) + } + } + + if let toolContext, !toolContext.definitions.isEmpty { + if let systemIndex = messages.firstIndex(where: { $0.role == "system" }) { + messages[systemIndex].content = toolContext.format.systemMessage( + existingText: messages[systemIndex].content, + tools: toolContext.definitions + ) + } else { + let systemText = toolContext.format.systemMessage( + existingText: "", + tools: toolContext.definitions + ) + messages.insert(("system", systemText), at: 0) } } @@ -1517,6 +2203,9 @@ import Foundation ) guard requiredSize > 0 else { + if let tmpl, String(cString: tmpl).contains("<|turn>") { + return renderGemma4Prompt(messages: messages, assistantPrefill: assistantPrefill) + } throw LlamaLanguageModelError.encodingFailed } @@ -1546,6 +2235,53 @@ import Foundation return rendered } + /// Renders the Gemma 4 canonical chat format, which + /// `llama_chat_apply_template` does not recognize: turns open with + /// `<|turn>role`, close with ``, and the assistant role is named + /// `model`. The BOS token is applied during tokenization. + private func renderGemma4Prompt( + messages: [(role: String, content: String)], + assistantPrefill: String? + ) -> String { + var rendered = "" + var openModelTurn = false + for (index, message) in messages.enumerated() { + if message.role == "tool" { + rendered += message.content + continue + } + let role = message.role == "assistant" ? "model" : message.role + let content = message.content.trimmingCharacters(in: .whitespacesAndNewlines) + if role == "model" && openModelTurn { + rendered += content + } else { + if openModelTurn { + rendered += "\n" + openModelTurn = false + } + rendered += "<|turn>\(role)\n\(content)" + } + if role == "model" { + let nextRole = index + 1 < messages.count ? messages[index + 1].role : nil + if nextRole == "tool" || nextRole == "assistant" { + openModelTurn = true + } else { + rendered += "\n" + openModelTurn = false + } + } else { + rendered += "\n" + } + } + if !openModelTurn { + rendered += "<|turn>model\n" + } + if let assistantPrefill, !assistantPrefill.isEmpty { + rendered += assistantPrefill + } + return rendered + } + private func extractText(from segments: [Transcript.Segment]) -> String { segments.compactMap { segment -> String? in if case .text(let t) = segment { return t.content } @@ -1553,6 +2289,42 @@ import Foundation }.joined() } + /// Extracts message content from segments, replacing each image segment + /// with `imageMarker` and collecting its payload in order. Image segments + /// throw ``LlamaLanguageModelError/unsupportedFeature`` when no marker is + /// provided. + private func extractContent( + from segments: [Transcript.Segment], + imageMarker: String?, + images: inout [Data] + ) throws -> String { + var parts: [String] = [] + for segment in segments { + switch segment { + case .text(let t): + parts.append(t.content) + case .image(let image): + guard let imageMarker else { + throw LlamaLanguageModelError.unsupportedFeature + } + switch image.source { + case .data(let data, _): + images.append(data) + parts.append(imageMarker) + case .url(let url): + guard url.isFileURL, let data = try? Data(contentsOf: url) else { + throw LlamaLanguageModelError.unsupportedFeature + } + images.append(data) + parts.append(imageMarker) + } + default: + break + } + } + return parts.joined() + } + private func tokenizeText(vocab: OpaquePointer, text: String) throws -> [llama_token] { let utf8Count = text.utf8.count let maxTokens = Int32(max(utf8Count * 2, 8)) // Rough estimate, minimum capacity diff --git a/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift b/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift new file mode 100644 index 0000000..837e3a5 --- /dev/null +++ b/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift @@ -0,0 +1,740 @@ +import Foundation + +/// The tool-calling syntax a GGUF model was trained on, detected from its +/// embedded chat template. +/// +/// `llama_chat_apply_template` renders chat messages but has no parameter for +/// tool definitions, so tool support is implemented at this layer: definitions +/// are rendered into the system prompt, past tool turns are replayed in the +/// model's native markup, and calls are parsed back out of generated text. +enum LlamaToolCallFormat: Sendable, Equatable { + /// Hermes-style JSON calls, used by Qwen 2.5/3 and many community fine-tunes: + /// `{"name": ..., "arguments": {...}}`. + case hermesJSON + + /// Qwen 3.5 XML calls: + /// `value`. + case qwenXML + + /// Gemma 4 canonical calls: + /// `<|tool_call>call:name{key:value}`, with `<|"|>`-quoted strings. + case gemma + + /// Detects the format from a model's embedded chat template text. + /// Unrecognized templates fall back to the Hermes JSON convention. + static func detect(template: String?) -> LlamaToolCallFormat { + guard let template else { return .hermesJSON } + if template.contains("<|turn>") { return .gemma } + if template.contains("" + case .gemma: return "" + } + } + + /// The marker that starts a tool-call block in generated text. + var callStartMarker: String { + switch self { + case .hermesJSON, .qwenXML: return "" + case .gemma: return "<|tool_call>" + } + } + + /// Markers that open a Gemma 4 thought channel. The canonical template + /// writes `<|channel>`, but deployed quantizations have been observed + /// emitting `<|channel|>`, so both spellings are recognized. + static let gemmaChannelOpenMarkers = ["<|channel>", "<|channel|>"] + + /// The marker that closes a Gemma 4 thought channel. + static let gemmaChannelCloseMarker = "" + + /// Removes Gemma 4 thought-channel spans from generated text. Thinking is + /// opt-in via `<|think|>`, but the model volunteers thought channels + /// anyway; the canonical template ships a `strip_thinking` macro for the + /// same reason. A span left unclosed at the end of the text is removed + /// through the end. + static func stripGemmaThoughtChannels(from text: String) -> String { + var result = "" + var remainder = Substring(text) + while let open = earliestRange(of: gemmaChannelOpenMarkers, in: remainder) { + result += remainder[.. Range? { + var earliest: Range? + for marker in markers { + if let range = text.range(of: marker), + earliest == nil || range.lowerBound < earliest!.lowerBound + { + earliest = range + } + } + return earliest + } + + /// The portion of partially generated text that is safe to show while + /// streaming: completed thought channels are removed (Gemma only), text + /// from a tool-call start onward is withheld when tools are active, and a + /// trailing partial match of either marker is held back until the next + /// token confirms or breaks it. + func streamingVisibleText(in raw: String, withholdToolCalls: Bool) -> String { + var text = raw + if self == .gemma { + text = Self.stripGemmaThoughtChannels(from: text) + } + if withholdToolCalls, let range = text.range(of: callStartMarker) { + text = String(text[.. 0 else { continue } + for length in stride(from: maxLength, through: 1, by: -1) + where text.hasSuffix(String(marker.prefix(length))) { + cut = max(cut, length) + break + } + } + if cut > 0 { + text.removeLast(cut) + } + return text + } +} + +/// A tool definition rendered into the system prompt. +struct LlamaToolDefinition { + let name: String + let description: String + let parameters: [String: Any]? +} + +/// A tool call parsed out of generated text. +struct LlamaParsedToolCall: Equatable { + let name: String + let argumentsJSON: String +} + +// MARK: - System prompt rendering + +extension LlamaToolCallFormat { + /// Renders the tool section of the system prompt and merges it with any + /// existing system text, following each template's own ordering. + func systemMessage(existingText: String, tools: [LlamaToolDefinition]) -> String { + guard !tools.isEmpty else { return existingText } + switch self { + case .hermesJSON: + let block = hermesToolsBlock(tools: tools) + return existingText.isEmpty ? block : existingText + "\n\n" + block + case .qwenXML: + let block = qwenXMLToolsBlock(tools: tools) + return existingText.isEmpty ? block : block + "\n\n" + existingText + case .gemma: + let declarations = tools.map { "<|tool>" + gemmaDeclaration(for: $0) + "" }.joined() + return existingText + declarations + } + } + + private func toolSpecJSON(_ tool: LlamaToolDefinition) -> String { + var function: [String: Any] = [ + "name": tool.name, + "description": tool.description, + ] + if let parameters = tool.parameters { + function["parameters"] = parameters + } + let spec: [String: Any] = ["type": "function", "function": function] + guard + let data = try? JSONSerialization.data(withJSONObject: spec, options: [.sortedKeys]), + let json = String(data: data, encoding: .utf8) + else { + return "{}" + } + return json + } + + private func hermesToolsBlock(tools: [LlamaToolDefinition]) -> String { + var block = "# Tools\n\n" + block += "You may call one or more functions to assist with the user query.\n\n" + block += "You are provided with function signatures within XML tags:\n" + for tool in tools { + block += "\n" + toolSpecJSON(tool) + } + block += "\n\n\n" + block += + "For each function call, return a json object with function name and arguments within " + + " XML tags:\n\n" + + "{\"name\": , \"arguments\": }\n" + return block + } + + private func qwenXMLToolsBlock(tools: [LlamaToolDefinition]) -> String { + var block = "# Tools\n\n" + block += "You have access to the following functions:\n\n" + for tool in tools { + block += "\n" + toolSpecJSON(tool) + } + block += "\n\n\n" + block += "If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" + block += "\n\n" + block += "\nvalue_1\n\n" + block += "\nThis is the value for the second parameter\n" + block += "that can span\nmultiple lines\n\n\n\n\n" + block += "\nReminder:\n" + block += + "- Function calls MUST follow the specified format: an inner " + + "block must be nested within XML tags\n" + block += "- Required parameters MUST be specified\n" + block += + "- You may provide optional reasoning for your function call in natural language " + + "BEFORE the function call, but NOT after\n" + block += + "- If there is no function call available, answer the question like normal with your " + + "current knowledge and do not tell the user about function calls\n" + block += "" + return block + } +} + +// MARK: - Gemma declaration and argument notation + +extension LlamaToolCallFormat { + /// Renders one Gemma 4 function declaration: + /// `declaration:name{description:<|"|>...<|"|>,parameters:{...}}`. + /// Types are uppercased and strings are quoted with the `<|"|>` token, per + /// the canonical template's `format_function_declaration` macro. + fileprivate func gemmaDeclaration(for tool: LlamaToolDefinition) -> String { + var rendered = "declaration:\(tool.name){description:\(gemmaQuote(tool.description))" + if let parameters = tool.parameters { + rendered += ",parameters:{" + var parts: [String] = [] + if let properties = parameters["properties"] as? [String: Any], !properties.isEmpty { + parts.append("properties:{" + gemmaProperties(properties) + "}") + } + if let required = parameters["required"] as? [Any], !required.isEmpty { + let items = required.map { gemmaQuote("\($0)") }.joined(separator: ",") + parts.append("required:[\(items)]") + } + if let type = parameters["type"] as? String { + parts.append("type:\(gemmaQuote(type.uppercased()))") + } + rendered += parts.joined(separator: ",") + "}" + } + rendered += "}" + return rendered + } + + private func gemmaProperties(_ properties: [String: Any]) -> String { + var parts: [String] = [] + for key in properties.keys.sorted() { + guard let value = properties[key] as? [String: Any] else { continue } + var fields: [String] = [] + if let description = value["description"] as? String { + fields.append("description:\(gemmaQuote(description))") + } + let type = (value["type"] as? String)?.uppercased() ?? "STRING" + if type == "STRING", let enumValues = value["enum"] as? [Any] { + let items = enumValues.map { gemmaArgument($0) }.joined(separator: ",") + fields.append("enum:[\(items)]") + } + if type == "ARRAY", let items = value["items"] as? [String: Any], !items.isEmpty { + var itemFields: [String] = [] + for itemKey in items.keys.sorted() { + guard let itemValue = items[itemKey] else { continue } + if itemKey == "type", let itemType = itemValue as? String { + itemFields.append("type:\(gemmaQuote(itemType.uppercased()))") + } else if itemKey == "properties", let nested = itemValue as? [String: Any] { + itemFields.append("properties:{" + gemmaProperties(nested) + "}") + } else if itemKey == "required", let required = itemValue as? [Any] { + let names = required.map { gemmaQuote("\($0)") }.joined(separator: ",") + itemFields.append("required:[\(names)]") + } else { + itemFields.append("\(itemKey):\(gemmaArgument(itemValue))") + } + } + fields.append("items:{" + itemFields.joined(separator: ",") + "}") + } + if type == "OBJECT", let nested = value["properties"] as? [String: Any] { + fields.append("properties:{" + gemmaProperties(nested) + "}") + if let required = value["required"] as? [Any], !required.isEmpty { + let names = required.map { gemmaQuote("\($0)") }.joined(separator: ",") + fields.append("required:[\(names)]") + } + } + fields.append("type:\(gemmaQuote(type))") + parts.append("\(key):{" + fields.joined(separator: ",") + "}") + } + return parts.joined(separator: ",") + } + + fileprivate func gemmaQuote(_ string: String) -> String { + "<|\"|>\(string)<|\"|>" + } + + /// Renders one JSON value in Gemma 4 argument notation: unquoted keys, + /// `<|"|>`-quoted strings, and dictionary keys in sorted order. + fileprivate func gemmaArgument(_ value: Any) -> String { + switch value { + case is NSNull: + return "null" + case let string as String: + return gemmaQuote(string) + case let number as NSNumber: + if CFGetTypeID(number) == CFBooleanGetTypeID() { + return number.boolValue ? "true" : "false" + } + if number.doubleValue == number.doubleValue.rounded(), + number.doubleValue.magnitude < 1e15, + !"\(number)".contains(".") + { + return "\(number.int64Value)" + } + return "\(number)" + case let dictionary as [String: Any]: + let fields = dictionary.keys.sorted().map { "\($0):\(gemmaArgument(dictionary[$0]!))" } + return "{" + fields.joined(separator: ",") + "}" + case let array as [Any]: + return "[" + array.map { gemmaArgument($0) }.joined(separator: ",") + "]" + default: + return gemmaQuote("\(value)") + } + } + + /// Renders a JSON object string as Gemma 4 call arguments (the text between + /// the braces of `call:name{...}`). + fileprivate func gemmaArgumentsBody(fromJSON json: String) -> String { + guard + let data = json.data(using: .utf8), + let object = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any] + else { + return "" + } + return object.keys.sorted().map { "\($0):\(gemmaArgument(object[$0]!))" }.joined(separator: ",") + } +} + +// MARK: - Transcript replay rendering + +extension LlamaToolCallFormat { + /// Renders past tool calls as the assistant-message text the model + /// originally produced, so multi-turn history replays faithfully. + func assistantText(for calls: [LlamaParsedToolCall], precededByContent: Bool) -> String { + var parts: [String] = [] + for call in calls { + switch self { + case .hermesJSON: + parts.append( + "\n{\"name\": \"\(call.name)\", \"arguments\": \(call.argumentsJSON)}\n" + ) + case .qwenXML: + var block = "\n\n" + if let data = call.argumentsJSON.data(using: .utf8), + let object = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any] + { + for key in object.keys.sorted() { + block += "\n\(qwenXMLParameterValue(object[key]!))\n\n" + } + } + block += "\n" + parts.append(block) + case .gemma: + parts.append( + "<|tool_call>call:\(call.name){\(gemmaArgumentsBody(fromJSON: call.argumentsJSON))}" + ) + } + } + let joined = parts.joined(separator: "\n") + if precededByContent && self != .gemma { + return "\n" + joined + } + return joined + } + + private func qwenXMLParameterValue(_ value: Any) -> String { + if let string = value as? String { return string } + if let number = value as? NSNumber { + if CFGetTypeID(number) == CFBooleanGetTypeID() { + return number.boolValue ? "true" : "false" + } + return "\(number)" + } + guard + let data = try? JSONSerialization.data(withJSONObject: value, options: [.sortedKeys]), + let json = String(data: data, encoding: .utf8) + else { + return "\(value)" + } + return json + } + + /// Renders one tool output as the message that carries it back to the model. + /// Hermes and Qwen XML formats deliver results inside a user turn; Gemma 4 + /// continues the open model turn with a `<|tool_response>` block. + func toolResponseMessage(toolName: String, content: String) -> (role: String, content: String) { + switch self { + case .hermesJSON, .qwenXML: + return ("user", "\n\(content)\n") + case .gemma: + let body: String + if let data = content.data(using: .utf8), + let object = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any] + { + body = object.keys.sorted().map { "\($0):\(gemmaArgument(object[$0]!))" }.joined(separator: ",") + } else { + body = "value:\(gemmaArgument(content))" + } + return ("tool", "<|tool_response>response:\(toolName){\(body)}") + } + } +} + +// MARK: - Parsing generated text + +extension LlamaToolCallFormat { + /// Splits generated text into the visible response and any tool calls, + /// removing the call markup from the visible portion. + func parseToolCalls(in text: String) -> (visibleText: String, calls: [LlamaParsedToolCall]) { + switch self { + case .hermesJSON: + return parseMarkedBlocks(in: text, start: "", end: "") { body in + parseHermesCall(body) + } + case .qwenXML: + return parseMarkedBlocks(in: text, start: "", end: "") { body in + parseQwenXMLCall(body) + } + case .gemma: + return parseGemmaCalls(in: text) + } + } + + private func parseMarkedBlocks( + in text: String, + start: String, + end: String, + parse: (String) -> LlamaParsedToolCall? + ) -> (String, [LlamaParsedToolCall]) { + var visible = "" + var calls: [LlamaParsedToolCall] = [] + var remainder = Substring(text) + while let startRange = remainder.range(of: start) { + visible += remainder[.. LlamaParsedToolCall? { + let trimmed = body.trimmingCharacters(in: .whitespacesAndNewlines) + guard + let data = trimmed.data(using: .utf8), + let object = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any], + let name = object["name"] as? String + else { + return nil + } + var argumentsJSON = "{}" + if let arguments = object["arguments"] { + if let nested = arguments as? String { + argumentsJSON = nested + } else if let argumentsData = try? JSONSerialization.data( + withJSONObject: arguments, + options: [.sortedKeys] + ), let json = String(data: argumentsData, encoding: .utf8) { + argumentsJSON = json + } + } + return LlamaParsedToolCall(name: name, argumentsJSON: argumentsJSON) + } + + private func parseQwenXMLCall(_ body: String) -> LlamaParsedToolCall? { + guard let nameStart = body.range(of: "") else { return nil } + let name = String(afterName[..") else { break } + let key = String(afterParam[..") else { break } + var value = String(afterParam[valueStart ..< paramEnd.lowerBound]) + if value.hasPrefix("\n") { value.removeFirst() } + if value.hasSuffix("\n") { value.removeLast() } + arguments[key] = qwenXMLDecodedValue(value) + remainder = afterParam[paramEnd.upperBound...] + } + + guard + let data = try? JSONSerialization.data(withJSONObject: arguments, options: [.sortedKeys]), + let json = String(data: data, encoding: .utf8) + else { + return nil + } + return LlamaParsedToolCall(name: name, argumentsJSON: json) + } + + /// The XML format writes objects and arrays as JSON but scalars as raw + /// text, so structured values are decoded and everything else stays a + /// string. + private func qwenXMLDecodedValue(_ raw: String) -> Any { + let trimmed = raw.trimmingCharacters(in: .whitespaces) + guard trimmed.hasPrefix("{") || trimmed.hasPrefix("[") else { return raw } + guard + let data = trimmed.data(using: .utf8), + let object = try? JSONSerialization.jsonObject(with: data, options: [.fragmentsAllowed]) + else { + return raw + } + return object + } + + private func parseGemmaCalls(in text: String) -> (String, [LlamaParsedToolCall]) { + var visible = "" + var calls: [LlamaParsedToolCall] = [] + var remainder = Substring(text) + while let startRange = remainder.range(of: "<|tool_call>call:") { + visible += remainder[.."), + rest[..`-quoted strings and counting nested structures. + private func gemmaBalancedBodyEnd( + in text: Substring, + from start: Substring.Index + ) -> Substring.Index? { + var depth = 0 + var index = start + while index < text.endIndex { + if text[index...].hasPrefix("<|\"|>") { + let afterQuote = text.index(index, offsetBy: 5) + guard let closeQuote = text[afterQuote...].range(of: "<|\"|>") else { return nil } + index = closeQuote.upperBound + continue + } + let character = text[index] + if character == "{" || character == "[" { + depth += 1 + } else if character == "]" { + depth -= 1 + } else if character == "}" { + if depth == 0 { return index } + depth -= 1 + } + index = text.index(after: index) + } + return nil + } +} + +/// Parses Gemma 4 argument notation into canonical JSON: unquoted keys, +/// `<|"|>`-quoted strings, nested objects and arrays, and bare +/// number/boolean/null literals. +struct LlamaGemmaArgumentParser { + private let characters: [Character] + private var index = 0 + + init(_ text: String) { + self.characters = Array(text) + } + + /// Parses the full input as an object body (`key:value,...`) and returns + /// it as a JSON object string, or `nil` if the input is malformed. + mutating func parseObjectJSON() -> String? { + guard let object = parseObjectBody(terminators: []) else { return nil } + guard + let data = try? JSONSerialization.data(withJSONObject: object, options: [.sortedKeys]), + let json = String(data: data, encoding: .utf8) + else { + return nil + } + skipWhitespace() + guard index >= characters.count else { return nil } + return json + } + + private mutating func parseObjectBody(terminators: Set) -> [String: Any]? { + var object: [String: Any] = [:] + skipWhitespace() + while index < characters.count, !terminators.contains(characters[index]) { + guard let key = parseKey() else { return nil } + guard consume(":") else { return nil } + guard let value = parseValue() else { return nil } + object[key] = value + skipWhitespace() + if index < characters.count, characters[index] == "," { + index += 1 + skipWhitespace() + } else { + break + } + } + return object + } + + private mutating func parseKey() -> String? { + skipWhitespace() + if let quoted = parseQuotedString() { return quoted } + var key = "" + while index < characters.count { + let character = characters[index] + if character == ":" || character == "," || character == "}" { break } + key.append(character) + index += 1 + } + let trimmed = key.trimmingCharacters(in: .whitespaces) + return trimmed.isEmpty ? nil : trimmed + } + + private mutating func parseValue() -> Any? { + skipWhitespace() + if let string = parseQuotedString() { return string } + guard index < characters.count else { return nil } + switch characters[index] { + case "{": + index += 1 + guard let object = parseObjectBody(terminators: ["}"]) else { return nil } + guard consume("}") else { return nil } + return object + case "[": + index += 1 + var array: [Any] = [] + skipWhitespace() + while index < characters.count, characters[index] != "]" { + guard let element = parseValue() else { return nil } + array.append(element) + skipWhitespace() + if index < characters.count, characters[index] == "," { + index += 1 + skipWhitespace() + } + } + guard consume("]") else { return nil } + return array + default: + var literal = "" + while index < characters.count { + let character = characters[index] + if character == "," || character == "}" || character == "]" { break } + literal.append(character) + index += 1 + } + let trimmed = literal.trimmingCharacters(in: .whitespaces) + switch trimmed { + case "true": return true + case "false": return false + case "null": return NSNull() + default: + if let integer = Int64(trimmed) { return integer } + if let double = Double(trimmed) { return double } + return trimmed + } + } + } + + private mutating func parseQuotedString() -> String? { + guard remainingHasPrefix("<|\"|>") else { return nil } + index += 5 + var value = "" + while index < characters.count { + if remainingHasPrefix("<|\"|>") { + index += 5 + return value + } + value.append(characters[index]) + index += 1 + } + return nil + } + + private func remainingHasPrefix(_ prefix: String) -> Bool { + let prefixCharacters = Array(prefix) + guard index + prefixCharacters.count <= characters.count else { return false } + for offset in 0 ..< prefixCharacters.count + where characters[index + offset] != prefixCharacters[offset] { + return false + } + return true + } + + private mutating func skipWhitespace() { + while index < characters.count, characters[index].isWhitespace { + index += 1 + } + } + + private mutating func consume(_ character: Character) -> Bool { + skipWhitespace() + guard index < characters.count, characters[index] == character else { return false } + index += 1 + return true + } +} diff --git a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift index 3d8f012..fbc632e 100644 --- a/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift @@ -1103,54 +1103,125 @@ import Foundation let userInputProcessing = options[custom: MLXLanguageModel.self]?.processingForUserInput ?? .init(resize: nil) - let chat = convertTranscriptToMLXChat( + let toolSpecs = mlxToolSpecs(for: session) + var chat = convertTranscriptToMLXChat( session: session, fallbackPrompt: prompt.description ) - let userInput = makeUserInput( - chat: chat, - tools: nil, - processing: userInputProcessing, - additionalContext: additionalContext - ) - let lmInput = try await context.processor.prepare(input: userInput) - let resolved = resolveCache( - session: session, - lmInput: lmInput, - generateParameters: generateParameters, - context: context - ) + // Accumulators live outside the tool loop so streamed snapshots stay + // monotonic across rounds: text never shrinks, entries only grow. + var accumulatedText = "" + var accumulatedEntries: [Transcript.Entry] = [] + let maxToolIterations = 8 + var toolIteration = 0 + var previousToolCallSignature: String? + + // Yields a snapshot carrying the cumulative text and tool entries so far. + func yieldSnapshot() { + let raw = GeneratedContent(accumulatedText) + let content: Content.PartiallyGenerated = (accumulatedText as! Content) + .asPartiallyGenerated() + continuation.yield( + .init( + content: content, + rawContent: raw, + transcriptEntries: ArraySlice(accumulatedEntries) + ) + ) + } - let mlxStream = try MLXLMCommon.generate( - input: resolved.input, - cache: resolved.cache, - parameters: generateParameters, - context: context - ) + // Loop until the model stops without pending tool calls (mirrors `respond()`). + toolLoop: while true { + let userInput = makeUserInput( + chat: chat, + tools: toolSpecs, + processing: userInputProcessing, + additionalContext: additionalContext + ) + let lmInput = try await context.processor.prepare(input: userInput) + let resolved = resolveCache( + session: session, + lmInput: lmInput, + generateParameters: generateParameters, + context: context + ) + + let mlxStream = try MLXLMCommon.generate( + input: resolved.input, + cache: resolved.cache, + parameters: generateParameters, + context: context + ) + + let roundStartTextCount = accumulatedText.count + var collectedToolCalls: [MLXLMCommon.ToolCall] = [] + + for await item in mlxStream { + if Task.isCancelled { break toolLoop } + + switch item { + case .chunk(let text): + accumulatedText += text + yieldSnapshot() + case .toolCall(let call): + collectedToolCalls.append(call) + case .info: + break + } + } - var accumulatedText = "" - for await item in mlxStream { - if Task.isCancelled { break } - - switch item { - case .chunk(let text): - accumulatedText += text - let raw = GeneratedContent(accumulatedText) - let content: Content.PartiallyGenerated = (accumulatedText as! Content) - .asPartiallyGenerated() - continuation.yield(.init(content: content, rawContent: raw)) - case .info, .toolCall: - break + storeSessionCache( + cache: resolved.cache, + fullTokens: resolved.fullTokens, + generateParameters: generateParameters, + session: session + ) + + // Feed this round's assistant text back into the chat history. + let roundText = String(accumulatedText.dropFirst(roundStartTextCount)) + if !roundText.isEmpty { + chat.append(.assistant(roundText)) + } + + guard !collectedToolCalls.isEmpty else { break } + + toolIteration += 1 + if toolIteration > maxToolIterations { + throw Self.maxToolIterationsExceededError(limit: maxToolIterations) + } + + let signature = + collectedToolCalls + .map { "\($0.function.name):\($0.function.arguments)" } + .joined(separator: "|") + if signature == previousToolCallSignature { + throw Self.repeatedToolCallLoopError() + } + previousToolCallSignature = signature + + let resolution = try await resolveToolCalls(collectedToolCalls, session: session) + switch resolution { + case .stop(let calls): + if !calls.isEmpty { + accumulatedEntries.append(.toolCalls(Transcript.ToolCalls(calls))) + yieldSnapshot() + } + break toolLoop + case .invocations(let invocations): + if invocations.isEmpty { break toolLoop } + + accumulatedEntries.append( + .toolCalls(Transcript.ToolCalls(invocations.map(\.call))) + ) + for invocation in invocations { + accumulatedEntries.append(.toolOutput(invocation.output)) + chat.append(.tool(toolOutputToJSON(invocation.output))) + } + yieldSnapshot() } } - storeSessionCache( - cache: resolved.cache, - fullTokens: resolved.fullTokens, - generateParameters: generateParameters, - session: session - ) finishScope() finishGenerationSlot() continuation.finish() diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index a319b82..e1aa464 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -47,6 +47,26 @@ import Testing #expect(!response.content.isEmpty) } + @Test func reusesSessionContextAcrossTurns() async throws { + let session = LanguageModelSession(model: model) + var options = GenerationOptions(maximumResponseTokens: 24) + options[custom: LlamaLanguageModel.self] = .init(contextSize: 2048, batchSize: 512) + + let first = try await session.respond( + to: "My favorite color is blue. Reply with OK.", + options: options + ) + #expect(!first.content.isEmpty) + #expect(model.lastReusedTokenCount == 0) + + let second = try await session.respond( + to: "What is my favorite color? Answer with one word.", + options: options + ) + #expect(!second.content.isEmpty) + #expect(model.lastReusedTokenCount > 0) + } + @Test func customGenerationOptionsRoundTrip() { var options = GenerationOptions( temperature: 0.6, @@ -425,3 +445,78 @@ import Testing } } #endif // Llama + +#if Llama + @Suite( + "LlamaLanguageModel vision", + .serialized, + .enabled( + if: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"] != nil + && ProcessInfo.processInfo.environment["LLAMA_VISION_MMPROJ_PATH"] != nil + ) + ) + struct LlamaLanguageModelVisionTests { + static let redSquarePNG = Data( + base64Encoded: "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAIAAABt+uBvAAABC0lEQVR4nO3OMQ0AIAAEsfdvGhyw9gaS" + + "CujO9j34QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UHcBWwZ3g5gacwjAAAAAElFTkSuQmCC" + )! + + let model = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"]!, + mmprojPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MMPROJ_PATH"]! + ) + + @Test func describesImageData() async throws { + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .text(.init(content: "What is the dominant color of this image? Answer with one word.")), + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")), + ]) + ) + ]) + let session = LanguageModelSession(model: model, transcript: transcript) + let response = try await session.respond(to: "") + #expect(response.content.lowercased().contains("red")) + } + + @Test func streamsImageDescription() async throws { + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .text(.init(content: "What is the dominant color of this image? Answer with one word.")), + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")), + ]) + ) + ]) + let session = LanguageModelSession(model: model, transcript: transcript) + let stream = session.streamResponse(to: "") + var last = "" + for try await snapshot in stream { + last = snapshot.content + } + #expect(last.lowercased().contains("red")) + } + + @Test func rejectsImagesWithoutProjector() async throws { + let textOnlyModel = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"]! + ) + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")) + ]) + ) + ]) + let session = LanguageModelSession(model: textOnlyModel, transcript: transcript) + await #expect(throws: LlamaLanguageModelError.self) { + _ = try await session.respond(to: "") + } + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift b/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift new file mode 100644 index 0000000..d171695 --- /dev/null +++ b/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift @@ -0,0 +1,455 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if Llama + @Suite("LlamaToolCallFormat") + struct LlamaToolCallFormatTests { + private let weatherTool = LlamaToolDefinition( + name: "get_weather", + description: "Get the current weather for a city", + parameters: [ + "type": "object", + "properties": [ + "city": [ + "type": "string", + "description": "The city name", + ] + ], + "required": ["city"], + ] + ) + + // MARK: - Detection + + @Test func detectsGemmaFromTurnMarker() { + let template = "{{- '<|turn>' + role + '\\n' }}" + #expect(LlamaToolCallFormat.detect(template: template) == .gemma) + } + + @Test func detectsQwenXMLFromFunctionMarker() { + let template = "{{- '\\n\\n' }}" + #expect(LlamaToolCallFormat.detect(template: template) == .qwenXML) + } + + @Test func defaultsToHermesJSON() { + #expect(LlamaToolCallFormat.detect(template: "<|im_start|>{{ role }}") == .hermesJSON) + #expect(LlamaToolCallFormat.detect(template: nil) == .hermesJSON) + } + + // MARK: - System prompt rendering + + @Test func hermesSystemMessageWrapsToolSpecs() { + let message = LlamaToolCallFormat.hermesJSON.systemMessage( + existingText: "You are helpful.", + tools: [weatherTool] + ) + #expect(message.hasPrefix("You are helpful.\n\n# Tools")) + #expect(message.contains("")) + #expect(message.contains("\"name\":\"get_weather\"")) + #expect(message.contains("{\"name\": , \"arguments\": }")) + } + + @Test func qwenXMLSystemMessagePutsToolsFirst() { + let message = LlamaToolCallFormat.qwenXML.systemMessage( + existingText: "You are helpful.", + tools: [weatherTool] + ) + #expect(message.hasPrefix("# Tools")) + #expect(message.hasSuffix("You are helpful.")) + #expect(message.contains("")) + } + + @Test func gemmaSystemMessageAppendsDeclarations() { + let message = LlamaToolCallFormat.gemma.systemMessage( + existingText: "You are helpful.", + tools: [weatherTool] + ) + #expect(message.hasPrefix("You are helpful.<|tool>declaration:get_weather{")) + #expect(message.hasSuffix("")) + #expect(message.contains("description:<|\"|>Get the current weather for a city<|\"|>")) + #expect(message.contains("city:{description:<|\"|>The city name<|\"|>,type:<|\"|>STRING<|\"|>}")) + #expect(message.contains("required:[<|\"|>city<|\"|>]")) + #expect(message.contains("type:<|\"|>OBJECT<|\"|>")) + } + + @Test func emptyToolListLeavesSystemTextUntouched() { + let message = LlamaToolCallFormat.hermesJSON.systemMessage(existingText: "Hi.", tools: []) + #expect(message == "Hi.") + } + + // MARK: - Hermes JSON parsing + + @Test func parsesHermesCall() { + let text = """ + Let me check that for you. + + {"name": "get_weather", "arguments": {"city": "Paris"}} + + """ + let (visible, calls) = LlamaToolCallFormat.hermesJSON.parseToolCalls(in: text) + #expect(visible == "Let me check that for you.") + #expect(calls.count == 1) + #expect(calls.first?.name == "get_weather") + #expect(calls.first?.argumentsJSON == "{\"city\":\"Paris\"}") + } + + @Test func parsesHermesCallWithStringEncodedArguments() { + let text = "{\"name\": \"f\", \"arguments\": \"{\\\"a\\\": 1}\"}" + let (_, calls) = LlamaToolCallFormat.hermesJSON.parseToolCalls(in: text) + #expect(calls.first?.argumentsJSON == "{\"a\": 1}") + } + + @Test func parsesMultipleHermesCalls() { + let text = """ + + {"name": "a", "arguments": {}} + + + {"name": "b", "arguments": {"x": 2}} + + """ + let (visible, calls) = LlamaToolCallFormat.hermesJSON.parseToolCalls(in: text) + #expect(visible.isEmpty) + #expect(calls.map(\.name) == ["a", "b"]) + } + + @Test func plainTextHasNoHermesCalls() { + let (visible, calls) = LlamaToolCallFormat.hermesJSON.parseToolCalls(in: "Just an answer.") + #expect(visible == "Just an answer.") + #expect(calls.isEmpty) + } + + @Test func unterminatedHermesBlockStaysVisible() { + let text = "Answer {\"name\": \"a\"" + let (visible, calls) = LlamaToolCallFormat.hermesJSON.parseToolCalls(in: text) + #expect(calls.isEmpty) + #expect(visible.contains("")) + } + + // MARK: - Qwen XML parsing + + @Test func parsesQwenXMLCall() { + let text = """ + I will look that up. + + + + Paris + + + + """ + let (visible, calls) = LlamaToolCallFormat.qwenXML.parseToolCalls(in: text) + #expect(visible == "I will look that up.") + #expect(calls.count == 1) + #expect(calls.first?.name == "get_weather") + #expect(calls.first?.argumentsJSON == "{\"city\":\"Paris\"}") + } + + @Test func qwenXMLPreservesMultilineParameterValues() { + let text = """ + + + + line one + line two + + + + """ + let (_, calls) = LlamaToolCallFormat.qwenXML.parseToolCalls(in: text) + #expect(calls.first?.argumentsJSON == "{\"body\":\"line one\\nline two\"}") + } + + @Test func qwenXMLDecodesStructuredParameterValues() { + let text = """ + + + + ["a", "b"] + + + + """ + let (_, calls) = LlamaToolCallFormat.qwenXML.parseToolCalls(in: text) + #expect(calls.first?.argumentsJSON == "{\"items\":[\"a\",\"b\"]}") + } + + // MARK: - Gemma parsing + + @Test func parsesGemmaCall() { + let text = "<|tool_call>call:get_weather{city:<|\"|>Paris<|\"|>}" + let (visible, calls) = LlamaToolCallFormat.gemma.parseToolCalls(in: text) + #expect(visible.isEmpty) + #expect(calls.count == 1) + #expect(calls.first?.name == "get_weather") + #expect(calls.first?.argumentsJSON == "{\"city\":\"Paris\"}") + } + + @Test func gemmaQuotedStringsMayContainStructuralCharacters() { + let text = "<|tool_call>call:f{note:<|\"|>a, {b}: [c]<|\"|>}" + let (_, calls) = LlamaToolCallFormat.gemma.parseToolCalls(in: text) + #expect(calls.first?.argumentsJSON == "{\"note\":\"a, {b}: [c]\"}") + } + + @Test func gemmaParsesScalarAndNestedArguments() { + let text = + "<|tool_call>call:f{count:3,enabled:true,tags:[<|\"|>a<|\"|>,<|\"|>b<|\"|>],meta:{k:<|\"|>v<|\"|>}}" + let (_, calls) = LlamaToolCallFormat.gemma.parseToolCalls(in: text) + #expect( + calls.first?.argumentsJSON + == "{\"count\":3,\"enabled\":true,\"meta\":{\"k\":\"v\"},\"tags\":[\"a\",\"b\"]}" + ) + } + + @Test func gemmaCallWithoutTerminatorStillParses() { + let text = "<|tool_call>call:f{city:<|\"|>Paris<|\"|>}" + let (_, calls) = LlamaToolCallFormat.gemma.parseToolCalls(in: text) + #expect(calls.first?.name == "f") + } + + // MARK: - Transcript replay round trips + + @Test func hermesAssistantTextRoundTrips() { + let call = LlamaParsedToolCall(name: "get_weather", argumentsJSON: "{\"city\":\"Paris\"}") + let text = LlamaToolCallFormat.hermesJSON.assistantText(for: [call], precededByContent: false) + let (_, parsed) = LlamaToolCallFormat.hermesJSON.parseToolCalls(in: text) + #expect(parsed == [call]) + } + + @Test func qwenXMLAssistantTextRoundTrips() { + let call = LlamaParsedToolCall(name: "get_weather", argumentsJSON: "{\"city\":\"Paris\"}") + let text = LlamaToolCallFormat.qwenXML.assistantText(for: [call], precededByContent: false) + let (_, parsed) = LlamaToolCallFormat.qwenXML.parseToolCalls(in: text) + #expect(parsed == [call]) + } + + @Test func gemmaAssistantTextRoundTrips() { + let call = LlamaParsedToolCall(name: "get_weather", argumentsJSON: "{\"city\":\"Paris\"}") + let text = LlamaToolCallFormat.gemma.assistantText(for: [call], precededByContent: false) + #expect(text == "<|tool_call>call:get_weather{city:<|\"|>Paris<|\"|>}") + let (_, parsed) = LlamaToolCallFormat.gemma.parseToolCalls(in: text) + #expect(parsed == [call]) + } + + // MARK: - Gemma thought channels + + @Test func stripsCompletedThoughtChannels() { + let text = "<|channel>thought\nThe user said hi.\nHello there!" + #expect(LlamaToolCallFormat.stripGemmaThoughtChannels(from: text) == "Hello there!") + } + + @Test func stripsAlternateChannelSpelling() { + let text = "<|channel|>thought\nReasoning.\nAnswer." + #expect(LlamaToolCallFormat.stripGemmaThoughtChannels(from: text) == "Answer.") + } + + @Test func stripsUnclosedThoughtChannelToEnd() { + let text = "Partial<|channel>thought\nstill thinking" + #expect(LlamaToolCallFormat.stripGemmaThoughtChannels(from: text) == "Partial") + } + + @Test func stripsMultipleThoughtChannels() { + let text = "<|channel>thought\na\nX<|channel>thought\nb\nY" + #expect(LlamaToolCallFormat.stripGemmaThoughtChannels(from: text) == "XY") + } + + @Test func gemmaParseStripsThoughtChannels() { + let text = "<|channel>thought\nplan\nDone.<|tool_call>call:f{}" + let (visible, calls) = LlamaToolCallFormat.gemma.parseToolCalls(in: text) + #expect(visible == "Done.") + #expect(calls.count == 1) + } + + // MARK: - Streaming visibility + + @Test func streamingWithholdsPartialToolCallMarker() { + let visible = LlamaToolCallFormat.hermesJSON.streamingVisibleText( + in: "The answer is\n{\"name\":", + withholdToolCalls: true + ) + #expect(visible == "Checking.") + } + + @Test func streamingIgnoresToolMarkersWhenToolsInactive() { + let visible = LlamaToolCallFormat.hermesJSON.streamingVisibleText( + in: "text more", + withholdToolCalls: false + ) + #expect(visible == "text more") + } + + @Test func streamingWithholdsGemmaThoughtChannel() { + let format = LlamaToolCallFormat.gemma + #expect(format.streamingVisibleText(in: "<|chan", withholdToolCalls: false) == "") + #expect( + format.streamingVisibleText( + in: "<|channel>thought\nhmm", + withholdToolCalls: false + ) == "" + ) + #expect( + format.streamingVisibleText( + in: "<|channel>thought\nhmm\nHi", + withholdToolCalls: false + ) == "Hi" + ) + } + + @Test func streamingWithholdsGemmaPartialToolMarkerAfterThought() { + let format = LlamaToolCallFormat.gemma + let raw = "<|channel>thought\nplan\nSure.<|tool_" + #expect(format.streamingVisibleText(in: raw, withholdToolCalls: true) == "Sure.") + } + + // MARK: - Tool response messages + + @Test func hermesToolResponseIsAUserTurn() { + let message = LlamaToolCallFormat.hermesJSON.toolResponseMessage( + toolName: "get_weather", + content: "{\"temperature\": 21}" + ) + #expect(message.role == "user") + #expect(message.content == "\n{\"temperature\": 21}\n") + } + + @Test func gemmaToolResponseContinuesTheModelTurn() { + let message = LlamaToolCallFormat.gemma.toolResponseMessage( + toolName: "get_weather", + content: "{\"temperature\": 21}" + ) + #expect(message.role == "tool") + #expect( + message.content + == "<|tool_response>response:get_weather{temperature:21}" + ) + } + + @Test func gemmaScalarToolResponseWrapsInValue() { + let message = LlamaToolCallFormat.gemma.toolResponseMessage(toolName: "f", content: "done") + #expect(message.content == "<|tool_response>response:f{value:<|\"|>done<|\"|>}") + } + } + + @Suite( + "LlamaLanguageModel tools", + .serialized, + .enabled(if: ProcessInfo.processInfo.environment["LLAMA_TOOL_MODEL_PATH"] != nil) + ) + struct LlamaLanguageModelToolTests { + let model = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_TOOL_MODEL_PATH"]! + ) + + @Test func executesToolAndAnswersFromItsOutput() async throws { + let weatherTool = spy(on: WeatherTool()) + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + var options = GenerationOptions(temperature: 0.0, maximumResponseTokens: 1024) + options[custom: LlamaLanguageModel.self] = .init(contextSize: 4096) + let response = try await session.respond( + to: "How's the weather in Paris? Use the getWeather tool.", + options: options + ) + + var foundToolOutput = false + for case let .toolOutput(toolOutput) in response.transcriptEntries { + #expect(toolOutput.toolName == weatherTool.name) + foundToolOutput = true + } + #expect(foundToolOutput) + + let calls = await weatherTool.calls + #expect(calls.count == 1) + #expect(calls.first?.arguments.city.contains("Paris") == true) + #expect(response.content.lowercased().contains("72") || response.content.lowercased().contains("sunny")) + #expect(!response.content.contains("")) + } + + @Test func replaysToolExchangeInFollowUpTurns() async throws { + let weatherTool = spy(on: WeatherTool()) + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + var options = GenerationOptions(temperature: 0.0, maximumResponseTokens: 1024) + options[custom: LlamaLanguageModel.self] = .init(contextSize: 4096) + _ = try await session.respond( + to: "How's the weather in Paris? Use the getWeather tool.", + options: options + ) + let followUp = try await session.respond( + to: "What temperature did you just report, in Fahrenheit? Answer with just the number.", + options: options + ) + + let calls = await weatherTool.calls + #expect(calls.count == 1) + #expect(followUp.content.contains("72")) + } + + @Test func streamsToolExchangeProgressively() async throws { + let weatherTool = spy(on: WeatherTool()) + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + var options = GenerationOptions(temperature: 0.0, maximumResponseTokens: 1024) + options[custom: LlamaLanguageModel.self] = .init(contextSize: 4096) + let stream = session.streamResponse( + to: "How's the weather in Paris? Use the getWeather tool.", + options: options + ) + var snapshots: [String] = [] + var sawToolOutputEntry = false + for try await snapshot in stream { + snapshots.append(snapshot.content) + for case .toolOutput(_) in snapshot.transcriptEntries { + sawToolOutputEntry = true + } + } + + let calls = await weatherTool.calls + #expect(calls.count == 1) + #expect(sawToolOutputEntry) + #expect(snapshots.count > 3) + let final = snapshots.last ?? "" + #expect(final.lowercased().contains("72") || final.lowercased().contains("sunny")) + for content in snapshots { + #expect(!content.contains("")) + #expect(!content.contains("<|tool_call>")) + #expect(!content.contains("<|channel")) + } + } + + @Test func answersDirectlyWhenNoToolApplies() async throws { + let weatherTool = spy(on: WeatherTool()) + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + var options = GenerationOptions(temperature: 0.0, maximumResponseTokens: 1024) + options[custom: LlamaLanguageModel.self] = .init(contextSize: 4096) + let response = try await session.respond( + to: "What is 2 + 2? Answer with just the number.", + options: options + ) + + let calls = await weatherTool.calls + #expect(calls.isEmpty) + #expect(response.content.contains("4")) + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift index bb048d3..2ac68bc 100644 --- a/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift @@ -144,6 +144,45 @@ import Testing } } + @Test func streamingWithTools() async throws { + let weatherTool = spy(on: WeatherTool()) + let session = LanguageModelSession( + model: model, + tools: [weatherTool], + instructions: "You are a helpful assistant. Use available tools when needed." + ) + + let stream = session.streamResponse(to: "How's the weather in San Francisco?") + + // Iterate the stream, keeping the last snapshot as the final state. + var snapshotCount = 0 + var lastSnapshot: LanguageModelSession.ResponseStream.Snapshot? + for try await snapshot in stream { + snapshotCount += 1 + lastSnapshot = snapshot + } + + // The stream yielded incremental snapshots and produced text. + #expect(snapshotCount >= 1) + #expect(!(lastSnapshot?.content.isEmpty ?? true)) + + // The tool actually executed. + let calls = await weatherTool.calls + #expect(calls.count >= 1) + if let first = calls.first { + #expect(first.arguments.city.contains("San Francisco")) + } + + // Tool activity surfaces through the stream's transcript entries. + var foundToolOutput = false + for case let .toolOutput(toolOutput) in lastSnapshot?.transcriptEntries ?? [] { + #expect(!toolOutput.id.isEmpty) + #expect(toolOutput.toolName == weatherTool.name) + foundToolOutput = true + } + #expect(foundToolOutput) + } + @Test func multimodalWithImageURL() async throws { let transcript = Transcript(entries: [ .prompt(