From 97b59109bad276ef0190e8be7f45d281e28f6b5c Mon Sep 17 00:00:00 2001 From: Claude Opus 5 Date: Mon, 3 Aug 2026 11:57:07 -0400 Subject: [PATCH 1/5] Add tool calling to LlamaLanguageModel respond and streamResponse MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both paths previously ignored `session.tools` entirely. They now advertise tools, parse tool calls out of the generated text, execute them through the session's tool-execution delegate, and loop until the model answers without requesting tools. Tools cannot be injected through the model's own chat template here. `llama.h` exposes no tool API at all: `llama_chat_message` carries only a role and a content string, and `llama_chat_apply_template` explicitly "does not use a jinja parser" — it sniffs the GGUF template string to select one of a fixed set of built-in C++ templates, none of which accept tool definitions. The tool-aware renderer (`common_chat_templates_apply`) lives in llama.cpp's `common` library, which is not part of the `llama` xcframework that LlamaSwift re-exports. So tools are described in a synthesized system message that requests the Hermes/Qwen `` format. Parsing is deliberately more permissive than emission, and its scope is explicit: `{...}` (Hermes/Qwen/ChatML), `<|python_tag|>{...}` including `;`-separated multi-calls (Llama 3.x), `[TOOL_CALLS] [{...}]` (Mistral), and `{...}`. A bare leading JSON object is also accepted, but only when every name matches a registered tool, so an ordinary JSON answer is not swallowed. Any other family-specific syntax is unsupported and is returned to the caller as ordinary text. Brace matching is string- and escape-aware, and both `arguments` and `parameters` keys are honored, including arguments emitted as a JSON string. Streaming matches the Anthropic baseline: text growth goes through `growStreamingTranscript(text:)`, `.toolCalls` is appended to the live transcript before the corresponding `.toolOutput` entries, and a `.stop` decision records the calls, finishes the stream, and executes nothing. Because tool markup arrives inline with prose rather than on a separate channel, streaming withholds any text that has begun — or could still grow into — a tool-call marker, so markup never reaches the caller. Loop protection mirrors MLX: an 8-iteration cap plus a repeated-tool-call-signature check, both of which record the offending calls before throwing. Transcript tool entries are replayed into the prompt as `` / `` text so a follow-up turn on a session that already used tools sees a coherent history. Tool results use the `user` role rather than a `tool` role because several built-in templates do not recognize a `tool` role, and consecutive same-role messages are merged to preserve the strict alternation those templates assume. The KV cache is cleared between tool iterations since each one re-decodes the whole prompt from position 0. --- .../Models/LlamaLanguageModel.swift | 812 +++++++++++++++++- 1 file changed, 777 insertions(+), 35 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 90e28de0..81b557c9 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -505,32 +505,113 @@ import Foundation llama_set_warmup(context, false) llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) - let fullPrompt: String - if includeSchemaInPrompt, type != String.self { - fullPrompt = try formatPrompt( - for: session, - extraSystemMessage: schemaPrompt(for: type.generationSchema) - ) - } else { - fullPrompt = try formatPrompt(for: session) - } - 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 knownToolNames = Set(session.tools.map(\.name)) + + // Tools are advertised through a synthesized system message because the binding + // cannot pass them to the chat template. See `toolInstructionMessage(for:)`. + var messages = chatMessages(for: session) + if let instruction = toolInstructionMessage(for: session) { + messages = mergingToolInstruction(instruction, into: messages) + } + + var transcriptEntries: [Transcript.Entry] = [] + var visibleChunks: [String] = [] + var toolIteration = 0 + var previousToolCallSignature: String? + + while true { + try Task.checkCancellation() + + if toolIteration > 0, let memory = llama_get_memory(context) { + // Each iteration re-decodes the whole prompt from position 0, so the KV + // cache from the previous iteration has to be dropped first. + llama_memory_clear(memory, true) + } + + let iterationPrompt = try renderPrompt(messages) + let rawText = try await generateText( + context: context, + model: model!, + prompt: iterationPrompt, + maxTokens: maxTokens, + options: runtimeOptions + ) + + let split = llamaSplitToolCalls(from: rawText, knownToolNames: knownToolNames) + if !split.visibleText.isEmpty { + visibleChunks.append(split.visibleText) + } + + guard !split.toolCalls.isEmpty else { break } + + toolIteration += 1 + if toolIteration > Self.maximumToolIterations { + transcriptEntries.append( + .toolCalls(Transcript.ToolCalls(makeTranscriptToolCalls(from: split.toolCalls))) + ) + throw Self.maxToolIterationsExceededError(limit: Self.maximumToolIterations) + } + + let signature = Self.toolCallSignature(for: split.toolCalls) + if signature == previousToolCallSignature { + transcriptEntries.append( + .toolCalls(Transcript.ToolCalls(makeTranscriptToolCalls(from: split.toolCalls))) + ) + throw Self.repeatedToolCallLoopError() + } + previousToolCallSignature = signature + + // Replay the model's own turn verbatim so it sees the calls it just made. + messages.append(("assistant", rawText)) + + switch try await resolveToolCalls(split.toolCalls, session: session) { + case .stop(let calls): + if !calls.isEmpty { + transcriptEntries.append(.toolCalls(Transcript.ToolCalls(calls))) + } + let stoppedText = visibleChunks.joined(separator: "\n") + return LanguageModelSession.Response( + content: stoppedText as! Content, + rawContent: GeneratedContent(stoppedText), + transcriptEntries: ArraySlice(transcriptEntries) + ) + + case .invocations(let invocations): + guard !invocations.isEmpty else { break } + + transcriptEntries.append( + .toolCalls(Transcript.ToolCalls(invocations.map(\.call))) + ) + for invocation in invocations { + transcriptEntries.append(.toolOutput(invocation.output)) + } + messages.append( + ( + "user", + invocations.map { toolResponseMarkup($0.output) }.joined(separator: "\n") + ) + ) + messages = mergingConsecutiveRoles(messages) + continue + } + break + } + + let text = visibleChunks.joined(separator: "\n") return LanguageModelSession.Response( content: text as! Content, rawContent: GeneratedContent(text), - transcriptEntries: ArraySlice([]) + transcriptEntries: ArraySlice(transcriptEntries) ) } else { + let fullPrompt = + includeSchemaInPrompt + ? try formatPrompt(for: session, extraSystemMessage: schemaPrompt(for: type.generationSchema)) + : try formatPrompt(for: session) + let maxTokens = structuredOptions.maximumResponseTokens ?? 512 let jsonString = try await generateStructuredJSON( context: context, @@ -599,28 +680,143 @@ import Foundation llama_set_warmup(context, false) llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) - var accumulatedText = "" - let fullPrompt = try self.formatPrompt(for: session) + let knownToolNames = Set(session.tools.map(\.name)) + let hasTools = !knownToolNames.isEmpty + + // Snapshot the prompt messages before generating: the transcript is + // mutated during the stream, and in-flight turns are tracked here + // instead so they are never counted twice. + var messages = self.chatMessages(for: session) + if let instruction = self.toolInstructionMessage(for: session) { + messages = self.mergingToolInstruction(instruction, into: messages) + } + + // Text emitted across all completed tool iterations. + var emittedText = "" + var toolIteration = 0 + var previousToolCallSignature: String? + + func yieldText(_ text: String) { + session.growStreamingTranscript(text: text) + continuation.yield( + LanguageModelSession.ResponseStream.Snapshot( + content: (text as! Content).asPartiallyGenerated(), + rawContent: GeneratedContent(text) + ) + ) + } + + while true { + try Task.checkCancellation() + + if toolIteration > 0, let memory = llama_get_memory(context) { + // Drop the previous iteration's KV cache before re-decoding + // the extended prompt from position 0. + llama_memory_clear(memory, true) + } + + let iterationPrompt = try self.renderPrompt(messages) + var rawText = "" + var lastVisible = "" - do { for try await tokenText in generateTextStream( context: context, model: model!, - prompt: fullPrompt, + prompt: iterationPrompt, maxTokens: maxTokens, options: runtimeOptions ) { - accumulatedText += tokenText + if Task.isCancelled { break } + rawText += tokenText + + // Withhold anything that has begun, or might still become, + // tool-call markup. + let visible = + hasTools + ? llamaStreamableVisiblePrefix(of: rawText) + : rawText + guard visible != lastVisible else { continue } + lastVisible = visible + yieldText(emittedText + visible) + } + + let split = llamaSplitToolCalls( + from: rawText, + knownToolNames: knownToolNames + ) + + // Flush the final text for this iteration now that the full + // response is known and tool markup can be stripped exactly. + if split.visibleText != lastVisible { + yieldText(emittedText + split.visibleText) + } + emittedText += split.visibleText + + guard !split.toolCalls.isEmpty else { break } + + toolIteration += 1 + if toolIteration > Self.maximumToolIterations { + session.appendTranscriptEntry( + .toolCalls( + Transcript.ToolCalls( + makeTranscriptToolCalls(from: split.toolCalls) + ) + ) + ) + throw Self.maxToolIterationsExceededError( + limit: Self.maximumToolIterations + ) + } - let snapshot = LanguageModelSession.ResponseStream.Snapshot( - content: (accumulatedText as! Content).asPartiallyGenerated(), - rawContent: GeneratedContent(accumulatedText) + let signature = Self.toolCallSignature(for: split.toolCalls) + if signature == previousToolCallSignature { + session.appendTranscriptEntry( + .toolCalls( + Transcript.ToolCalls( + makeTranscriptToolCalls(from: split.toolCalls) + ) + ) ) - continuation.yield(snapshot) + throw Self.repeatedToolCallLoopError() } - } catch { - continuation.finish(throwing: error) - return + previousToolCallSignature = signature + + // Replay the model's own turn verbatim. + messages.append(("assistant", rawText)) + + switch try await resolveToolCalls(split.toolCalls, session: session) { + case .stop(let calls): + // Record the calls, then end the stream without executing them. + if !calls.isEmpty { + session.appendTranscriptEntry(.toolCalls(Transcript.ToolCalls(calls))) + } + continuation.finish() + return + + case .invocations(let invocations): + guard !invocations.isEmpty else { break } + + // Tool calls must land in the transcript before their outputs. + session.appendTranscriptEntry( + .toolCalls(Transcript.ToolCalls(invocations.map(\.call))) + ) + for invocation in invocations { + session.appendTranscriptEntry(.toolOutput(invocation.output)) + } + + messages.append( + ( + "user", + invocations + .map { self.toolResponseMarkup($0.output) } + .joined(separator: "\n") + ) + ) + messages = self.mergingConsecutiveRoles(messages) + continue + } + + break } continuation.finish() @@ -1370,10 +1566,20 @@ import Foundation for session: LanguageModelSession, extraSystemMessage: String? = nil ) throws -> String { - guard let model = self.model else { - throw LlamaLanguageModelError.modelLoadFailed - } + try renderPrompt(chatMessages(for: session, extraSystemMessage: extraSystemMessage)) + } + /// Flattens the session transcript into role/content pairs for the chat template. + /// + /// Tool activity is rendered back into the text formats this provider emits and parses, + /// so that a follow-up turn on a session that already used tools sees a coherent history. + /// Tool results are replayed with the `user` role rather than a `tool` role because + /// `llama_chat_apply_template` only implements a fixed set of built-in templates, and + /// several of them (Llama 2, Mistral) do not recognize a `tool` role at all. + private func chatMessages( + for session: LanguageModelSession, + extraSystemMessage: String? = nil + ) -> [(role: String, content: String)] { var messages: [(role: String, content: String)] = [] for entry in session.transcript { @@ -1396,8 +1602,14 @@ import Foundation messages.append(("assistant", text)) } - default: - break + case .toolCalls(let toolCalls): + let markup = toolCallMarkup(Array(toolCalls)) + if !markup.isEmpty { + messages.append(("assistant", markup)) + } + + case .toolOutput(let output): + messages.append(("user", toolResponseMarkup(output))) } } @@ -1405,6 +1617,33 @@ import Foundation messages.append(("system", extraSystemMessage)) } + return mergingConsecutiveRoles(messages) + } + + /// Collapses runs of same-role messages into one. + /// + /// Several built-in llama.cpp templates assume strictly alternating user/assistant turns, + /// and a run of consecutive tool-result messages would otherwise break them. + private func mergingConsecutiveRoles( + _ messages: [(role: String, content: String)] + ) -> [(role: String, content: String)] { + var merged: [(role: String, content: String)] = [] + for message in messages { + if var last = merged.last, last.role == message.role { + last.content += "\n" + message.content + merged[merged.count - 1] = last + } else { + merged.append(message) + } + } + return merged + } + + private func renderPrompt(_ messages: [(role: String, content: String)]) throws -> String { + guard let model = self.model else { + throw LlamaLanguageModelError.modelLoadFailed + } + // Keep C strings alive while using them let cRoles = messages.map { strdup($0.role) } let cContents = messages.map { strdup($0.content) } @@ -1457,6 +1696,144 @@ import Foundation } } + // MARK: - Tool Prompting + + /// The maximum number of tool round-trips allowed in a single response. + private static let maximumToolIterations = 8 + + 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." + ) + ) + } + + /// A stable fingerprint of a batch of tool calls, used to detect a stuck model. + private static func toolCallSignature(for calls: [ParsedLlamaToolCall]) -> String { + calls.map { "\($0.name):\($0.arguments.jsonString)" }.joined(separator: "|") + } + + /// Places the tool instructions in the leading system message. + /// + /// Appending a trailing system message (the way schema prompts are handled) would break the + /// strict user/assistant alternation that several built-in llama.cpp templates require. + private func mergingToolInstruction( + _ instruction: String, + into messages: [(role: String, content: String)] + ) -> [(role: String, content: String)] { + var messages = messages + if let first = messages.first, first.role == "system" { + messages[0] = ("system", first.content + "\n\n" + instruction) + } else { + messages.insert(("system", instruction), at: 0) + } + return messages + } + + /// Builds the system message that advertises the session's tools to the model. + /// + /// - Important: `llama.h` exposes no tool API. `llama_chat_message` carries only a role and + /// a content string, and `llama_chat_apply_template` explicitly "does not use a jinja + /// parser" — it sniffs the GGUF template string to pick one of a fixed set of built-in + /// C++ templates, none of which accept tool definitions. The tool-aware renderer + /// (`common_chat_templates_apply`) lives in llama.cpp's `common` library, which is not + /// part of the `llama` xcframework that LlamaSwift re-exports. + /// + /// So tools cannot be injected through the model's own chat template here. Instead we + /// describe them in a system message and ask for the Hermes/Qwen `` format, + /// which is the most widely trained-on text convention across GGUF chat models. Parsing + /// is more permissive than emission: see ``llamaSplitToolCalls(from:knownToolNames:)``. + private func toolInstructionMessage(for session: LanguageModelSession) -> String? { + guard !session.tools.isEmpty else { return nil } + + var specs: [String] = [] + for tool in session.tools { + let resolvedSchema = tool.parameters.withResolvedRoot() ?? tool.parameters + guard + let schemaData = try? JSONEncoder().encode(resolvedSchema), + let schemaObject = try? JSONSerialization.jsonObject(with: schemaData) + else { continue } + + let spec: [String: Any] = [ + "name": tool.name, + "description": tool.description, + "parameters": schemaObject, + ] + guard + let specData = try? JSONSerialization.data(withJSONObject: spec, options: [.sortedKeys]), + let specString = String(data: specData, encoding: .utf8) + else { continue } + + specs.append(specString) + } + + guard !specs.isEmpty else { return nil } + + return """ + You have access to the following tools. Each is described by a JSON object with a \ + name, a description, and a JSON Schema for its arguments. + + \(specs.joined(separator: "\n")) + + To call a tool, reply with nothing but one or more blocks in exactly this form: + + + {"name": "", "arguments": {}} + + + Results come back as blocks. Once you have the results you need, \ + answer the user in plain text and do not emit any more tool calls. If no tool is \ + needed, just answer in plain text. + """ + } + + /// Renders transcript tool calls back into the `` text form the model emits. + private func toolCallMarkup(_ calls: [Transcript.ToolCall]) -> String { + calls.map { call in + "\n{\"name\": \"\(call.toolName)\", \"arguments\": \(call.arguments.jsonString)}\n" + } + .joined(separator: "\n") + } + + /// Renders a tool output into the `` text form fed back to the model. + private func toolResponseMarkup(_ output: Transcript.ToolOutput) -> String { + let text = toolOutputText(output) + let payload = + (try? JSONSerialization.data( + withJSONObject: ["name": output.toolName, "content": text], + options: [.sortedKeys] + )) + .flatMap { String(data: $0, encoding: .utf8) } + ?? "{\"name\": \"\(output.toolName)\"}" + return "\n\(payload)\n" + } + + private func toolOutputText(_ output: Transcript.ToolOutput) -> String { + output.segments.compactMap { segment -> String? in + switch segment { + case .text(let text): + return text.content + case .structure(let structured): + return structured.content.jsonString + case .image: + // Image tool output has no textual form to feed back through the chat template. + return nil + } + } + .joined(separator: "\n") + } + private func extractText(from segments: [Transcript.Segment]) -> String { segments.compactMap { segment -> String? in if case .text(let t) = segment { return t.content } @@ -1528,6 +1905,371 @@ import Foundation } } + // MARK: - Tool Call Parsing + + /// A tool call recovered from the model's generated text. + private struct ParsedLlamaToolCall { + let name: String + let arguments: GeneratedContent + } + + /// Opening markers for every tool-call text format this provider recognizes. + /// + /// Local GGUF models express tool calls as generated text in a family-specific format, and the + /// binding gives us no structured tool channel to read instead. This list is the explicit, + /// deliberately narrow scope of what is supported: + /// + /// - `{...}` — Hermes / Qwen / ChatML-with-tools. This is also the + /// format we ask for in the system message, so it is the primary path. + /// - `<|python_tag|>{...}` — Llama 3.1 / 3.2, including `;`-separated multi-calls. + /// - `[TOOL_CALLS] [{...}]` — Mistral / Mixtral. + /// - `{...}` — Llama 3.2 functionary-style. + /// + /// A bare leading JSON object is also accepted, but only when its name matches a registered + /// tool (see ``llamaSplitToolCalls(from:knownToolNames:)``). + /// + /// Any other family-specific syntax is *not* supported and its text will be returned to the + /// caller as ordinary model output. + private let llamaToolCallMarkers = ["", "<|python_tag|>", "[TOOL_CALLS]", " Range? { + var index = start + while index < text.endIndex, text[index].isWhitespace { + index = text.index(after: index) + } + guard index < text.endIndex else { return nil } + + let opening = text[index] + let closing: Character + switch opening { + case "{": closing = "}" + case "[": closing = "]" + default: return nil + } + + let begin = index + var depth = 0 + var inString = false + var escaped = false + + while index < text.endIndex { + let character = text[index] + if escaped { + escaped = false + } else if inString { + if character == "\\" { + escaped = true + } else if character == "\"" { + inString = false + } + } else if character == "\"" { + inString = true + } else if character == opening { + depth += 1 + } else if character == closing { + depth -= 1 + if depth == 0 { + return begin ..< text.index(after: index) + } + } + index = text.index(after: index) + } + + return nil + } + + /// Decodes one tool call from a already-parsed JSON value. + private func llamaToolCall(fromObject object: Any, defaultName: String? = nil) -> ParsedLlamaToolCall? { + guard var dictionary = object as? [String: Any] else { return nil } + + // OpenAI-shaped nesting: {"type": "function", "function": {"name": ..., "arguments": ...}} + if let function = dictionary["function"] as? [String: Any] { + dictionary = function + } + + guard let name = (dictionary["name"] as? String) ?? defaultName, !name.isEmpty else { return nil } + + // Families disagree on the argument key, and some emit the payload as a JSON *string*. + let rawArguments = dictionary["arguments"] ?? dictionary["parameters"] ?? [String: Any]() + let argumentsJSON: String + if let string = rawArguments as? String { + argumentsJSON = string.isEmpty ? "{}" : string + } else if let data = try? JSONSerialization.data(withJSONObject: rawArguments, options: [.sortedKeys]), + let string = String(data: data, encoding: .utf8) + { + argumentsJSON = string + } else { + return nil + } + + guard let arguments = try? GeneratedContent(json: argumentsJSON) else { return nil } + return ParsedLlamaToolCall(name: name, arguments: arguments) + } + + /// Decodes one or more tool calls from a JSON object or array. + private func llamaToolCalls(fromJSON json: String, defaultName: String? = nil) -> [ParsedLlamaToolCall] { + guard let object = try? JSONSerialization.jsonObject(with: Data(json.utf8)) else { return [] } + if let array = object as? [Any] { + return array.compactMap { llamaToolCall(fromObject: $0, defaultName: defaultName) } + } + return [llamaToolCall(fromObject: object, defaultName: defaultName)].compactMap { $0 } + } + + /// Splits generated text into the prose a caller should see and the tool calls embedded in it. + /// + /// See ``llamaToolCallMarkers`` for the exact set of supported formats. + private func llamaSplitToolCalls( + from text: String, + knownToolNames: Set + ) -> (visibleText: String, toolCalls: [ParsedLlamaToolCall]) { + var calls: [ParsedLlamaToolCall] = [] + var visible = "" + var cursor = text.startIndex + + while cursor < text.endIndex { + // Find the earliest marker at or after the cursor. + var found: (range: Range, marker: String)? + for marker in llamaToolCallMarkers { + guard let range = text.range(of: marker, range: cursor ..< text.endIndex) else { continue } + if found == nil || range.lowerBound < found!.range.lowerBound { + found = (range, marker) + } + } + + guard let match = found else { + visible += text[cursor...] + break + } + + visible += text[cursor ..< match.range.lowerBound] + var next = match.range.upperBound + + switch match.marker { + case "{arguments} — the payload is the arguments themselves. + guard let nameEnd = text.range(of: ">", range: next ..< text.endIndex) else { + cursor = next + continue + } + let name = String(text[next ..< nameEnd.lowerBound]) + guard + !name.isEmpty, + let jsonRange = llamaBalancedJSONRange(in: text, from: nameEnd.upperBound) + else { + cursor = nameEnd.upperBound + continue + } + if let arguments = try? GeneratedContent(json: String(text[jsonRange])) { + calls.append(ParsedLlamaToolCall(name: name, arguments: arguments)) + } + next = jsonRange.upperBound + next = llamaConsuming("", in: text, from: next) + + case "": + guard let jsonRange = llamaBalancedJSONRange(in: text, from: next) else { + cursor = next + continue + } + calls.append(contentsOf: llamaToolCalls(fromJSON: String(text[jsonRange]))) + next = jsonRange.upperBound + next = llamaConsuming("", in: text, from: next) + + default: + // <|python_tag|> and [TOOL_CALLS]: one or more payloads, optionally ";"-separated. + var scan = next + while let jsonRange = llamaBalancedJSONRange(in: text, from: scan) { + calls.append(contentsOf: llamaToolCalls(fromJSON: String(text[jsonRange]))) + scan = jsonRange.upperBound + + var probe = scan + while probe < text.endIndex, text[probe].isWhitespace { + probe = text.index(after: probe) + } + guard probe < text.endIndex, text[probe] == ";" else { break } + scan = text.index(after: probe) + } + if scan == next { + // Marker with no decodable payload; treat it as ordinary text and move on. + cursor = next + continue + } + next = llamaConsuming("<|eom_id|>", in: text, from: scan) + } + + cursor = next + } + + if calls.isEmpty { + // Some templates (notably Llama 3.1/3.2 without <|python_tag|>) emit a bare JSON object. + // Only treat that as a tool call when every name matches a registered tool, so that an + // ordinary JSON answer is not silently swallowed. + if let jsonRange = llamaBalancedJSONRange(in: visible, from: visible.startIndex) { + let candidates = llamaToolCalls(fromJSON: String(visible[jsonRange])) + if !candidates.isEmpty, candidates.allSatisfy({ knownToolNames.contains($0.name) }) { + var remaining = visible + remaining.removeSubrange(jsonRange) + return (remaining.trimmingCharacters(in: .whitespacesAndNewlines), candidates) + } + } + } + + return (visible.trimmingCharacters(in: .whitespacesAndNewlines), calls) + } + + /// Advances past `token` if it is the next non-whitespace content, otherwise returns `index`. + private func llamaConsuming(_ token: String, in text: String, from index: String.Index) -> String.Index { + var probe = index + while probe < text.endIndex, text[probe].isWhitespace { + probe = text.index(after: probe) + } + guard text[probe...].hasPrefix(token) else { return index } + return text.index(probe, offsetBy: token.count) + } + + /// Returns the portion of a partially generated response that is safe to show a caller. + /// + /// Tool-call markup arrives inline with prose, so streaming must withhold any text that has + /// begun — or might still turn into — a tool call. + private func llamaStreamableVisiblePrefix(of text: String) -> String { + // A response opening with a JSON container may be a bare-JSON tool call, which cannot be + // recognized until generation finishes. Withhold all of it until then. + let leading = text.drop(while: { $0.isWhitespace }) + if leading.first == "{" || leading.first == "[" { return "" } + + var earliest: String.Index? + for marker in llamaToolCallMarkers { + guard let range = text.range(of: marker) else { continue } + if earliest == nil || range.lowerBound < earliest! { + earliest = range.lowerBound + } + } + if let earliest { return String(text[.. 0 { + let tail = String(text.suffix(holdback)) + if llamaToolCallMarkers.contains(where: { $0.hasPrefix(tail) }) { + return String(text.dropLast(holdback)) + } + holdback -= 1 + } + + return text + } + + // MARK: - Tool Invocation Handling + + private struct LlamaToolInvocationResult { + let call: Transcript.ToolCall + let output: Transcript.ToolOutput + } + + private enum LlamaToolResolutionOutcome { + case stop(calls: [Transcript.ToolCall]) + case invocations([LlamaToolInvocationResult]) + } + + private func makeTranscriptToolCalls(from parsed: [ParsedLlamaToolCall]) -> [Transcript.ToolCall] { + parsed.map { + Transcript.ToolCall(id: UUID().uuidString, toolName: $0.name, arguments: $0.arguments) + } + } + + /// Runs the session's tool-execution delegate protocol over a batch of parsed tool calls. + /// + /// - Note: This mirrors the `resolveToolCalls` / `resolveToolUses` functions that each other + /// provider keeps file-private in its own file. It is duplicated rather than shared because + /// those are `private` and this task may not introduce a shared abstraction. + private func resolveToolCalls( + _ parsedCalls: [ParsedLlamaToolCall], + session: LanguageModelSession + ) async throws -> LlamaToolResolutionOutcome { + 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 = makeTranscriptToolCalls(from: parsedCalls) + + if let delegate = session.toolExecutionDelegate { + await delegate.didGenerateToolCalls(transcriptCalls, in: session) + } + + guard !transcriptCalls.isEmpty else { return .invocations([]) } + + var decisions: [ToolExecutionDecision] = [] + decisions.reserveCapacity(transcriptCalls.count) + + if let delegate = session.toolExecutionDelegate { + for call in transcriptCalls { + let decision = await delegate.toolCallDecision(for: call, in: session) + if case .stop = decision { + return .stop(calls: transcriptCalls) + } + decisions.append(decision) + } + } else { + decisions = Array(repeating: .execute, count: transcriptCalls.count) + } + + var results: [LlamaToolInvocationResult] = [] + results.reserveCapacity(transcriptCalls.count) + + for (index, call) in transcriptCalls.enumerated() { + switch decisions[index] { + case .stop: + // Unreachable: `.stop` returns while decisions are collected. Defensive only. + 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(LlamaToolInvocationResult(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(LlamaToolInvocationResult(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(LlamaToolInvocationResult(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) + } + /// Errors that can occur when using LlamaLanguageModel public enum LlamaLanguageModelError: Error, LocalizedError { case modelLoadFailed From a392ec46f1fd165cbe3abea53ce9057bfb1374a7 Mon Sep 17 00:00:00 2001 From: Claude Opus 5 Date: Mon, 3 Aug 2026 11:57:14 -0400 Subject: [PATCH 2/5] Add withTools and streamWithTools tests for LlamaLanguageModel Modeled on the Anthropic suite's equivalents. `streamWithTools` additionally asserts that `.toolCalls` precedes `.toolOutput` in the live transcript and that tool-call markup never leaks into streamed response text, since for a local model that markup arrives inline with prose. Both follow the suite's existing gating on `LLAMA_MODEL_PATH` and so skip without a local GGUF. They also need a model actually trained to emit tool calls in one of the supported text formats. --- .../LlamaLanguageModelTests.swift | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index 3d9e32c7..0faa3020 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -401,5 +401,85 @@ import Testing #expect(response.content.choices.count == 4) #expect(!response.content.answer.isEmpty) } + + // MARK: - Tool Calling + // + // These require a GGUF model that has actually been trained to emit tool calls in one of + // the text formats `LlamaLanguageModel` parses (Hermes/Qwen ``, Llama 3.x + // `<|python_tag|>`, Mistral `[TOOL_CALLS]`, or ``). Like the rest of this + // suite they are gated on `LLAMA_MODEL_PATH`; pointing that at a model without tool + // training will fail here rather than skip, which is the honest signal. + + @Test func withTools() async throws { + let weatherTool = WeatherTool() + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + let response = try await session.respond( + to: "How's the weather in San Francisco?", + options: GenerationOptions(maximumResponseTokens: 512) + ) + + var foundToolOutput = false + for case let .toolOutput(toolOutput) in response.transcriptEntries { + #expect(!toolOutput.id.isEmpty) + #expect(toolOutput.toolName == "getWeather") + foundToolOutput = true + } + #expect(foundToolOutput) + } + + @Test func streamWithTools() async throws { + let weatherTool = WeatherTool() + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + let stream = session.streamResponse( + to: "How's the weather in San Francisco?", + options: GenerationOptions(maximumResponseTokens: 512) + ) + + var streamedText: [String] = [] + + var toolAppearedInTranscript = false + var toolResponseAppearedInTranscript = false + + for try await snapshot in stream { + streamedText.append(snapshot.content) + + for entry in session.transcript { + switch entry { + case .toolCalls: + toolAppearedInTranscript = true + case .toolOutput: + toolResponseAppearedInTranscript = true + default: break + } + } + } + + #expect(toolAppearedInTranscript, "Expected a tool call to appear in the transcript during streaming.") + #expect( + toolResponseAppearedInTranscript, + "Expected a tool output to appear in the transcript during streaming." + ) + + // Tool calls must be recorded before the outputs they produced. + let firstToolCallIndex = session.transcript.firstIndex { entry in + if case .toolCalls = entry { return true } + return false + } + let firstToolOutputIndex = session.transcript.firstIndex { entry in + if case .toolOutput = entry { return true } + return false + } + if let firstToolCallIndex, let firstToolOutputIndex { + #expect(firstToolCallIndex < firstToolOutputIndex) + } + + // Tool-call markup must never be surfaced to the caller as response text. + let leakedMarkup = streamedText.filter { text in + text.contains("") || text.contains("<|python_tag|>") || text.contains("[TOOL_CALLS]") + } + #expect(leakedMarkup.isEmpty, "Tool-call markup leaked into streamed response text.") + } } #endif // Llama From 9bb9d018f804e32f3f623fa35176f15e62891475 Mon Sep 17 00:00:00 2001 From: Claude Opus 5 Date: Mon, 3 Aug 2026 12:02:06 -0400 Subject: [PATCH 3/5] Add ungated unit tests for the Llama tool-call parser MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tool-call text parser is the substance of Llama tool calling and had no committed coverage: the integration tests need a local GGUF, so nothing exercised the parsing logic in CI. The pure parsing helpers (`ParsedLlamaToolCall`, `llamaToolCallMarkers`, `llamaBalancedJSONRange`, `llamaToolCall`, `llamaToolCalls`, `llamaSplitToolCalls`, `llamaConsuming`, `llamaStreamableVisiblePrefix`) are widened from file-private to `internal` so the test target can reach them via `@testable import`. Nothing becomes public API. The tool *invocation* helpers — `resolveToolCalls`, `makeTranscriptToolCalls`, `LlamaToolInvocationResult`, `LlamaToolResolutionOutcome` — stay file-private, since they need a live session and are not what these tests cover. The new `LlamaToolCallParsing` suite is deliberately not gated on `LLAMA_MODEL_PATH` and runs with no model present. 22 cases cover every supported format (Hermes ``, Llama 3.x `<|python_tag|>` including `;`-separated multi-calls, Mistral `[TOOL_CALLS]`, ``), escaped braces inside string literals, arguments emitted as a JSON string, OpenAI-style function nesting, bare-JSON acceptance only when the name is a registered tool, and the eight streaming holdback cases that keep tool markup from leaking into streamed text. --- .../Models/LlamaLanguageModel.swift | 21 ++- .../LlamaLanguageModelTests.swift | 175 ++++++++++++++++++ 2 files changed, 188 insertions(+), 8 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 81b557c9..b207276a 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -1906,9 +1906,14 @@ import Foundation } // MARK: - Tool Call Parsing + // + // The pure parsing helpers below are `internal` rather than file-private so that + // `LlamaLanguageModelTests` can unit-test them without a local GGUF model. They are the only + // substantial new logic here that can be exercised in CI. Nothing becomes public API, and the + // tool *invocation* helpers further down stay file-private since they need a live session. /// A tool call recovered from the model's generated text. - private struct ParsedLlamaToolCall { + internal struct ParsedLlamaToolCall { let name: String let arguments: GeneratedContent } @@ -1930,13 +1935,13 @@ import Foundation /// /// Any other family-specific syntax is *not* supported and its text will be returned to the /// caller as ordinary model output. - private let llamaToolCallMarkers = ["", "<|python_tag|>", "[TOOL_CALLS]", "", "<|python_tag|>", "[TOOL_CALLS]", " Range? { + internal func llamaBalancedJSONRange(in text: String, from start: String.Index) -> Range? { var index = start while index < text.endIndex, text[index].isWhitespace { index = text.index(after: index) @@ -1983,7 +1988,7 @@ import Foundation } /// Decodes one tool call from a already-parsed JSON value. - private func llamaToolCall(fromObject object: Any, defaultName: String? = nil) -> ParsedLlamaToolCall? { + internal func llamaToolCall(fromObject object: Any, defaultName: String? = nil) -> ParsedLlamaToolCall? { guard var dictionary = object as? [String: Any] else { return nil } // OpenAI-shaped nesting: {"type": "function", "function": {"name": ..., "arguments": ...}} @@ -2011,7 +2016,7 @@ import Foundation } /// Decodes one or more tool calls from a JSON object or array. - private func llamaToolCalls(fromJSON json: String, defaultName: String? = nil) -> [ParsedLlamaToolCall] { + internal func llamaToolCalls(fromJSON json: String, defaultName: String? = nil) -> [ParsedLlamaToolCall] { guard let object = try? JSONSerialization.jsonObject(with: Data(json.utf8)) else { return [] } if let array = object as? [Any] { return array.compactMap { llamaToolCall(fromObject: $0, defaultName: defaultName) } @@ -2022,7 +2027,7 @@ import Foundation /// Splits generated text into the prose a caller should see and the tool calls embedded in it. /// /// See ``llamaToolCallMarkers`` for the exact set of supported formats. - private func llamaSplitToolCalls( + internal func llamaSplitToolCalls( from text: String, knownToolNames: Set ) -> (visibleText: String, toolCalls: [ParsedLlamaToolCall]) { @@ -2121,7 +2126,7 @@ import Foundation } /// Advances past `token` if it is the next non-whitespace content, otherwise returns `index`. - private func llamaConsuming(_ token: String, in text: String, from index: String.Index) -> String.Index { + internal func llamaConsuming(_ token: String, in text: String, from index: String.Index) -> String.Index { var probe = index while probe < text.endIndex, text[probe].isWhitespace { probe = text.index(after: probe) @@ -2134,7 +2139,7 @@ import Foundation /// /// Tool-call markup arrives inline with prose, so streaming must withhold any text that has /// begun — or might still turn into — a tool call. - private func llamaStreamableVisiblePrefix(of text: String) -> String { + internal func llamaStreamableVisiblePrefix(of text: String) -> String { // A response opening with a JSON container may be a bare-JSON tool call, which cannot be // recognized until generation finishes. Withhold all of it until then. let leading = text.drop(while: { $0.isWhitespace }) diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index 0faa3020..20f7da64 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -482,4 +482,179 @@ import Testing #expect(leakedMarkup.isEmpty, "Tool-call markup leaked into streamed response text.") } } + + /// Unit tests for the tool-call text parser. + /// + /// Deliberately **not** gated on `LLAMA_MODEL_PATH`: local models express tool calls as + /// generated text in family-specific formats, so this parser is the substance of Llama tool + /// calling, and it must be verifiable with no model present. + @Suite("LlamaToolCallParsing") + struct LlamaToolCallParsingTests { + /// The single registered tool used by these cases. + static let knownTools: Set = ["getWeather"] + + private func split( + _ text: String, + tools: Set = LlamaToolCallParsingTests.knownTools + ) -> (names: [String], arguments: [String], visible: String) { + let result = llamaSplitToolCalls(from: text, knownToolNames: tools) + return ( + result.toolCalls.map(\.name), + result.toolCalls.map { $0.arguments.jsonString }, + result.visibleText + ) + } + + // MARK: Supported formats + + @Test func hermesToolCall() { + let result = split( + "\n{\"name\": \"getWeather\", \"arguments\": {\"city\": \"SF\"}}\n" + ) + #expect(result.names == ["getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}"]) + #expect(result.visible == "") + } + + @Test func hermesToolCallKeepsSurroundingProse() { + let result = split( + "Let me check.\n{\"name\":\"getWeather\",\"arguments\":{\"city\":\"SF\"}}" + ) + #expect(result.names == ["getWeather"]) + #expect(result.visible == "Let me check.") + } + + @Test func multipleHermesToolCalls() { + let result = split( + "{\"name\":\"getWeather\",\"arguments\":{\"city\":\"SF\"}}" + + "{\"name\":\"getWeather\",\"arguments\":{\"city\":\"NY\"}}" + ) + #expect(result.names == ["getWeather", "getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}", "{\"city\":\"NY\"}"]) + #expect(result.visible == "") + } + + @Test func llamaPythonTagToolCall() { + let result = split("<|python_tag|>{\"name\": \"getWeather\", \"parameters\": {\"city\": \"SF\"}}") + #expect(result.names == ["getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}"]) + #expect(result.visible == "") + } + + @Test func llamaPythonTagSemicolonSeparatedCalls() { + let result = split( + "<|python_tag|>{\"name\":\"getWeather\",\"parameters\":{\"city\":\"SF\"}}; " + + "{\"name\":\"getWeather\",\"parameters\":{\"city\":\"NY\"}}<|eom_id|>" + ) + #expect(result.names == ["getWeather", "getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}", "{\"city\":\"NY\"}"]) + #expect(result.visible == "") + } + + @Test func mistralToolCallsArray() { + let result = split("[TOOL_CALLS] [{\"name\": \"getWeather\", \"arguments\": {\"city\": \"SF\"}}]") + #expect(result.names == ["getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}"]) + #expect(result.visible == "") + } + + @Test func functionEqualsToolCall() { + // Here the JSON payload *is* the arguments; the name comes from the tag. + let result = split("{\"city\": \"SF\"}") + #expect(result.names == ["getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}"]) + #expect(result.visible == "") + } + + // MARK: Argument shapes + + @Test func bracesInsideStringLiteralsDoNotTerminateScan() { + let result = split("{\"name\":\"getWeather\",\"arguments\":{\"city\":\"a}b\"}}") + #expect(result.names == ["getWeather"]) + #expect(result.arguments == ["{\"city\":\"a}b\"}"]) + #expect(result.visible == "") + } + + @Test func openAIStyleFunctionNesting() { + let result = split( + "{\"type\":\"function\",\"function\":" + + "{\"name\":\"getWeather\",\"arguments\":{\"city\":\"SF\"}}}" + ) + #expect(result.names == ["getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}"]) + } + + @Test func stringifiedArguments() { + let result = split( + "{\"name\":\"getWeather\",\"arguments\":\"{\\\"city\\\":\\\"SF\\\"}\"}" + ) + #expect(result.names == ["getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}"]) + } + + // MARK: Negative cases + + @Test func bareJSONIsAToolCallWhenNameIsRegistered() { + let result = split("{\"name\": \"getWeather\", \"parameters\": {\"city\": \"SF\"}}") + #expect(result.names == ["getWeather"]) + #expect(result.arguments == ["{\"city\":\"SF\"}"]) + #expect(result.visible == "") + } + + @Test func bareJSONStaysTextWhenNameIsNotRegistered() { + // An ordinary JSON answer must not be swallowed as a tool call. + let text = "{\"name\": \"other\", \"parameters\": {}}" + let result = split(text) + #expect(result.names.isEmpty) + #expect(result.visible == text) + } + + @Test func plainProseIsUntouched() { + let result = split("The weather is sunny.") + #expect(result.names.isEmpty) + #expect(result.visible == "The weather is sunny.") + } + + @Test func markerWithoutDecodablePayloadYieldsNoCalls() { + // The stray marker is dropped rather than surfaced, and no call is invented. + let result = split(" oops") + #expect(result.names.isEmpty) + #expect(result.visible == "oops") + } + + // MARK: Streaming holdback + + @Test func streamingEmitsPlainTextImmediately() { + #expect(llamaStreamableVisiblePrefix(of: "Hello there") == "Hello there") + } + + @Test func streamingWithholdsPartialMarker() { + #expect(llamaStreamableVisiblePrefix(of: "Hello {\"x\":1}") == "Hello ") + } + + @Test func streamingWithholdsPartialPythonTag() { + #expect(llamaStreamableVisiblePrefix(of: "Sure <|python_") == "Sure ") + } + + @Test func streamingWithholdsLeadingJSON() { + // A bare-JSON tool call is unrecognizable until generation finishes. + #expect(llamaStreamableVisiblePrefix(of: "{\"name\":") == "") + } + + @Test func streamingWithholdsLeadingJSONAfterWhitespace() { + #expect(llamaStreamableVisiblePrefix(of: " \n{\"na") == "") + } + + @Test func streamingWithholdsLeadingBracket() { + #expect(llamaStreamableVisiblePrefix(of: "[TOOL_C") == "") + } + + @Test func streamingDoesNotTreatLoneAngleBracketAsMarker() { + #expect(llamaStreamableVisiblePrefix(of: "5 < 3 is false") == "5 < 3 is false") + } + } #endif // Llama From 585f25c6d6edd89602072028032298754588f785 Mon Sep 17 00:00:00 2001 From: Claude Opus 5 Date: Mon, 3 Aug 2026 20:58:48 -0400 Subject: [PATCH 4/5] Reuse the KV cache prefix across tool iterations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each tool iteration re-rendered the whole conversation and cleared the KV cache, so every turn re-decoded the prompt, the instructions and all earlier turns from position 0 — the work grew with the conversation even though each iteration only appended a message or two. The cache now keeps the tokens it shares with the next prompt and trims only the divergent tail, so a follow-up turn decodes roughly the newly appended messages instead of everything. The guards mirror what `llama-server` checks before reusing a prefix, because the same conditions make reuse unsound here: recurrent models (Mamba, RWKV) carry a rolling state with no per-token prefix to keep; encoder models hold state these tokens don't describe; a sliding-window cache that has evicted the head reports a non-zero `pos_min`, so position 0 is already gone; and `llama_memory_seq_rm` returns false when a partial removal isn't supported for the memory type. Any of those falls back to discarding the sequence and decoding from scratch. At least one token is always re-decoded so the batch still produces logits to sample from. --- .../Models/LlamaLanguageModel.swift | 127 +++++++++++++++--- 1 file changed, 111 insertions(+), 16 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index b207276a..d2a998d5 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -706,15 +706,13 @@ import Foundation ) } + // Carried across tool iterations so each re-render keeps the prefix it + // shares with the last one instead of decoding the prompt again. + let kvCache = KVCacheState() + while true { try Task.checkCancellation() - if toolIteration > 0, let memory = llama_get_memory(context) { - // Drop the previous iteration's KV cache before re-decoding - // the extended prompt from position 0. - llama_memory_clear(memory, true) - } - let iterationPrompt = try self.renderPrompt(messages) var rawText = "" var lastVisible = "" @@ -724,7 +722,8 @@ import Foundation model: model!, prompt: iterationPrompt, maxTokens: maxTokens, - options: runtimeOptions + options: runtimeOptions, + cache: kvCache ) { if Task.isCancelled { break } rawText += tokenText @@ -1320,12 +1319,86 @@ import Foundation } } + /// The tokens currently resident in a context's KV cache. + /// + /// A tool-calling turn re-renders the whole conversation each iteration, and every + /// iteration shares a long prefix with the last one — the instructions, the user's + /// prompt, and every earlier turn. Carrying the decoded tokens between iterations lets + /// the cache be trimmed back to that shared prefix rather than cleared, so only the + /// newly appended messages are decoded. + private final class KVCacheState { + var tokens: [llama_token] = [] + } + + /// The number of leading tokens two sequences share. + private func commonPrefixLength(_ lhs: [llama_token], _ rhs: [llama_token]) -> Int { + var count = 0 + let limit = min(lhs.count, rhs.count) + while count < limit, lhs[count] == rhs[count] { + count += 1 + } + return count + } + + /// How much of `promptTokens` is already resident and can be left in place, trimming the + /// rest of the sequence away. + /// + /// This mirrors the checks `llama-server` makes before reusing a cached prefix, because + /// the same conditions make reuse unsound here: + /// + /// - Recurrent models (Mamba, RWKV) carry a rolling state rather than per-token entries, + /// so there is no prefix to keep. + /// - Encoder models hold state that these tokens don't describe. + /// - A sliding-window cache evicts the head of the sequence, so a non-zero `pos_min` + /// means position 0 is already gone and the prefix can't be trusted. + /// - `llama_memory_seq_rm` reports `false` when a partial removal isn't supported, in + /// which case the whole sequence has to go. + /// + /// Returns the number of leading tokens left resident; the caller decodes from there. + private func reusableCachePrefix( + cache: KVCacheState?, + promptTokens: [llama_token], + model: OpaquePointer, + context: OpaquePointer + ) -> Int { + guard let memory = llama_get_memory(context) else { return 0 } + + func discardEverything() -> Int { + llama_memory_seq_rm(memory, 0, -1, -1) + cache?.tokens.removeAll() + return 0 + } + + guard let cache, !cache.tokens.isEmpty else { return discardEverything() } + guard !llama_model_is_recurrent(model), !llama_model_has_encoder(model) else { + return discardEverything() + } + + // A sliding-window cache that has already evicted the head can't anchor a prefix. + guard llama_memory_seq_pos_min(memory, 0) == 0 else { return discardEverything() } + + // Re-decode at least one token so the batch produces logits to sample from. + let shared = min( + commonPrefixLength(cache.tokens, promptTokens), + promptTokens.count - 1 + ) + guard shared > 0 else { return discardEverything() } + + guard llama_memory_seq_rm(memory, 0, Int32(shared), -1) else { + // Partial removal unsupported for this memory type. + return discardEverything() + } + + return shared + } + private func generateTextStream( context: OpaquePointer, model: OpaquePointer, prompt: String, maxTokens: Int, - options: ResolvedGenerationOptions + options: ResolvedGenerationOptions, + cache: KVCacheState? = nil ) -> AsyncThrowingStream { return AsyncThrowingStream { continuation in self.performTextGeneration( @@ -1334,6 +1407,7 @@ import Foundation prompt: prompt, maxTokens: maxTokens, options: options, + cache: cache, continuation: continuation ) } @@ -1345,6 +1419,7 @@ import Foundation prompt: String, maxTokens: Int, options: ResolvedGenerationOptions, + cache: KVCacheState? = nil, continuation: AsyncThrowingStream.Continuation ) { do { @@ -1364,15 +1439,26 @@ import Foundation var batch = llama_batch_init(Int32(options.batchSize), 0, 1) defer { llama_batch_free(batch) } + let reusedPrefix = reusableCachePrefix( + cache: cache, + promptTokens: promptTokens, + model: model, + context: context + ) + let hasEncoder = try prepareInitialBatch( batch: &batch, promptTokens: promptTokens, + reusedPrefix: reusedPrefix, model: model, vocab: vocab, context: context, batchSize: options.batchSize ) + // Everything through the prompt is now resident, whatever was reused. + cache?.tokens = promptTokens + // Initialize sampler chain with options guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { throw LlamaLanguageModelError.decodingFailed @@ -1407,7 +1493,7 @@ import Foundation // Generate tokens one by one // Track position - for encoder-decoder models, we start from position 1 (after decoder start token) // For decoder-only models, we continue from the end of the prompt - var n_cur: Int32 = hasEncoder ? 1 : batch.n_tokens + var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) for _ in 0 ..< maxTokens { // Sample next token from logits of the last token we just decoded @@ -1424,6 +1510,9 @@ import Foundation continuation.yield(tokenText) } + // The sampled token is about to be decoded, so it joins the resident prefix. + cache?.tokens.append(nextToken) + // Prepare batch for next token batch.n_tokens = 1 batch.token[0] = nextToken @@ -1480,13 +1569,18 @@ import Foundation private func prepareInitialBatch( batch: inout llama_batch, promptTokens: [llama_token], + reusedPrefix: Int = 0, model: OpaquePointer, vocab: OpaquePointer, context: OpaquePointer, batchSize: UInt32 ) throws -> Bool { - // Validate that prompt token count doesn't exceed batch capacity to prevent buffer overflow - guard promptTokens.count <= batchSize else { + // Tokens up to `reusedPrefix` are already in the KV cache, so only the remainder is + // decoded — but it still occupies positions `reusedPrefix.. Date: Mon, 3 Aug 2026 20:58:56 -0400 Subject: [PATCH 5/5] Replay tool results under a tool role when the template has one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tool results were always replayed as `user` messages. Models trained for tool calling expect them under a `tool` role, and the other local providers already do this — MLX and Core ML both append tool-role messages — so Llama was the odd one out. `llama_chat_apply_template` implements a fixed set of built-in templates: ChatML, Qwen and Hermes render a `tool` role explicitly, while Llama 2 and Mistral fold any unrecognized role into the user turn, where the role name would leak into the prompt. Rather than hardcode which templates qualify, the same content is rendered under both roles and compared — a template that distinguishes them produces different output — and the result is memoized. Templates without tool support keep the previous `user` behavior. --- .../Models/LlamaLanguageModel.swift | 39 +++++++++++++++++-- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index d2a998d5..10993f3a 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -415,6 +415,9 @@ import Foundation /// The loaded model instance private var model: OpaquePointer? + /// Memoized result of probing the chat template for `tool` role support. + private var cachedToolResultRole: String? + /// The model's vocabulary private var vocab: OpaquePointer? @@ -1668,9 +1671,8 @@ import Foundation /// /// Tool activity is rendered back into the text formats this provider emits and parses, /// so that a follow-up turn on a session that already used tools sees a coherent history. - /// Tool results are replayed with the `user` role rather than a `tool` role because - /// `llama_chat_apply_template` only implements a fixed set of built-in templates, and - /// several of them (Llama 2, Mistral) do not recognize a `tool` role at all. + /// Tool results are replayed under whichever role the model's template understands — + /// see ``toolResultRole``. private func chatMessages( for session: LanguageModelSession, extraSystemMessage: String? = nil @@ -1704,7 +1706,7 @@ import Foundation } case .toolOutput(let output): - messages.append(("user", toolResponseMarkup(output))) + messages.append((toolResultRole, toolResponseMarkup(output))) } } @@ -1715,6 +1717,35 @@ import Foundation return mergingConsecutiveRoles(messages) } + /// The role to replay tool results under. + /// + /// `llama_chat_apply_template` implements a fixed set of built-in templates. Some of them + /// (ChatML, Qwen, Hermes) render a `tool` role explicitly, which is what a tool-trained + /// model expects to see; others (Llama 2, Mistral) don't recognize it and fold any + /// unknown role into the user turn, where the raw role name would leak into the prompt. + /// + /// Rather than hardcode which templates support it, this renders the same content under + /// both roles and compares: a template that distinguishes them produces different output. + private var toolResultRole: String { + if let cached = cachedToolResultRole { + return cached + } + let role = templateDistinguishesToolRole() ? "tool" : "user" + cachedToolResultRole = role + return role + } + + private func templateDistinguishesToolRole() -> Bool { + let probe = "__tool_role_probe__" + guard + let asTool = try? renderPrompt([("user", "probe"), ("tool", probe)]), + let asUser = try? renderPrompt([("user", "probe"), ("user", probe)]) + else { + return false + } + return asTool != asUser + } + /// Collapses runs of same-role messages into one. /// /// Several built-in llama.cpp templates assume strictly alternating user/assistant turns,