Support tool calling in MLXLanguageModel.streamResponse() - #181
Conversation
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>
|
Hi @noorbhatia. Thanks for this, and sorry for not reviewing it sooner. This and #176 each add a way to observe transcript entries during streaming, for MLX and Anthropic respectively, which is what I asked for in #103. Rather than land two shapes, I'd like to settle on one in #103 and then have both PRs adopt it. I've posted a proposal there that's essentially this PR's shape; would you take a look? |
|
I like the implementation! It seems to match with the rest of the streaming tool calls |
There was a problem hiding this comment.
🟡 Changes recommended
Stream completion can race ahead of transcript finalization, causing subsequent turns to observe incomplete history.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds tool-calling support to MLX streaming while preserving cumulative response snapshots.
Changes:
- Executes tool calls across streaming generation rounds.
- Exposes streamed tool transcript entries.
- Adds MLX streaming tool coverage.
File summaries
| File | Description |
|---|---|
MLXLanguageModel.swift |
Implements streaming tool-call cycles. |
LanguageModelSession.swift |
Propagates streamed transcript entries. |
MLXLanguageModelTests.swift |
Tests streamed tool execution. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| ) | ||
| ) | ||
| session.withMutation(keyPath: \.transcript) { | ||
| session.state.withLock { $0.transcript.append(responseEntry) } | ||
| session.state.withLock { | ||
| $0.transcript.append(contentsOf: lastSnapshot.transcriptEntries) |
|
Thanks, @noorbhatia, and sorry this sat so long. This is the #103 shape exactly, so merging as-is. Streaming tool calls for the other providers can follow the same pattern. One small follow-up on my side: |
…ted (#224) wrapStream reported completion before appending the response (and, since #181, the streamed tool entries) to the session transcript, so a caller that drains the stream and immediately starts the next turn could build its prompt from an incomplete transcript. Commit the entries and clear isResponding first, then finish the stream, on both the success and the error path.
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 #164