From f6d15148e9e7b494b04931f1c948ef5a218b5d91 Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 14:36:42 -0700 Subject: [PATCH 1/6] Add vision support to LlamaLanguageModel via mtmd Image segments threw unsupportedFeature because the backend had no multimodal path, even though the prebuilt llama.cpp binaries ship the mtmd library and its helpers. Accept an mmprojPath at initialization and load the projector next to the model. When a projector is present, prompt formatting replaces each image segment with the mtmd media marker and collects payloads in order, then generation tokenizes the marker-annotated prompt with mtmd_tokenize and evaluates text and image chunks through mtmd_helper_eval_chunks before sampling continues from the resulting position. Both respond and streaming support images, and models without a projector keep rejecting image input. Adds live tests generating from an embedded test image through both paths. --- .../Models/LlamaLanguageModel.swift | 353 ++++++++++++++++-- .../LlamaLanguageModelTests.swift | 75 ++++ 2 files changed, 397 insertions(+), 31 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index d2f29b03..4dd41aa0 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,6 +458,9 @@ 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 @@ -462,9 +470,18 @@ import Foundation /// - 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,6 +522,9 @@ import Foundation } deinit { + if let mtmdContext = mtmdContext { + mtmd_free(mtmdContext) + } if let model = model { llama_model_free(model) } @@ -517,8 +537,9 @@ 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) @@ -542,6 +563,9 @@ import Foundation llama_set_causal_attn(context, true) llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) + var promptImages: [Data] = [] + let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil + let fullPrompt: String if includeSchemaInPrompt, type != String.self { fullPrompt = try formatPrompt( @@ -550,18 +574,39 @@ import Foundation assistantPrefill: runtimeOptions.assistantPrefill ) } else { - fullPrompt = try formatPrompt(for: session, assistantPrefill: runtimeOptions.assistantPrefill) + fullPrompt = try formatPrompt( + for: session, + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages + ) } 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 text: String + if promptImages.isEmpty { + text = try await generateText( + context: context, + model: model!, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions + ) + } else { + var accumulated = "" + try performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions + ) { tokenText in + accumulated += tokenText + } + text = accumulated + } return LanguageModelSession.Response( content: text as! Content, @@ -599,15 +644,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> = @@ -637,18 +683,18 @@ import Foundation llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) var accumulatedText = "" + var promptImages: [Data] = [] + let imageMarker = + self.mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil let fullPrompt = try self.formatPrompt( for: session, - assistantPrefill: runtimeOptions.assistantPrefill + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages ) - try self.performTextGeneration( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) { tokenText in + let yieldToken: (String) -> Void = { tokenText in accumulatedText += tokenText let snapshot = LanguageModelSession.ResponseStream.Snapshot( @@ -658,6 +704,26 @@ import Foundation continuation.yield(snapshot) } + if promptImages.isEmpty { + try self.performTextGeneration( + context: context, + model: model!, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: yieldToken + ) + } else { + try self.performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: yieldToken + ) + } + continuation.finish() } catch { continuation.finish(throwing: error) @@ -686,6 +752,10 @@ import Foundation llama_backend_init() // Free any existing model before loading a new one + if let existingContext = mtmdContext { + mtmd_free(existingContext) + self.mtmdContext = nil + } if let existingModel = model { llama_model_free(existingModel) self.model = nil @@ -696,6 +766,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 = legacyDefaults.threads + 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 @@ -1285,6 +1371,146 @@ import Foundation } } + /// 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) -> Void + ) throws { + guard let mtmdContext, let model = self.model, + let vocab = llama_model_get_vocab(model) + else { + throw LlamaLanguageModelError.contextInitializationFailed + } + + var bitmaps: [OpaquePointer?] = [] + defer { + for bitmap in bitmaps { + if let bitmap { + mtmd_bitmap_free(bitmap) + } + } + } + for imageData in images { + // Pinned to the current llama.swift signature. llama.cpp master adds a + // trailing options argument to this helper; update alongside the dependency. + 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) + } + + 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 + } + + 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 + } + + 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) + + 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) + + var batch = llama_batch_init(1, 0, 1) + defer { llama_batch_free(batch) } + + 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) + + if llama_vocab_is_eog(vocab, nextToken) { + break + } + + if let tokenText = tokenToText(vocab: vocab, token: nextToken) { + onToken(tokenText) + } + + 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 + } + } + // MARK: - Image Validation private func validateNoImageSegments(in session: LanguageModelSession) throws { @@ -1422,6 +1648,23 @@ import Foundation for session: LanguageModelSession, extraSystemMessage: String? = nil, assistantPrefill: String? = nil + ) throws -> String { + var images: [Data] = [] + return try formatPrompt( + for: session, + extraSystemMessage: extraSystemMessage, + assistantPrefill: assistantPrefill, + imageMarker: nil, + images: &images + ) + } + + private func formatPrompt( + for session: LanguageModelSession, + extraSystemMessage: String?, + assistantPrefill: String?, + imageMarker: String?, + images: inout [Data] ) throws -> String { guard let model = self.model else { throw LlamaLanguageModelError.modelLoadFailed @@ -1432,19 +1675,31 @@ import Foundation for entry in session.transcript { 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)) } @@ -1522,6 +1777,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/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index a319b82c..ed5a18f6 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -425,3 +425,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 From 03ada51f829675e59f4aa6629f4c13624374119d Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 15:13:49 -0700 Subject: [PATCH 2/6] Render the Gemma 4 chat format when template detection fails Gemma 4's canonical chat template no longer contains the start_of_turn marker that llama_chat_apply_template keys its Gemma detection on, so formatting threw encodingFailed for every Gemma 4 GGUF. When template application fails and the embedded template carries the Gemma 4 turn syntax, render it directly: turns open with a turn marker and role, close with the reverse marker, the assistant role is named model, and generation opens a model turn. The BOS token is applied during tokenization, and thinking is opt-in in this format so no suppression is needed. --- .../Models/LlamaLanguageModel.swift | 24 ++++++++++++ .../LlamaGemma4TemplateTests.swift | 37 +++++++++++++++++++ 2 files changed, 61 insertions(+) create mode 100644 Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 4dd41aa0..bc5d3340 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -1741,6 +1741,9 @@ import Foundation ) guard requiredSize > 0 else { + if let tmpl, String(cString: tmpl).contains("<|turn>") { + return Self.renderGemma4Prompt(messages: messages, assistantPrefill: assistantPrefill) + } throw LlamaLanguageModelError.encodingFailed } @@ -1770,6 +1773,27 @@ 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. + static func renderGemma4Prompt( + messages: [(role: String, content: String)], + assistantPrefill: String? + ) -> String { + var rendered = "" + for message in messages { + let role = message.role == "assistant" ? "model" : message.role + let content = message.content.trimmingCharacters(in: .whitespacesAndNewlines) + rendered += "<|turn>\(role)\n\(content)\n" + } + 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 } diff --git a/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift b/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift new file mode 100644 index 00000000..756eb412 --- /dev/null +++ b/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift @@ -0,0 +1,37 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if Llama + @Suite("LlamaLanguageModel Gemma 4 template") + struct LlamaGemma4TemplateTests { + @Test func rendersTurnsAndOpensTheModelTurn() { + let rendered = LlamaLanguageModel.renderGemma4Prompt( + messages: [ + (role: "system", content: "Be brief."), + (role: "user", content: "Hi"), + (role: "assistant", content: "Hello"), + (role: "user", content: "Bye"), + ], + assistantPrefill: nil + ) + #expect( + rendered + == "<|turn>system\nBe brief.\n" + + "<|turn>user\nHi\n" + + "<|turn>model\nHello\n" + + "<|turn>user\nBye\n" + + "<|turn>model\n" + ) + } + + @Test func appendsTheAssistantPrefill() { + let rendered = LlamaLanguageModel.renderGemma4Prompt( + messages: [(role: "user", content: "Hi")], + assistantPrefill: "" + ) + #expect(rendered == "<|turn>user\nHi\n<|turn>model\n") + } + } +#endif From e79f8f1019f0ff129e928572ac9f082fd5eec4ab Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 20:09:26 -0700 Subject: [PATCH 3/6] Reuse a per-session context across chat exchanges Every generation created a fresh llama_context and prefilled the full rendered conversation from token zero, so multi-turn chat cost grew with the square of the transcript and long conversations spent most of their time re-decoding history. Keep one context alive per session for plain chat generations. Each exchange tokenizes the rendered prompt, keeps the longest token prefix shared with the context's recorded state, removes diverged state with llama_memory_seq_rm, and decodes only the remainder. Backends that cannot rewind, such as recurrent models, fall back to clearing memory and decoding the full prompt, and appends need no rewind on any backend. The final prompt token is always re-decoded so sampling has fresh logits, generated tokens extend the recorded state as they decode, and any generation error discards the cached context. Structured generation, image prompts, and encoder models keep single-use contexts, and clearCachedContext lets consumers free the cached state under memory pressure. Adds a live test asserting prefix reuse on the second turn of a session. --- .../Models/LlamaLanguageModel.swift | 429 +++++++++++------- .../LlamaLanguageModelTests.swift | 20 + 2 files changed, 289 insertions(+), 160 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index bc5d3340..e6204a1f 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -464,6 +464,234 @@ import Foundation /// 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 + /// Whether a generation is currently decoding on the context. + var isCheckedOut: Bool + /// Whether the context should be freed once the current generation releases it. + var discardWhenReleased: Bool + } + + /// Guards `cachedSessionContext`. A generation checks the cached context out + /// for its whole run, so a concurrent generation for another session never + /// frees a context that is still decoding: it runs on a transient context + /// instead and leaves the cache untouched. + private let sessionContextLock = NSLock() + 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 cached context, including its KV state, otherwise lives as long as the + /// model. The next chat generation prefills its full prompt again. Call this + /// under memory pressure or when a session is discarded. If a generation is + /// running on the cached context, it is freed as soon as that generation ends. + public func clearCachedContext() { + discardCachedSessionContext() + } + + private func discardCachedSessionContext() { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + guard var cached = cachedSessionContext else { return } + if cached.isCheckedOut { + cached.discardWhenReleased = true + cachedSessionContext = cached + return + } + llama_free(cached.context) + cachedSessionContext = nil + } + + private func recordCachedTokens(_ tokens: [llama_token], context: OpaquePointer) { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + guard var cached = cachedSessionContext, cached.context == context else { return } + cached.tokens = tokens + cachedSessionContext = cached + } + + /// Returns a context obtained from `acquireSessionContext`. The cached + /// context is checked back in (or freed, if a discard was requested while + /// it was busy); a transient context is freed. + private func releaseSessionContext(_ context: OpaquePointer) { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + if var cached = cachedSessionContext, cached.context == context { + if cached.discardWhenReleased { + llama_free(cached.context) + cachedSessionContext = nil + } else { + cached.isCheckedOut = false + cachedSessionContext = cached + } + return + } + llama_free(context) + } + + /// Returns a context for the session along with the index of the first + /// prompt token that still needs to be decoded. Pair every call with + /// `releaseSessionContext(_:)` once generation ends. + /// + /// 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. + /// + /// While another generation holds the cached context, the caller gets a + /// transient context that decodes the full prompt and is not cached. + private func acquireSessionContext( + for session: LanguageModelSession, + promptTokens: [llama_token], + options: ResolvedGenerationOptions + ) throws -> (context: OpaquePointer, startIndex: Int) { + let sessionID = ObjectIdentifier(session) + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + + if var cached = cachedSessionContext, + !cached.isCheckedOut, + 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)) + cached.isCheckedOut = true + cachedSessionContext = cached + return (cached.context, common) + } + + if let cached = cachedSessionContext, cached.isCheckedOut { + return (try makeFreshContext(options: options), 0) + } + + if let cached = cachedSessionContext { + llama_free(cached.context) + cachedSessionContext = nil + } + 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, + isCheckedOut: true, + discardWhenReleased: false + ) + 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) -> Void + ) 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 + ) + defer { releaseSessionContext(context) } + 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: @@ -522,6 +750,9 @@ import Foundation } deinit { + if let cached = cachedSessionContext { + llama_free(cached.context) + } if let mtmdContext = mtmdContext { mtmd_free(mtmdContext) } @@ -544,24 +775,6 @@ import Foundation 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_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) var promptImages: [Data] = [] let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil @@ -587,14 +800,20 @@ import Foundation let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 let text: String if promptImages.isEmpty { - text = try await generateText( - context: context, - model: model!, + var accumulated = "" + try generateChatText( + session: session, prompt: fullPrompt, maxTokens: maxTokens, options: runtimeOptions - ) + ) { tokenText in + accumulated += tokenText + } + text = accumulated } else { + discardCachedSessionContext() + let context = try makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } var accumulated = "" try performMultimodalGeneration( context: context, @@ -614,6 +833,8 @@ import Foundation transcriptEntries: ArraySlice([]) ) } else { + let context = try makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } let maxTokens = structuredOptions.maximumResponseTokens ?? 512 let jsonString = try await generateStructuredJSON( context: context, @@ -664,23 +885,6 @@ 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 - } - - // Stabilize runtime behavior per-context - llama_set_causal_attn(context, true) - llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) var accumulatedText = "" var promptImages: [Data] = [] @@ -705,15 +909,17 @@ import Foundation } if promptImages.isEmpty { - try self.performTextGeneration( - context: context, - model: model!, + try self.generateChatText( + session: session, prompt: fullPrompt, maxTokens: maxTokens, options: runtimeOptions, onToken: yieldToken ) } else { + self.discardCachedSessionContext() + let context = try self.makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } try self.performMultimodalGeneration( context: context, prompt: fullPrompt, @@ -752,6 +958,7 @@ 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 @@ -897,110 +1104,6 @@ import Foundation 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 - } - } - - return generatedText - } - /// Builds a JSONSchema-informed prompt for structured output. private func schemaPrompt(for schema: GenerationSchema) -> String { let encoder = JSONEncoder() @@ -1265,23 +1368,21 @@ import Foundation } } - private func performTextGeneration( + private func performTokenGeneration( context: OpaquePointer, - model: OpaquePointer, - prompt: String, + vocab: OpaquePointer, + promptTokens: [llama_token], + startIndex: Int, maxTokens: Int, options: ResolvedGenerationOptions, onToken: (String) -> Void ) throws { - guard let vocab = llama_model_get_vocab(model) else { - throw LlamaLanguageModelError.contextInitializationFailed + guard let model = self.model else { + throw LlamaLanguageModelError.modelLoadFailed } - // Tokenize the prompt - let promptTokens = try tokenizeText(vocab: vocab, text: prompt) - guard !promptTokens.isEmpty else { - throw LlamaLanguageModelError.tokenizationFailed - } + lastReusedTokenCount = startIndex + lastPrefillTokenCount = promptTokens.count - startIndex // Initialize batch var batch = llama_batch_init(Int32(options.batchSize), 0, 1) @@ -1294,7 +1395,8 @@ import Foundation vocab: vocab, context: context, batchSize: options.batchSize, - contextSize: options.contextSize + contextSize: options.contextSize, + startIndex: startIndex ) // Initialize sampler chain with options @@ -1332,6 +1434,7 @@ import Foundation // 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 { @@ -1368,7 +1471,10 @@ import Foundation guard decodeResult == 0 else { break } + decodedTokens.append(nextToken) } + + recordCachedTokens(decodedTokens, context: context) } /// Evaluates a marker-annotated multimodal prompt through the projector, @@ -1542,6 +1648,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` @@ -1553,7 +1661,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 { @@ -1564,7 +1673,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 } @@ -1616,7 +1725,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) diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index ed5a18f6..e1aa4640 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, From 1c845c35c816e1089f4b59c038af33842ff321ee Mon Sep 17 00:00:00 2001 From: james-333i Date: Wed, 26 Aug 2026 17:59:06 -0700 Subject: [PATCH 4/6] Add tool calling to LlamaLanguageModel llama_chat_apply_template has no parameter for tool definitions, so tool support is implemented at the prompt layer. The tool syntax is detected from the model's embedded chat template: Hermes-style JSON (Qwen 2.5/3 and most ChatML fine-tunes), Qwen 3.5's XML function/parameter form, and the Gemma 4 canonical format with its token-quoted argument notation. Definitions are rendered into the system prompt following each template's own wording and placement, past tool turns are replayed in the native markup (including Gemma 4's open-model-turn continuation), and calls are parsed back out of generated text and run through the resolve-and-continue loop used by the MLX and Ollama backends, with the same delegate hooks, iteration cap, and repeated-signature guard. Generation stops early once a complete tool-call block is produced. Non-streaming respond() only; streamResponse() ignores session tools as before. --- .../Models/LlamaLanguageModel.swift | 444 ++++++++++-- .../Models/LlamaToolCallFormat.swift | 668 ++++++++++++++++++ .../LlamaToolCallFormatTests.swift | 337 +++++++++ 3 files changed, 1392 insertions(+), 57 deletions(-) create mode 100644 Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift create mode 100644 Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index e6204a1f..9b68b5d9 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -641,7 +641,7 @@ import Foundation prompt: String, maxTokens: Int, options: ResolvedGenerationOptions, - onToken: (String) -> Void + onToken: (String) -> Bool ) throws { guard let model = self.model, let vocab = llama_model_get_vocab(model) else { throw LlamaLanguageModelError.contextInitializationFailed @@ -761,6 +761,176 @@ import Foundation } } + // 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 makeToolPromptContext(for session: LanguageModelSession) throws -> LlamaToolPromptContext? { + guard !session.tools.isEmpty, let model = self.model else { return nil } + let template = llama_model_chat_template(model, nil).map { String(cString: $0) } + let format = LlamaToolCallFormat.detect(template: template) + 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, @@ -776,63 +946,141 @@ import Foundation let runtimeOptions = resolvedOptions(from: options) let structuredOptions = resolvedStructuredOptions(from: options) - var promptImages: [Data] = [] let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil - 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, - extraSystemMessage: nil, - assistantPrefill: runtimeOptions.assistantPrefill, - imageMarker: imageMarker, - images: &promptImages - ) - } - if type == String.self { let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - let text: String - if promptImages.isEmpty { + 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 = "" - try generateChatText( - session: session, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) { tokenText in + 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 } - text = accumulated - } else { - discardCachedSessionContext() - let context = try makeFreshContext(options: runtimeOptions) - defer { llama_free(context) } - var accumulated = "" - try performMultimodalGeneration( - context: context, - prompt: fullPrompt, - images: promptImages, - maxTokens: maxTokens, - options: runtimeOptions - ) { tokenText in - accumulated += tokenText + + 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 { + 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) + } } - text = accumulated } 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 @@ -898,7 +1146,7 @@ import Foundation images: &promptImages ) - let yieldToken: (String) -> Void = { tokenText in + let yieldToken: (String) -> Bool = { tokenText in accumulatedText += tokenText let snapshot = LanguageModelSession.ResponseStream.Snapshot( @@ -906,6 +1154,7 @@ import Foundation rawContent: GeneratedContent(accumulatedText) ) continuation.yield(snapshot) + return true } if promptImages.isEmpty { @@ -1375,7 +1624,7 @@ import Foundation startIndex: Int, maxTokens: Int, options: ResolvedGenerationOptions, - onToken: (String) -> Void + onToken: (String) -> Bool ) throws { guard let model = self.model else { throw LlamaLanguageModelError.modelLoadFailed @@ -1452,7 +1701,9 @@ import Foundation // Convert token to text and yield it if let tokenText = tokenToText(vocab: vocab, token: nextToken) { - onToken(tokenText) + guard onToken(tokenText) else { + break + } } // Prepare batch for next token @@ -1485,7 +1736,7 @@ import Foundation images: [Data], maxTokens: Int, options: ResolvedGenerationOptions, - onToken: (String) -> Void + onToken: (String) -> Bool ) throws { guard let mtmdContext, let model = self.model, let vocab = llama_model_get_vocab(model) @@ -1596,7 +1847,9 @@ import Foundation } if let tokenText = tokenToText(vocab: vocab, token: nextToken) { - onToken(tokenText) + guard onToken(tokenText) else { + break + } } batch.n_tokens = 1 @@ -1756,7 +2009,8 @@ 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( @@ -1764,7 +2018,8 @@ import Foundation extraSystemMessage: extraSystemMessage, assistantPrefill: assistantPrefill, imageMarker: nil, - images: &images + images: &images, + toolContext: toolContext ) } @@ -1773,7 +2028,8 @@ import Foundation extraSystemMessage: String?, assistantPrefill: String?, imageMarker: String?, - images: inout [Data] + images: inout [Data], + toolContext: LlamaToolPromptContext? = nil ) throws -> String { guard let model = self.model else { throw LlamaLanguageModelError.modelLoadFailed @@ -1781,7 +2037,7 @@ 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 = try extractContent( @@ -1813,8 +2069,56 @@ import Foundation 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) } } @@ -1891,12 +2195,38 @@ import Foundation assistantPrefill: String? ) -> String { var rendered = "" - for message in messages { + 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) - rendered += "<|turn>\(role)\n\(content)\n" + 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" } - rendered += "<|turn>model\n" if let assistantPrefill, !assistantPrefill.isEmpty { rendered += assistantPrefill } diff --git a/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift b/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift new file mode 100644 index 00000000..073915e7 --- /dev/null +++ b/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift @@ -0,0 +1,668 @@ +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 "" + } + } +} + +/// 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 isBooleanNumber(number) { + 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 isBooleanNumber(number) { + 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 + } +} + +/// Whether an `NSNumber` produced by JSON decoding holds a boolean. +/// +/// Core Foundation type identity is the exact check on Darwin. swift-corelibs-foundation +/// has no `CFBoolean`, so other platforms fall back to the encoded Objective-C type, +/// which JSON decoding sets to `c` only for booleans. +func isBooleanNumber(_ number: NSNumber) -> Bool { + #if canImport(Darwin) + return CFGetTypeID(number) == CFBooleanGetTypeID() + #else + return String(cString: number.objCType) == "c" + #endif +} diff --git a/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift b/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift new file mode 100644 index 00000000..ecbe6c7b --- /dev/null +++ b/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift @@ -0,0 +1,337 @@ +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: - 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 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 From eb4ba6f8915e79b71e347345b38d72f307c0d85a Mon Sep 17 00:00:00 2001 From: noorbhatia Date: Wed, 12 Aug 2026 10:38:28 +0530 Subject: [PATCH 5/6] Support tool calling in MLXLanguageModel.streamResponse() streamResponse() hardcoded tools: nil and discarded .toolCall stream items, so MLX callers had to choose between streamed tokens and tool calling. respond() already ran the full tool cycle; this ports that while-loop into the streaming path, reusing mlxToolSpecs, resolveToolCalls, makeTranscriptToolCalls, and the maxToolIterations / repeated-signature guards. Text and tool entries accumulate across rounds so snapshots stay monotonic. Also surface streamed tool activity: ResponseStream.Snapshot gains a defaulted transcriptEntries field (ArraySlice), wrapStream appends it to the session transcript before the response entry, and collect() returns it instead of []. The field defaults to empty, so the other providers keep their current behavior; the shared plumbing is ready for them to populate later. Closes #164 Co-Authored-By: Claude Opus 4.8 --- .../LanguageModelSession.swift | 21 ++- .../Models/MLXLanguageModel.swift | 149 +++++++++++++----- .../MLXLanguageModelTests.swift | 39 +++++ 3 files changed, 166 insertions(+), 43 deletions(-) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 688a126c..7dbb9fdb 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/MLXLanguageModel.swift b/Sources/AnyLanguageModel/Models/MLXLanguageModel.swift index 3d8f0125..fbc632ec 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/MLXLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MLXLanguageModelTests.swift index bb048d3e..2ac68bc5 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( From d990360d791a759efa26c75785916504dd67df17 Mon Sep 17 00:00:00 2001 From: james-333i Date: Thu, 27 Aug 2026 10:07:51 -0700 Subject: [PATCH 6/6] Stream tool-call exchanges and strip Gemma 4 thought channels streamResponse() now runs the same resolve-and-continue loop as respond(), yielding snapshots that carry the cumulative visible text and the tool-call and tool-output entries produced so far. Because llama tool calls arrive as text rather than parsed events, snapshots withhold any trailing partial match of a call-start marker until the next token confirms or breaks it, so markup never appears mid-stream. Gemma 4 emits thought-channel spans without being asked: thinking is opt-in via a system-turn token this backend never injects, and the canonical template ships a strip_thinking macro for consumers. Both respond() and streamResponse() now remove completed spans and withhold unclosed ones, recognizing the canonical marker spelling and the variant observed from deployed quantizations. --- .../Models/LlamaLanguageModel.swift | 184 ++++++++++++++---- .../Models/LlamaToolCallFormat.swift | 91 ++++++++- .../LlamaToolCallFormatTests.swift | 118 +++++++++++ 3 files changed, 354 insertions(+), 39 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 9b68b5d9..3cd66608 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -800,10 +800,15 @@ import Foundation ) } - private func makeToolPromptContext(for session: LanguageModelSession) throws -> LlamaToolPromptContext? { - guard !session.tools.isEmpty, let model = self.model else { return nil } + private func currentToolCallFormat() -> LlamaToolCallFormat { + guard let model = self.model else { return .hermesJSON } let template = llama_model_chat_template(model, nil).map { String(cString: $0) } - let format = LlamaToolCallFormat.detect(template: template) + 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) @@ -950,6 +955,7 @@ import Foundation if type == String.self { let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 + let outputFormat = currentToolCallFormat() var toolContext = try makeToolPromptContext(for: session) let maxToolIterations = 8 var toolIteration = 0 @@ -1003,7 +1009,12 @@ import Foundation } guard let format = toolContext?.format else { - text = accumulated + if outputFormat == .gemma { + text = LlamaToolCallFormat.stripGemmaThoughtChannels(from: accumulated) + .trimmingCharacters(in: .whitespacesAndNewlines) + } else { + text = accumulated + } break generationLoop } let (visibleText, parsedCalls) = format.parseToolCalls(in: accumulated) @@ -1133,52 +1144,149 @@ import Foundation let runtimeOptions = resolvedOptions(from: options) let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - - var accumulatedText = "" - var promptImages: [Data] = [] + 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 - let fullPrompt = try self.formatPrompt( - for: session, - extraSystemMessage: nil, - assistantPrefill: runtimeOptions.assistantPrefill, - imageMarker: imageMarker, - images: &promptImages - ) - - let yieldToken: (String) -> Bool = { tokenText in - accumulatedText += tokenText + func yieldSnapshot(_ text: String) { + lastYieldedText = text let snapshot = LanguageModelSession.ResponseStream.Snapshot( - content: (accumulatedText as! Content).asPartiallyGenerated(), - rawContent: GeneratedContent(accumulatedText) + content: (text as! Content).asPartiallyGenerated(), + rawContent: GeneratedContent(text), + transcriptEntries: ArraySlice(accumulatedEntries) ) continuation.yield(snapshot) - return true } - if promptImages.isEmpty { - try self.generateChatText( - session: session, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions, - onToken: yieldToken + generationLoop: while true { + var promptImages: [Data] = [] + let fullPrompt = try self.formatPrompt( + for: session, + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages, + toolContext: toolContext ) - } 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: yieldToken + + 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): + emittedBase += roundVisible + if !calls.isEmpty { + accumulatedEntries.append(.toolCalls(Transcript.ToolCalls(calls))) + yieldSnapshot(emittedBase) + } + break generationLoop + case .invocations(let invocations): + guard !invocations.isEmpty else { + emittedBase += roundVisible + break generationLoop + } + let callsEntry = Transcript.Entry.toolCalls( + Transcript.ToolCalls(invocations.map(\.call)) + ) + 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) + } } + if emittedBase != lastYieldedText { + yieldSnapshot(emittedBase) + } continuation.finish() } catch { continuation.finish(throwing: error) diff --git a/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift b/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift index 073915e7..0e287057 100644 --- a/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift +++ b/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift @@ -40,6 +40,94 @@ enum LlamaToolCallFormat: Sendable, Equatable { 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. @@ -477,7 +565,8 @@ extension LlamaToolCallFormat { remainder = rest } visible += remainder - return (visible.trimmingCharacters(in: .whitespacesAndNewlines), calls) + let stripped = Self.stripGemmaThoughtChannels(from: visible) + return (stripped.trimmingCharacters(in: .whitespacesAndNewlines), calls) } /// Finds the closing brace of a Gemma call body, skipping braces inside diff --git a/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift b/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift index ecbe6c7b..d171695e 100644 --- a/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaToolCallFormatTests.swift @@ -234,6 +234,92 @@ import Testing #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() { @@ -318,6 +404,38 @@ import Testing #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])