Adoption-ready long-document reasoning for MLX Swift apps#7
Merged
joelnishanth merged 19 commits intoApr 27, 2026
Merged
Conversation
- New hero: Swift-native long-document reasoning for private, on-device MLX apps - Add "How This Fits" section explaining relationship to MLX Swift / MLX Swift LM - Add "Target Use Cases" section (transcripts, markdown, logs, offline docs) - Fix Requirements: Swift 6.1+ (was 5.9+, but Package.swift requires 6.1) Made-with: Cursor
- Add ChainExecutionOptions for controlling map/reduce behavior - Update DocumentChain protocol with stuffPrompt and options params - Backward-compatible overloads preserve existing callers - AdaptiveChain now accounts for prompt + text + reserved output in budget check, not just raw text word count - StuffChain uses stuffPrompt when provided, falls back to reducePrompt - ContextBudget gains fitsInBudget helper method - Add tests for prompt overhead, stuffPrompt, and reserved output Made-with: Cursor
TranscriptChunker: - Splits at speaker-turn boundaries, not sentence boundaries - Recognizes HH:MM:SS Speaker: and Speaker: patterns - overlapTurns for context continuity across chunks - Populates metadata with speaker labels and timestamps MarkdownHeadingChunker: - Splits at # heading boundaries up to configurable level - Preserves heading text in each chunk - Falls back to sentence-aware chunking for oversized sections LogChunker: - Splits at ISO 8601 and common log timestamp boundaries - Detects stack traces (indented at/frame lines) and keeps them intact - Extracts timestamps into chunk metadata These three chunkers are the adoption wedge for the library, targeting meeting transcripts, document processing, and log analysis. Made-with: Cursor
When combined summaries exceed the context budget, MapReduceChain now groups summaries, reduces each group, and repeats recursively until the final result fits in a single reduce call. - Add ChainError.reduceDepthExceeded for safety limit on recursion - Chunk labels ([Chunk N]) propagate through reduce levels as [Chunks X-Y] for source traceability - MapReduceChain gains optional contextBudget init parameter - Without a budget, single reduce is used (backward compatible) - maxReduceGroupSize and maxReduceDepth in ChainExecutionOptions control hierarchical reduce behavior Made-with: Cursor
- ChainPromptTemplate bundles map/reduce/stuff prompts for reuse - PromptTemplates provides pre-built templates: - transcriptSummary: meeting notes with speaker attribution - actionItems: extract tasks with owners and deadlines - logRootCause: error analysis with debugging steps - markdownBrief: section-aware document summarization - All templates instruct the model to cite [Chunk N] labels - DocumentChain gains run(_:template:) convenience method Made-with: Cursor
@mainactor ChainRunner provides reactive state for SwiftUI views: - phase: current ChainProgress.Phase - isRunning: execution state - result: final output - error: any failure Supports both explicit prompts and ChainPromptTemplate convenience. Handles cancellation and progress observation automatically. Made-with: Cursor
- TokenCounter protocol for pluggable token counting - WordHeuristicTokenCounter as lightweight fallback (1.33 tokens/word) - TokenAwareBackend protocol extends LLMBackend with context window and token counter info (opt-in, keeps LLMBackend clean) - PromptBudgeter resolves the best counting strategy: 1. TokenAwareBackend -> real token counting 2. .tokens budget -> word heuristic with ratio 3. .words budget -> raw word counting - AdaptiveChain refactored to use PromptBudgeter - Add MockTokenAwareBackend and tests verifying precise budgeting Made-with: Cursor
MLXBackend now accepts GenerateParameters at init, allowing callers to configure temperature, maxTokens, topP, and other MLXLMCommon sampling parameters. Parameters are passed to each ChatSession created per call. Keeps LLMBackend protocol clean (prompt in, text out) while giving MLX-specific callers full control over generation behavior. Made-with: Cursor
MemoryPressure enum provides .current() check using mach_task_basic_info to estimate available memory headroom. Returns .ok, .warning(<300MB), or .critical(<100MB). Includes os_log integration via logIfConstrained() for surfacing memory issues before inference. Important for iOS where jetsam kills memory-hungry processes. Made-with: Cursor
- Updated Quick Start with MLXBackend + GenerateParameters - Added prompt template usage example - Expanded Chunkers table with TranscriptChunker, MarkdownHeadingChunker, LogChunker - Added domain-specific examples: transcript, log, markdown, action items - Added SwiftUI integration example with ChainRunner - Added Production Options section: token budgeting, hierarchical reduce, concurrency, retries, memory awareness Made-with: Cursor
MapReduceChain now supports concurrent map execution via maxConcurrentMapTasks in ChainExecutionOptions. Default remains 1 (sequential), optimal for on-device MLX where the GPU serializes calls. When concurrency > 1, uses TaskGroup with bounded in-flight tasks. Results are always returned in original chunk order regardless of completion order. Task.checkCancellation() added before each map and reduce call for cooperative cancellation support. Made-with: Cursor
RetryPolicy struct with maxAttempts and delayMilliseconds, integrated into ChainExecutionOptions. Applied to all map and reduce calls in MapReduceChain via withRetry utility function. Default is .none (no retry), suited for on-device MLX where failures are deterministic. Useful for remote backends with transient errors. Made-with: Cursor
MockLLMBackend and MockTokenAwareBackend now use NSLock to protect all mutable state (promptsReceived, generateCallCount, cannedResponse, shouldThrow). This prevents data races when used with concurrent map phases (maxConcurrentMapTasks > 1) under Swift 6.1 strict concurrency. TransientFailBackend similarly hardened. MLXBackend's @unchecked Sendable is documented as safe (ModelContainer is an actor). Made-with: Cursor
- CI now runs resolve -> build -> test for better failure diagnosis - Updated architecture.md with all new components: domain chunkers, TokenAwareBackend, PromptBudgeter, ChainRunner, PromptTemplates, MemoryPressure, hierarchical reduce, and production reliability Made-with: Cursor
…ribution modes - Add AppleCrashReportChunker for .crash/.ips crash report analysis - Add DocumentStructureChunker for heading/page/table-aware document chunking - Add DiagnosticsMetadata types (LogChunkKind, CrashReportMetadata, LogMetadata) - Add DiagnosticSourceLabel for human-readable chunk citations - Add TranscriptAttributionMode (auto/speaker/temporal/topical) - Enhance LogChunker with diagnostic classification metadata - Enhance TranscriptChunker with single-speaker and topical attribution - Add diagnostic prompt templates (appleCrashTriage, simulatorLogRootCause, xcodeBuildFailure, testFailureAnalysis) - Add logMetadata field to TextChunkMetadata - Add 19 diagnostic tests and 14 document structure tests Made-with: Cursor
Reorganize README with adoption-focused sections: - Add "Choose Your Workflow" table mapping use cases to chunkers and templates - Add "Core Concepts" reference table - Add "Why Not Just Truncate?" and "Why Not a Generic LangChain?" sections - Move privacy note to standalone section - Add links section for CHANGELOG, CONTRIBUTING, architecture, and tests - Update CONTRIBUTING.md with resolve/build steps and domain contribution areas Made-with: Cursor
- Add CHANGELOG.md with categorized Unreleased section - Add docs/pr-summary.md with detailed reviewer-friendly PR summary - Add docs/pr-body.md with GitHub-ready PR body - Add docs/repo-adoption-metadata.md with GitHub sidebar, topics, and social preview recommendations Made-with: Cursor
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 3 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 2291033. Configure here.
|
|
||
| let mapReduce = MapReduceChain(backend: backend, chunker: chunker) | ||
| return try await mapReduce.run(text, mapPrompt: mapPrompt, reducePrompt: reducePrompt, systemPrompt: systemPrompt, progress: progress) | ||
| return try await mapReduce.run( |
There was a problem hiding this comment.
AdaptiveChain doesn't pass contextBudget to MapReduceChain
High Severity
AdaptiveChain creates a MapReduceChain without forwarding its contextBudget, using the init that sets contextBudget to nil. In MapReduceChain.fitsInSingleReduce, a nil budget causes it to always return true, which completely disables hierarchical reduce. This means large documents processed via AdaptiveChain can overflow the context window during the reduce phase — the exact scenario hierarchical reduce was designed to prevent.
Reviewed by Cursor Bugbot for commit 2291033. Configure here.
| } | ||
| } | ||
| } | ||
| throw lastError! |
There was a problem hiding this comment.
Retry logic doesn't respect cooperative task cancellation
Medium Severity
withRetry catches all errors including CancellationError, causing cancelled operations to be retried instead of propagating cancellation immediately. When delayMilliseconds is 0 (or the error occurs on the last attempt), there's no Task.sleep to re-check cancellation, so the function unnecessarily retries a cancelled task up to maxAttempts times before eventually re-throwing.
Reviewed by Cursor Bugbot for commit 2291033. Configure here.
| public init(container: ModelContainer) { | ||
| /// Parameters controlling token generation (temperature, maxTokens, topP, etc.). | ||
| /// These are passed to each `ChatSession` created per generation call. | ||
| public var generateParameters: GenerateParameters |
There was a problem hiding this comment.
Mutable generateParameters on @unchecked Sendable class causes data race
Medium Severity
generateParameters is a public var on MLXBackend, which is @unchecked Sendable. The doc comment claims safety because ModelContainer is an actor and GenerateParameters is Sendable, but Sendable doesn't protect against concurrent read/write of the property itself. One thread mutating generateParameters while generate() reads it on another creates a data race.
Reviewed by Cursor Bugbot for commit 2291033. Configure here.
- Forward contextBudget from AdaptiveChain to MapReduceChain so hierarchical reduce is enabled for long documents routed through the adaptive path. - Propagate CancellationError immediately in withRetry instead of retrying cancelled operations. Add Task.checkCancellation() at the start of each retry loop iteration. - Protect MLXBackend.generateParameters with NSLock to prevent data races on the @unchecked Sendable class. Read a snapshot of the parameters before creating each ChatSession. - Adjust AdaptiveChain routing tests to use realistic budgets that are compatible with hierarchical reduce. Made-with: Cursor
|
You have used all of your free Bugbot PR reviews. To receive reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial. |
- Map chunks are now budget-aware via PromptBudgeter.availableTextBudget and automatic rechunking when specialized chunkers emit oversized chunks. - Hierarchical reduce now uses PromptBudgeter for fit checks, including TokenAwareBackend when available. - preserveOrder now controls concurrent map result ordering while preserving correct source chunk labels via MapResult. - Reduce grouping is now budget-derived and capped by maxReduceGroupSize. - Default reservedOutputTokens changed from 0 to 512. - Added ChunkPromptFormatter for richer diagnostic chunk labels. - MLXBackend @unchecked Sendable comment now precisely describes lock protection and resource constraints. - README and docs now accurately frame PDF-extracted text and .ips support. - 105 tests passing, 0 build warnings. Made-with: Cursor
|
You have used all of your free Bugbot PR reviews. To receive reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR evolves
mlx-swift-chaininto a more complete Swift-native long-document reasoning layer for private, on-device MLX apps.It adds production-oriented primitives for long transcripts, Markdown/PDF-extracted docs, Xcode/simulator logs, crash reports, and offline reference material while preserving the package's lightweight Swift-first design.
Why this matters
MLX Swift and MLX Swift LM provide the model/runtime layer. This package now provides the app workflow layer above inference:
Highlights
Core reliability
TokenCounter,PromptBudgeter,TokenAwareBackend).Domain workflows
DocumentStructureChunkerfor Markdown/PDF-extracted documents with heading, page, table, and code block preservation.AppleCrashReportChunkerfor.crash/.ipscrash report analysis with section detection, metadata extraction, and symbolication heuristic.LogChunkerwith diagnostic classification (18LogChunkKindcases) and severity metadata.DiagnosticSourceLabelfor human-readable chunk citations.Prompt templates
SwiftUI / app integration
ChainRunnerfor reactive progress, result, and error state in SwiftUI apps.Docs and adoption
Compared with current repo
DocumentStructureChunkerwith page/heading/table preservationChainRunnerValidation
swift package resolveswift build— 0 errors, 0 warningsswift test— 96 tests, 0 failuresFinal local status:
Compatibility
nil).Limitations
Suggested review order
TextChunkMetadata,LogMetadata,DiagnosticSourceLabel)AdaptiveChain,MapReduceChain,PromptBudgeter)TranscriptChunker,DocumentStructureChunker,LogChunker,AppleCrashReportChunker)Note
Medium Risk
Touches core chain execution (budgeting/routing, map-reduce behavior, concurrency, retry/cancellation) and adds large new chunking/parsing code paths, which could change runtime behavior and performance on large inputs. Most changes are additive but the new routing/budget heuristics may affect when calls are made and how prompts are constructed.
Overview
Expands the library from basic stuff/map-reduce into a production-oriented long-document reasoning layer. Chains now support prompt-overhead-aware budgeting (with optional exact token counting via
TokenAwareBackend/TokenCounter), a configurableChainExecutionOptionssurface (reserved output tokens, retries, concurrency, reduce limits), and astuffPromptoverride for single-call execution.Strengthens
MapReduceChainexecution. The map phase can run with bounded concurrency, supports cooperative cancellation and retry, and the reduce phase gains hierarchical reduce with[Chunk N]-style source labels that propagate through intermediate reduction levels (withChainError.reduceDepthExceededsafety).Adds new domain capabilities and adoption docs. Introduces structure/diagnostic-aware chunkers (
TranscriptChunkerwith attribution modes,MarkdownHeadingChunker,DocumentStructureChunker,LogChunkerwithLogMetadata, andAppleCrashReportChunkerwithCrashReportMetadata/symbolication heuristics), plusPromptTemplates,DiagnosticSourceLabel,ChainRunnerfor SwiftUI,MemoryPressure, updated tests, and streamlined CI/docs including a newCHANGELOG.md.Reviewed by Cursor Bugbot for commit 2291033. Bugbot is set up for automated code reviews on this repo. Configure here.