Add tool calling to LlamaLanguageModel, including streaming - #216
james-333i wants to merge 6 commits into
Conversation
6d2ac2c to
d990360
Compare
Image segments threw unsupportedFeature because the backend had no multimodal path, even though the prebuilt llama.cpp binaries ship the mtmd library and its helpers. Accept an mmprojPath at initialization and load the projector next to the model. When a projector is present, prompt formatting replaces each image segment with the mtmd media marker and collects payloads in order, then generation tokenizes the marker-annotated prompt with mtmd_tokenize and evaluates text and image chunks through mtmd_helper_eval_chunks before sampling continues from the resulting position. Both respond and streaming support images, and models without a projector keep rejecting image input. Adds live tests generating from an embedded test image through both paths.
Gemma 4's canonical chat template no longer contains the start_of_turn marker that llama_chat_apply_template keys its Gemma detection on, so formatting threw encodingFailed for every Gemma 4 GGUF. When template application fails and the embedded template carries the Gemma 4 turn syntax, render it directly: turns open with a turn marker and role, close with the reverse marker, the assistant role is named model, and generation opens a model turn. The BOS token is applied during tokenization, and thinking is opt-in in this format so no suppression is needed.
Every generation created a fresh llama_context and prefilled the full rendered conversation from token zero, so multi-turn chat cost grew with the square of the transcript and long conversations spent most of their time re-decoding history. Keep one context alive per session for plain chat generations. Each exchange tokenizes the rendered prompt, keeps the longest token prefix shared with the context's recorded state, removes diverged state with llama_memory_seq_rm, and decodes only the remainder. Backends that cannot rewind, such as recurrent models, fall back to clearing memory and decoding the full prompt, and appends need no rewind on any backend. The final prompt token is always re-decoded so sampling has fresh logits, generated tokens extend the recorded state as they decode, and any generation error discards the cached context. Structured generation, image prompts, and encoder models keep single-use contexts, and clearCachedContext lets consumers free the cached state under memory pressure. Adds a live test asserting prefix reuse on the second turn of a session.
llama_chat_apply_template has no parameter for tool definitions, so tool support is implemented at the prompt layer. The tool syntax is detected from the model's embedded chat template: Hermes-style JSON (Qwen 2.5/3 and most ChatML fine-tunes), Qwen 3.5's XML function/parameter form, and the Gemma 4 canonical format with its token-quoted argument notation. Definitions are rendered into the system prompt following each template's own wording and placement, past tool turns are replayed in the native markup (including Gemma 4's open-model-turn continuation), and calls are parsed back out of generated text and run through the resolve-and-continue loop used by the MLX and Ollama backends, with the same delegate hooks, iteration cap, and repeated-signature guard. Generation stops early once a complete tool-call block is produced. Non-streaming respond() only; streamResponse() ignores session tools as before.
streamResponse() hardcoded tools: nil and discarded .toolCall stream items, so MLX callers had to choose between streamed tokens and tool calling. respond() already ran the full tool cycle; this ports that while-loop into the streaming path, reusing mlxToolSpecs, resolveToolCalls, makeTranscriptToolCalls, and the maxToolIterations / repeated-signature guards. Text and tool entries accumulate across rounds so snapshots stay monotonic. Also surface streamed tool activity: ResponseStream.Snapshot gains a defaulted transcriptEntries field (ArraySlice<Transcript.Entry>), wrapStream appends it to the session transcript before the response entry, and collect() returns it instead of []. The field defaults to empty, so the other providers keep their current behavior; the shared plumbing is ready for them to populate later. Closes huggingface#164 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
streamResponse() now runs the same resolve-and-continue loop as respond(), yielding snapshots that carry the cumulative visible text and the tool-call and tool-output entries produced so far. Because llama tool calls arrive as text rather than parsed events, snapshots withhold any trailing partial match of a call-start marker until the next token confirms or breaks it, so markup never appears mid-stream. Gemma 4 emits thought-channel spans without being asked: thinking is opt-in via a system-turn token this backend never injects, and the canonical template ships a strip_thinking macro for consumers. Both respond() and streamResponse() now remove completed spans and withhold unclosed ones, recognizing the canonical marker spelling and the variant observed from deployed quantizations.
d990360 to
436f229
Compare
There was a problem hiding this comment.
🟡 Changes recommended
Response accumulation, final streaming visibility, and Gemma BOS handling contain correctness issues, while multimodal scope is undocumented.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds prompt-layer tool calling and streaming support for Llama models, alongside MLX streaming integration and additional context caching and multimodal functionality.
Changes:
- Implements Hermes, Qwen XML, and Gemma tool-call formatting and parsing.
- Adds tool-aware streaming snapshots and transcript persistence.
- Adds Llama context reuse, Gemma rendering, and vision support.
File summaries
| File | Description |
|---|---|
MLXLanguageModelTests.swift |
Tests streamed MLX tool execution. |
LlamaToolCallFormatTests.swift |
Tests tool formats and Llama tool loops. |
LlamaLanguageModelTests.swift |
Tests context reuse and vision. |
LlamaGemma4TemplateTests.swift |
Tests manual Gemma prompt rendering. |
MLXLanguageModel.swift |
Adds streaming tool-resolution loops. |
LlamaToolCallFormat.swift |
Implements tool syntax rendering and parsing. |
LlamaLanguageModel.swift |
Adds tools, caching, Gemma rendering, and vision. |
LanguageModelSession.swift |
Persists streamed transcript entries. |
Review details
- Files reviewed: 8/8 changed files
- Comments generated: 4
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| messages: [(role: String, content: String)], | ||
| assistantPrefill: String? | ||
| ) -> String { | ||
| var rendered = "" |
| let (visibleText, parsedCalls) = format.parseToolCalls(in: accumulated) | ||
| if parsedCalls.isEmpty { | ||
| text = visibleText | ||
| break generationLoop | ||
| } |
| let roundVisible = outputFormat.streamingVisibleText( | ||
| in: roundRaw, | ||
| withholdToolCalls: withholdToolCalls | ||
| ) |
| /// The path to the multimodal projector GGUF file, when the model has one. | ||
| /// | ||
| /// Prompts may include image segments only when a projector is loaded. | ||
| public let mmprojPath: String? |
|
This landed on |
llama_chat_apply_template has no parameter for tool definitions, so tool support is implemented at the prompt layer with the syntax detected from the model's embedded template: Hermes-style JSON, Qwen XML, and the Gemma 4 canonical format. Calls run through the same resolve-and-continue loop as the MLX and Ollama backends. Streaming runs the full loop and withholds partial call markup until the next token confirms or breaks it. This includes the commit from #181, whose Snapshot.transcriptEntries field the streaming path requires, so #181 should merge first. Stacks on the session-context PR.