diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index 90e28de..10993f3 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? @@ -505,32 +508,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 +683,142 @@ 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) + ) + ) + } + + // 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() + + 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 + options: runtimeOptions, + cache: kvCache ) { - 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 + ) - let snapshot = LanguageModelSession.ResponseStream.Snapshot( - content: (accumulatedText as! Content).asPartiallyGenerated(), - rawContent: GeneratedContent(accumulatedText) + // 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 signature = Self.toolCallSignature(for: split.toolCalls) + if signature == previousToolCallSignature { + session.appendTranscriptEntry( + .toolCalls( + Transcript.ToolCalls( + makeTranscriptToolCalls(from: split.toolCalls) + ) + ) + ) + throw Self.repeatedToolCallLoopError() + } + 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") + ) ) - continuation.yield(snapshot) + messages = self.mergingConsecutiveRoles(messages) + continue } - } catch { - continuation.finish(throwing: error) - return + + break } continuation.finish() @@ -1124,12 +1322,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( @@ -1138,6 +1410,7 @@ import Foundation prompt: prompt, maxTokens: maxTokens, options: options, + cache: cache, continuation: continuation ) } @@ -1149,6 +1422,7 @@ import Foundation prompt: String, maxTokens: Int, options: ResolvedGenerationOptions, + cache: KVCacheState? = nil, continuation: AsyncThrowingStream.Continuation ) { do { @@ -1168,15 +1442,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 @@ -1211,7 +1496,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 @@ -1228,6 +1513,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 @@ -1284,13 +1572,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.. 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 under whichever role the model's template understands — + /// see ``toolResultRole``. + 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 +1699,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((toolResultRole, toolResponseMarkup(output))) } } @@ -1405,6 +1714,62 @@ import Foundation messages.append(("system", extraSystemMessage)) } + 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, + /// 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 +1822,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 +2031,376 @@ 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. + internal 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. + internal 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. + 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": ...}} + 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. + 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) } + } + 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. + internal 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`. + 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) + } + 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. + 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 }) + 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 diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index 3d9e32c..20f7da6 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -401,5 +401,260 @@ 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.") + } + } + + /// 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