Skip to content

Add tool calling to the Llama provider - #10

Open
ActuallyTaylor wants to merge 5 commits into
mainfrom
feature/llama-streaming-tool-calling
Open

Add tool calling to the Llama provider#10
ActuallyTaylor wants to merge 5 commits into
mainfrom
feature/llama-streaming-tool-calling

Conversation

@ActuallyTaylor

Copy link
Copy Markdown
Collaborator

LlamaLanguageModel ignored session.tools entirely — no tool calling in respond or streamResponse.

What the binding actually supports: nothing. LlamaSwift re-exports the raw llama.cpp C API. llama_chat_message is just { role, content }, and llama_chat_apply_template has no tools parameter and explicitly does not use a Jinja parser. llama_model_chat_template() returns the GGUF's template text, but nothing in the binding can render it. llama.cpp's tool-aware path lives in its common library, which isn't in the xcframework. So tools cannot be injected through the model's own chat template.

Tools are therefore advertised through a synthesized system message carrying each tool's name, description and JSON Schema. Parsing is deliberately more permissive than emission, covering Hermes/Qwen <tool_call>, Llama 3.x <|python_tag|> (including ;-separated calls), Mistral [TOOL_CALLS], <function=name>, and bare leading JSON — the last only when every parsed name matches a registered tool, so ordinary JSON answers aren't swallowed. Brace matching is string- and escape-aware.

Both paths honor the full delegate contract and the transcript ordering, with an 8-iteration cap and repeated-signature detection. Because tool markup arrives inline with prose, streaming withholds text that has begun or could still grow into a marker.

Also fixes a pre-existing break: the Llama trait did not compile at all. use_mmap/use_mlock were replaced upstream by a load_mode enum, and Package.resolved had no llama.swift entry, so the trait had never been resolved. This is the same fix already pushed to #3.

22 ungated parser unit tests run with no model present — the only new logic in this series that CI actually exercises. Verified live against Qwen2.5-3B-Instruct Q4_K_M: 54/54 pass including withTools and streamWithTools.

Tool calling applies to String generation only, matching MLX.

Based on #4; merge that first.

🤖 Generated with Claude Code

@ActuallyTaylor
ActuallyTaylor force-pushed the feature/central-tool-resolution branch from 3f72c2c to 6925e25 Compare August 3, 2026 19:35
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from cb9ba2d to 9e93327 Compare August 3, 2026 19:36
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/central-tool-resolution branch from 6925e25 to 28e5c10 Compare August 3, 2026 19:39
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch 2 times, most recently from 1e204cc to f0d25d1 Compare August 3, 2026 19:47
@ActuallyTaylor
ActuallyTaylor changed the base branch from feature/central-tool-resolution to feature/mlx-streaming-tool-calling August 3, 2026 19:48
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/mlx-streaming-tool-calling branch from cfc98fb to a5a0ea7 Compare August 3, 2026 20:03
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from f0d25d1 to df732ac Compare August 3, 2026 20:03
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/mlx-streaming-tool-calling branch from a5a0ea7 to 043c0ee Compare August 3, 2026 20:07
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from df732ac to 6cbb803 Compare August 3, 2026 20:07
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/mlx-streaming-tool-calling branch from 043c0ee to 66176ee Compare August 3, 2026 20:09
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from 6cbb803 to 7d02c02 Compare August 3, 2026 20:09
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/mlx-streaming-tool-calling branch from 66176ee to f2fb47c Compare August 3, 2026 20:17
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from 7d02c02 to 2cc5490 Compare August 3, 2026 20:17
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/mlx-streaming-tool-calling branch from f2fb47c to 9500897 Compare August 3, 2026 20:18
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from 2cc5490 to 3179cd5 Compare August 3, 2026 20:18
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/mlx-streaming-tool-calling branch from 9500897 to 8f4d598 Compare August 3, 2026 20:40
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from 3179cd5 to 84269fc Compare August 3, 2026 20:40
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/mlx-streaming-tool-calling branch from 8f4d598 to befc05d Compare August 3, 2026 20:44
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from 84269fc to 1c578c9 Compare August 3, 2026 20:44
Base automatically changed from feature/mlx-streaming-tool-calling to main August 3, 2026 22:23
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from 1c578c9 to fe249eb Compare August 3, 2026 22:23
claude added 3 commits August 3, 2026 20:10
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 `<tool_call>` format.

Parsing is deliberately more permissive than emission, and its scope is explicit: `<tool_call>{...}</tool_call>` (Hermes/Qwen/ChatML), `<|python_tag|>{...}` including `;`-separated multi-calls (Llama 3.x), `[TOOL_CALLS] [{...}]` (Mistral), and `<function=name>{...}</function>`. 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 `<tool_call>` / `<tool_response>` 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.
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.
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 `<tool_call>`, Llama 3.x `<|python_tag|>` including `;`-separated multi-calls, Mistral `[TOOL_CALLS]`, `<function=name>`), 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.
@ActuallyTaylor
ActuallyTaylor force-pushed the feature/llama-streaming-tool-calling branch from fe249eb to 9bb9d01 Compare August 4, 2026 00:10
claude added 2 commits August 3, 2026 20:58
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.
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.
@ActuallyTaylor

Copy link
Copy Markdown
Collaborator Author

Both review points addressed, in two commits.

KV cache should be trimmed back, not dropped

585f25c. Each tool iteration cleared the cache and re-decoded the entire conversation from position 0, so the work grew with the conversation even though an iteration only appends a message or two. The cache now keeps the prefix it shares with the next prompt and trims only the divergent tail.

Per the request to check this against llama-server, I read the llama.h contracts rather than assuming, and the first draft was unsound in three ways that the guards now cover:

  • llama_memory_seq_rm returns false when a partial removal isn't supported — the header says so explicitly. My first version ignored the return value, which would have left the cache inconsistent with the assumed prefix on recurrent models. It now falls back to discarding the sequence.
  • Sliding-window caches evict the head, so llama_memory_seq_pos_min is non-zero and position 0 is already gone; a prefix anchored at 0 can't be trusted.
  • Recurrent models (llama_model_is_recurrent — Mamba, RWKV) carry a rolling state with no per-token prefix, and encoder models hold state these tokens don't describe.

Any of those falls back to a full clear and a decode from scratch. At least one token is always re-decoded so the batch still produces logits to sample from.

Use a tool role where the template supports one

d22b9a2. Tool results were always replayed as user messages; this also brings the provider in line with MLX and Core ML, which both append tool-role messages.

llama_chat_apply_template implements a fixed set of built-in templates — ChatML, Qwen and Hermes render a tool role, while Llama 2 and Mistral fold unknown roles 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. The result is memoized, and templates without tool support keep the previous behavior.

Verification, including something worth flagging

Build and lint clean; the 22 ungated parser tests and withTools pass against Qwen2.5-3B.

streamWithTools is flaky against this model, and was already flaky before these changes — measured, not assumed: 1/3 on the pre-change baseline vs 4/5 with these commits. The earlier "54/54 passing" run was luckier than the code deserved, and I shouldn't have reported it without repeat runs. The change looks neutral-to-better on reliability, but the sample is small and the underlying flakiness is a separate problem: the 3B model doesn't always emit parseable tool markup in the streaming path even though the non-streaming path is consistent.

🤖 Generated with Claude Code

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants