mlx-swift-chain

Swift-native long-document reasoning for private, on-device MLX apps.

Long-document chunking, map-reduce, prompt budgeting, and source-grounded summarization for MLX Swift apps on macOS and iOS.

Process meeting transcripts, voice notes, Markdown documents, PDF-extracted text, Xcode logs, crash reports, and field manuals locally on macOS and iOS using MLX Swift-compatible backends.

Why This Exists

MLX Swift and MLX Swift LM handle model loading and inference. mlx-swift-chain handles everything above the model layer:

• Chunking — structure-aware splitting for transcripts, documents, logs, crash reports, and code-heavy Markdown
• Prompt/context budgeting — token-aware routing that accounts for prompt overhead and reserved output
• Map-reduce — process every chunk through the LLM, then combine results
• Hierarchical reduce — recursive reduction when combined summaries still exceed the context window
• Streaming — token-by-token output via StreamingLLMBackend protocol; non-streaming backends work unchanged
• Structured JSON extraction — runJSON<T: Decodable>(...) with automatic code-fence stripping and retry
• Rich results — ChainResult with source chunks, ChainMetrics (timing, call counts, token estimates)
• Safe prompt rendering — opt-in <source> tag wrapping with metadata via ChainPromptBuilder
• Retries and cancellation — configurable retry policy and cooperative Task cancellation
• Progress — AsyncStream-based phase reporting (stuffing, mapping N/M, reducing, complete)
• Source-grounded outputs — chunk labels ([Chunk N]) and metadata propagate through reduce levels
• SwiftUI integration — @Observable ChainRunner with streaming partialText support

Installation

Add to your Package.swift:

dependencies: [
    .package(url: "https://github.com/joelnishanth/mlx-swift-chain", from: "0.1.0"),
],
targets: [
    .target(
        name: "YourApp",
        dependencies: [
            .product(name: "MLXSwiftChain", package: "mlx-swift-chain"),
        ]
    ),
]

Quick Start

import MLXSwiftChain
import MLXLMCommon

// 1. Set up a backend
let backend = MLXBackend(
    container: modelContainer,
    generateParameters: GenerateParameters(maxTokens: 1024, temperature: 0.3)
)

// 2. Create an adaptive chain with a domain chunker
let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(4096),
    chunker: TranscriptChunker(targetWords: 800, overlapTurns: 1)
)

// 3. Run with a prompt template
let summary = try await chain.run(
    transcript, template: PromptTemplates.transcriptSummary
)

For short texts, AdaptiveChain uses a single LLM call (zero overhead). For long texts, it automatically chunks the input, maps each chunk through the LLM, and reduces the results.

Choose Your Workflow

Use Case	Chunker	Templates	Value
Meeting transcript / voice note	TranscriptChunker	transcriptSummary, actionItems, voiceNoteSummary, personalMemoActions	Speaker/temporal/topic attribution preserves who said what and when
Markdown / PDF-extracted text	DocumentStructureChunker or MarkdownHeadingChunker	markdownBrief, documentExecutiveSummary, documentStudyGuide	Heading/page/source-aware summaries with DocumentLocation metadata
Xcode / simulator / crash logs	LogChunker or AppleCrashReportChunker	logRootCause, appleCrashTriage, xcodeBuildFailure, testFailureAnalysis	Source-grounded private developer triage with diagnostic classification
Offline field manuals / emergency docs	DocumentStructureChunker	markdownBrief, documentExecutiveSummary	On-device reference material processing with page/section tracking
Generic prose	SentenceAwareChunker	Custom prompt	Safe fallback for any text

All templates are accessed via PromptTemplates.<name> and work with the chain.run(_:template:) convenience.

Core Concepts

Type	Role
AdaptiveChain	Auto-selects Stuff or MapReduce based on input length and prompt overhead. Start here.
StuffChain	Single LLM call when input fits in the context window.
MapReduceChain	Chunks, maps each through the LLM, reduces combined results. Supports hierarchical reduce for very large documents.
ChainResult	Rich result with .text, .sourceChunks, and .metrics. Returned by runWithMetadata(...).
ChainMetrics	Chunk count, map/reduce call counts, elapsed time, estimated tokens, throughput.
ChainEvent	Streaming event: .chunk(String), .progress(Update), .result(ChainResult).
TextChunk	A chunk of text with word count, index, and metadata.
TextChunkMetadata	Source word ranges, timestamps, speaker labels, DocumentLocation, LogMetadata.
PromptTemplates	Pre-built map/reduce/stuff prompt bundles for common workflows.
ChainRunner	@Observable @MainActor class for SwiftUI integration with streaming support.
ChainProgress	AsyncStream<Update> with phase info, elapsed time, and optional partial metrics.

Choose a Chunker

Chunker	Best For
SentenceAwareChunker	General text. Splits at sentence boundaries. Default.
FixedSizeChunker	Uniform chunks by word count with optional overlap.
TranscriptChunker	Meeting transcripts, voice notes, lectures, and memos. Auto-selects speaker, temporal, or topical attribution.
MarkdownHeadingChunker	Markdown documents. Splits at headings, preserves structure. Falls back to sentence splitting for large sections.
LogChunker	Xcode/simulator logs. Keeps stack traces intact, splits at timestamp boundaries. Classifies chunks by diagnostic kind.
AppleCrashReportChunker	Apple crash reports — translated .crash text and lightweight grouping for JSON-like .ips text. Preserves crashed thread, exception info, and binary images. Detects symbolication status. Does not fully interpret Apple's crash-report JSON schema.
CodeBlockAwareChunker	Markdown / code-heavy notes. Splits at paragraph boundaries but never inside fenced code blocks. Oversized code blocks become their own chunk.
DocumentStructureChunker	Markdown or PDF-extracted text. Preserves headings, page markers, Markdown-style tables, fenced code blocks, and lists. Populates DocumentLocation metadata. PDF parsing and OCR are out of scope — extract text first and preserve page markers when possible.

All chunkers populate TextChunkMetadata with chunk index, source word ranges, discovered timestamps, and speaker labels.

Examples

Summarize a meeting transcript:

let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(4096),
    chunker: TranscriptChunker(targetWords: 800, overlapTurns: 1)
)

let summary = try await chain.run(
    transcript, template: PromptTemplates.transcriptSummary
)

Summarize a voice note:

let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(4096),
    chunker: TranscriptChunker(targetWords: 800)
)

let summary = try await chain.run(
    voiceNote, template: PromptTemplates.voiceNoteSummary
)

Extract actions from a personal memo:

let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(4096),
    chunker: TranscriptChunker(targetWords: 800, attributionMode: .temporal)
)

let actions = try await chain.run(
    memo, template: PromptTemplates.personalMemoActions
)

TranscriptChunker supports adaptive attribution. Multi-speaker meetings use speaker-turn attribution, preserving who said what. Single-speaker voice notes and lectures use timestamp or topic attribution to avoid repetitive "Speaker 1" labels while preserving source grounding. Use TranscriptAttributionMode.auto (the default) to let the chunker select the best strategy, or force a specific mode when you know your input format.

Analyze Xcode logs for root cause:

let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(4096),
    chunker: LogChunker(targetWords: 1000)
)

let analysis = try await chain.run(
    logOutput, template: PromptTemplates.logRootCause
)

Triage an Apple crash report:

let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(8192),
    chunker: AppleCrashReportChunker(targetWords: 1200)
)

let triage = try await chain.run(
    crashReportText, template: PromptTemplates.appleCrashTriage
)

Create a section-aware brief from a Markdown document:

let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(4096),
    chunker: MarkdownHeadingChunker(targetWords: 1200)
)

let brief = try await chain.run(
    markdownDoc, template: PromptTemplates.markdownBrief
)

Summarize a PDF-extracted document with page tracking:

let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(4096),
    chunker: DocumentStructureChunker(targetWords: 1200, preserveTables: true)
)

let summary = try await chain.run(
    pdfExtractedText, template: PromptTemplates.markdownBrief
)
// Each chunk's metadata.documentLocation includes pageRange, headingPath, and blockTypes

Extract action items with source citations:

let actions = try await chain.run(
    notes, template: PromptTemplates.actionItems
)
// Output includes [Chunk N] references for traceability

Token-Accurate Budgeting with MLX

When you have access to the model's tokenizer, use MLXTokenAwareBackend for exact token counting instead of word heuristics:

import MLXSwiftChain
import MLXLMCommon

let container: ModelContainer = /* load your model */
let tokenizer = await container.tokenizer

let backend = MLXTokenAwareBackend(
    container: container,
    tokenizer: tokenizer,
    contextWindowTokens: 8192,
    generateParameters: GenerateParameters(maxTokens: 1024, temperature: 0.3)
)

let chain = AdaptiveChain(
    backend: backend,
    contextBudget: .tokens(8192)
)

let summary = try await chain.run(longDocument, template: PromptTemplates.markdownBrief)

Budget calculations (stuff vs. map-reduce routing, chunk sizing, reduce grouping) now use real token counts. MLXBackend without a tokenizer continues to work with word heuristics.

Streaming Output

Stream the final generation output token by token:

for try await event in chain.stream(transcript, template: PromptTemplates.transcriptSummary) {
    switch event {
    case .chunk(let fragment):
        print(fragment, terminator: "")  // incremental text
    case .progress(let update):
        print("\n[\(update.phase)]")
    case .result(let chainResult):
        print("\nDone — \(chainResult.metrics?.elapsedTime ?? .zero)")
    }
}

Streaming works with MLXBackend and any backend conforming to StreamingLLMBackend. Non-streaming backends emit a single .chunk with the full text.

Structured JSON Extraction

Extract typed values from LLM output with automatic JSON repair and retry:

struct ActionItem: Decodable {
    let owner: String
    let task: String
    let deadline: String?
}

let items = try await chain.runJSON(
    [ActionItem].self,
    text: meetingNotes,
    mapPrompt: "Extract action items as a JSON array:",
    reducePrompt: "Merge into a single JSON array:",
    options: ChainExecutionOptions(retryPolicy: RetryPolicy(maxAttempts: 2))
)

Handles markdown code fences, surrounding prose, and nested JSON. On decode failure, retries with an error hint appended to the prompt.

Rich Results and Metrics

Get source chunks and performance metrics alongside the generated text:

let result = try await chain.runWithMetadata(
    document,
    template: PromptTemplates.markdownBrief
)

print(result.text)
print("Chunks: \(result.sourceChunks.count)")
if let m = result.metrics {
    print("Map calls: \(m.mapCallCount), Reduce calls: \(m.reduceCallCount)")
    print("Elapsed: \(m.elapsedTime)")
    if let tps = m.tokensPerSecond { print("Throughput: \(tps) tok/s") }
}

Safe Prompt Rendering

Opt into <source> tag wrapping for better source traceability and reduced prompt injection surface:

let options = ChainExecutionOptions(promptStyle: .delimited)
let summary = try await chain.run(
    document, template: PromptTemplates.markdownBrief, options: options
)

In .delimited mode, source text is wrapped with metadata:

<source index="0" of="5" words="0-500" speaker="Alice" time="00:01:00-00:05:30">
...chunk text...
</source>

Default is .raw (legacy prompt + text concatenation) for backward compatibility.

SwiftUI Integration

import MLXSwiftChain

struct SummaryView: View {
    @State private var runner = ChainRunner()
    let chain: AdaptiveChain

    var body: some View {
        VStack {
            if runner.isRunning, let phase = runner.phase {
                ProgressView()
                Text(String(describing: phase))
            }
            if let partial = runner.partialText, runner.isRunning {
                ScrollView { Text(partial) }
            }
            if let result = runner.result {
                ScrollView { Text(result) }
            }
            Button("Summarize") {
                runner.run(chain, text: document,
                          template: PromptTemplates.transcriptSummary)
            }
            Button("Stream") {
                runner.runStreaming(chain, text: document,
                    mapPrompt: "Summarize:", reducePrompt: "Combine:")
            }
        }
    }
}

ChainRunner exposes partialText for streaming, chainResult for metrics and source chunks, and result for the final text.

Production Options

Token budgeting: AdaptiveChain accounts for system prompt, task prompt, input text, and reserved output tokens when deciding stuff vs. map-reduce routing. Backends conforming to TokenAwareBackend enable exact token counting. Map chunks are also budget-aware — if a specialized chunker emits oversized chunks, they are automatically re-split to fit within the map prompt budget. Fallback re-chunking converts token budgets into conservative word targets before using word-based chunkers.

Hierarchical reduce: When MapReduceChain is initialized with a contextBudget, it automatically groups and recursively reduces summaries that exceed the budget, preventing context overflow on large documents. Reduce-fit checks use PromptBudgeter (including TokenAwareBackend when available), and group sizes are budget-derived up to the maxReduceGroupSize cap.

Reserved output tokens: ChainExecutionOptions reserves 512 tokens for model output by default, providing a conservative margin for budget calculations. Set to 0 for maximum-input / legacy behavior.

Concurrency: ChainExecutionOptions(maxConcurrentMapTasks: 4) enables bounded parallel mapping. Default is 1, which is optimal for on-device MLX inference. Set preserveOrder: false to receive map results in completion order rather than original chunk order.

Retries: ChainExecutionOptions(retryPolicy: RetryPolicy(maxAttempts: 3, delayMilliseconds: 500)) adds retry logic, primarily useful for remote backends.

Memory awareness: MemoryPressure.current() checks available memory before inference — critical for iOS where jetsam kills memory-hungry processes.

Option	Type	Purpose
ChainExecutionOptions	struct	Bundles concurrency, retry, reduce depth, output token reservation, and prompt style
RetryPolicy	struct	Max attempts and delay for transient backend failures
ContextBudget	enum	.words(N) or .tokens(N) budget for adaptive routing
PromptBudgeter	struct	Checks whether system + task + text + reserved output fits in budget
PromptStyle	enum	.raw (legacy concatenation) or .delimited (<source> tag wrapping)
ChainPromptBuilder	struct	Centralized prompt rendering with metadata-enriched source tags
TokenCounter	protocol	Pluggable token counting; ships with WordHeuristicTokenCounter
TokenAwareBackend	protocol	Extends LLMBackend with context window size and token counter
MLXTokenCounter	struct	TokenCounter backed by a real MLXLMCommon Tokenizer
MLXTokenAwareBackend	class	MLX backend with TokenAwareBackend + StreamingLLMBackend conformance
StreamingLLMBackend	protocol	Extends LLMBackend with token-by-token streaming
MemoryPressure	enum	.current() returns .ok, .warning, or .critical based on available memory

Progress Reporting

let progress = ChainProgress()

Task {
    for await update in progress.updates {
        switch update.phase {
        case .stuffing: print("Processing in single call...")
        case .mapping(let step, let total): print("Chunk \(step)/\(total)...")
        case .reducing: print("Combining results...")
        case .complete: print("Done in \(update.elapsedTime)")
        }
    }
}

let result = try await chain.run(text, mapPrompt: "...", reducePrompt: "...", progress: progress)

Diagnostic Log Analysis

mlx-swift-chain can chunk and summarize long diagnostic logs entirely on-device. Use AppleCrashReportChunker for translated .crash text and lightweight grouping of JSON-like .ips text, and LogChunker for simulator logs, Xcode build output, and test failures. The crash chunker does not fully interpret Apple's crash-report JSON schema.

For diagnostic workflows, ChunkPromptFormatter.labeledText(for:) produces richer chunk labels using DiagnosticSourceLabel metadata when available (e.g. [Chunk 2, exceptionInfo, EXC_BAD_ACCESS]).

Input	Chunker	Prompt Template
Apple crash report (.crash / .ips)	AppleCrashReportChunker	PromptTemplates.appleCrashTriage
Simulator / Console.app logs	LogChunker	PromptTemplates.simulatorLogRootCause
Xcode build failure	LogChunker	PromptTemplates.xcodeBuildFailure
XCTest failure	LogChunker	PromptTemplates.testFailureAnalysis

Each chunk carries LogMetadata with a LogChunkKind (e.g. .crashedThread, .swiftCompilerError, .testFailure), process name, severity, and — for crash reports — full CrashReportMetadata including exception type, symbolication status, and crashed thread number.

Performance Guidance for Local MLX Inference

• maxConcurrentMapTasks: 1 (the default) is optimal for on-device MLX inference. Apple Silicon GPUs serialize inference calls, so parallelism adds overhead without throughput benefit.
• Increase maxConcurrentMapTasks only for remote or non-GPU-constrained backends.
• Streaming adds minimal overhead — MLXBackend uses ChatSession.streamResponse internally even for non-streaming generate() calls.
• Token budgeting with MLXTokenAwareBackend is more accurate but requires resolving the tokenizer at init time (await container.tokenizer). Falls back to word heuristics when not available.
• reservedOutputTokens: 512 is the default. Increase for tasks requiring longer outputs; set to 0 for maximum input budget.

Versioning Notes

mlx-swift-chain follows semantic versioning. All new APIs introduced in this version are additive — existing public types, protocols, and methods are unchanged:

• LLMBackend, TokenAwareBackend, DocumentChain, TextChunker protocols — unchanged
• StuffChain, MapReduceChain, AdaptiveChain — existing run(...) signatures unchanged
• ChainExecutionOptions — new fields have defaults that preserve existing behavior
• ChainProgress.Update — new partialMetrics field is optional with a nil default

New protocol requirements on DocumentChain (runWithMetadata, stream) have default implementations, so existing conformances compile without changes.

Why Not Just Truncate?

Local LLMs have limited context windows (e.g. 8,192 tokens for Gemma models). Long documents — research papers, legal contracts, codebases, transcripts — easily exceed this limit. Naive prefix truncation discards most of the content:

Document Size	Prefix Truncation Coverage	mlx-swift-chain Model-Visible Coverage
~2,500 words	100%	100%
~5,000 words	~60%	100%
~10,000 words	~30%	100%
~20,000 words	7%	100%

Coverage here means input/model-visible coverage across chunks, not guaranteed perfect retention in the final reduced answer.

• Truncation loses evidence. A crash report's root cause may be on page 5. A contract's liability clause may be in the appendix. Truncation throws it away.
• Map-reduce improves coverage. Every chunk is processed by the LLM. Nothing is silently dropped.
• Hierarchical reduce prevents overflow. When combined summaries still exceed the context window, recursive grouping and reduction keeps the final prompt within budget.
• Source labels help verification. [Chunk N] references let users trace claims back to the original text.

Why Not a Generic LangChain?

• Apple-native Swift package. No Python bridge, no HTTP overhead, no serialization layer. Just Swift protocols and structured concurrency.
• MLX-first / local-backend-first. Designed for on-device inference on Apple Silicon with MLXBackend, not cloud API wrappers.
• Offline and private. All processing can stay on-device. No telemetry, no network calls unless your backend makes them.
• SwiftUI-ready progress. ChainRunner is @Observable and @MainActor — drop it into a SwiftUI view.
• Domain chunkers. Purpose-built for transcripts, documents, logs, and crash reports — not generic text splitting.
• One problem, solved well. This is not a framework for agents, tool use, vector stores, or retrieval. It handles long-document reasoning above the model layer.

Privacy

MLX Swift Chain analyzes user-provided text. It does not acquire logs, read files, call Xcode APIs, or transmit data by itself. Whether inference stays local depends on the backend you provide.

Architecture

See docs/architecture.md for the component graph and docs/data-flow.md for the map-reduce sequence diagram.

Requirements

• macOS 14.0+ / iOS 17.0+
• Swift 6.1+
• mlx-swift-lm 3.31.3+

Links

• CHANGELOG
• Contributing
• Architecture
• Data Flow
• Diagnostic Log Analysis Review
• Tests

License

Apache 2.0 — see LICENSE.