From c6743974a6dbd1da8a7244fdad041b895b629ebc Mon Sep 17 00:00:00 2001 From: Ocean572 <45051223+Ocean572@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:22:10 -0600 Subject: [PATCH] Indexing helper: start the model text-only and bring in the media encoders only when a queue needs them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The indexing helper always loaded the full multimodal model — the text model plus the 555 MB image/audio projector (mmproj-embeddinggemma-2-Q8_0.gguf, this box's reference model) — even when every queued file was text or code, and kept it loaded through the idle linger after the run (issue #155). `SharedEmbedder` now starts text-only. Each run asks the store (`IndexStore.pendingKinds`, a DISTINCT-kind query on the queue) whether it holds a kind that embeds media (`FileKind.needsMediaEncoders`: images, screenshots, video, audio) before reading any file, and re-shapes the model then. A queue that said text and hands over media anyway — a scanned PDF's page, an unfinished long file — upgrades on the spot at the embedding call, where a model failure marks files the way it always has. Once media has been wanted the shape is sticky: encoders are not dropped and paid for twice. The text-only shape keeps the full 2048-token document budget: `LlamaEmbedder.init` takes `contextTokens` so skipping the media encoders no longer borrows the *query* encoder's 512-token cut. That decoupling is the whole safety of the change — vectors for long PDF pages, full doc chunks, long code stretches, and multilingual text sit past 512 tokens, and cutting them there changes their vectors under a fingerprint that cannot notice. Measured on this box with the two shapes side by side (9,840-char input, well past the cut, plus a short input, both embedders, one process, this box): at text-only-2048 the long and short vectors are byte-identical to the full model's; at the old text-only-512 the long vectors differ (from the first dimension) while the short ones match. Text-only indexing therefore changes no stored vector by construction; the query encoder keeps its 512 as always. (`contextTokens` is a llama.cpp-only knob: the mlx-vlm worker path in $DIGUP_WORKER_DIR sets its own width and ignores it — disclosed.) Measured on this box (Apple M1, 8 GB, macOS 26.5.1), a corpus of 60 text/code docs: peak resident size 1,045 MB → 499 MB (−546 MB), model load 1.47 s → 0.53 s, same 60 files done, same idle counts (60 searchable, 60 vectors); three meaning searches return identical hits and scores from indexes built by the old and the new helper. A mixed folder (images queued first, the usual order) starts full in both: 10.02 s vs 9.66 s wall, 1,003 MB vs 990 MB peak — within run-to-run noise, no regression. Text-only start was gated by the existing "waiting for model" check (the helper waits whenever the model set is incomplete, mmproj among them): that gate is unchanged here, disclosed. Also folded in from the pre-push review: a failed media-model swap now falls back to the text-only shape instead of stranding the session (`wantsFull` sticky + model gone made every later text embed answer "the model is closed"); `Kept.info` answers through the last-seen shape instead of force-unwrapping (a trap takes the whole helper down, and the embedder can close under a queued indexer); the `"model"` event joins the helper's protocol doc header; the sticky `wantsFull` read in `use` runs under the lock like every other access. The shape/upgrade logic lives in the executable target and stays testable only by box runs — disclosed. Suite: 127 green (base 125, +2: `pendingKindsSeeTheQueuedShapesOnly`, `mediaKindsAreTheOnesThatEmbedPicturesAndSound`). Shell JS tests 17 pass. The `digup index` one-shot CLI keeps loading the full model in-process: it has no idle linger to pay, and was left alone (disclosed). --- Sources/DigUpKit/IndexStore.swift | 14 +++ Sources/DigUpKit/Model.swift | 6 ++ Sources/LlamaRuntime/LlamaEmbedder.swift | 11 +- Sources/digup/IndexService.swift | 124 +++++++++++++++++++--- Sources/digup/main.swift | 17 ++- Tests/DigUpKitTests/IndexStoreTests.swift | 25 +++++ 6 files changed, 175 insertions(+), 22 deletions(-) diff --git a/Sources/DigUpKit/IndexStore.swift b/Sources/DigUpKit/IndexStore.swift index 7a9b8b82..e35aa011 100644 --- a/Sources/DigUpKit/IndexStore.swift +++ b/Sources/DigUpKit/IndexStore.swift @@ -346,6 +346,20 @@ public final class IndexStore { return count } + /// The kinds of the files waiting to be embedded (`only` narrows it to some files), without building + /// their records: the queue's look that `SharedEmbedder` decides the model's shape from (issue #155) + /// — a run of text and code alone never loads the image and audio encoders. + public func pendingKinds(only ids: Set? = nil) throws -> Set { + if let ids, ids.isEmpty { return [] } + let filter = ids.map { " AND id IN (\($0.map(String.init).joined(separator: ",")))" } ?? "" + var out: Set = [] + try db.query("SELECT DISTINCT kind FROM files WHERE state = 'pending'\(filter)\(scopeSQL)", + scopeValues) { row in + if let kind = FileKind(rawValue: row.text(0) ?? "") { out.insert(kind) } + } + return out + } + /// Long files with pages or chunks still to read after their first step, in indexing order (`pending`'s). public func unfinished() throws -> [FileRecord] { var out: [FileRecord] = [] diff --git a/Sources/DigUpKit/Model.swift b/Sources/DigUpKit/Model.swift index 72cb6726..87de6e3a 100644 --- a/Sources/DigUpKit/Model.swift +++ b/Sources/DigUpKit/Model.swift @@ -6,6 +6,12 @@ public enum FileKind: String, CaseIterable, Sendable, Codable { /// Source code, a repo's docs and notebooks: only in the code index (`IndexOptions.code`). case code + /// Kinds whose indexing embeds pictures or sound, so reading them needs the image and audio + /// encoders (issue #155: `SharedEmbedder` starts text-only and brings them in only for these). + /// A scanned PDF's pages are read as pictures page by page, so it's not on the list — its helper + /// upgrades when a page comes in that way, not from the queue's look. + public var needsMediaEncoders: Bool { [.image, .screenshot, .video, .audio].contains(self) } + /// The order of the first indexing pass: what people look for most comes first. static let indexingOrderSQL = """ CASE kind WHEN 'screenshot' THEN 0 WHEN 'image' THEN 1 WHEN 'pdf' THEN 2 WHEN 'doc' THEN 3 diff --git a/Sources/LlamaRuntime/LlamaEmbedder.swift b/Sources/LlamaRuntime/LlamaEmbedder.swift index 952c78f1..34e8fc3e 100644 --- a/Sources/LlamaRuntime/LlamaEmbedder.swift +++ b/Sources/LlamaRuntime/LlamaEmbedder.swift @@ -44,8 +44,13 @@ public final class LlamaEmbedder: Embedder { private var tokens = 0 // in `batch` private var position: llama_pos = 0 - /// `textOnly`: just the text model (queries). `log`: where llama.cpp's warnings and errors go (stderr when nil). - public init(models: LlamaModels, textOnly: Bool, log: URL? = nil) throws { + /// `textOnly`: skip the media encoders (queries, and issue #155's text-only indexing start). + /// `contextTokens`: the context/budget width — defaults like before (`textOnly ? 512 : 2048`); + /// the indexing helper passes 2048 with `textOnly` so a text-only start cuts document text at + /// the same 2048 the full model always used (512 is a *query* budget; indexing long pages, + /// chunks, code, or multilingual text at 512 would silently change its vectors). + /// `log`: where llama.cpp's warnings and errors go (stderr when nil). + public init(models: LlamaModels, textOnly: Bool, contextTokens: Int? = nil, log: URL? = nil) throws { let started = Date() LlamaLog.shared.send(to: log) _ = Self.backend @@ -59,7 +64,7 @@ public final class LlamaEmbedder: Embedder { } // One input must fit one micro-batch: 2048 tokens (the reference's cut), which also holds a 1120-token image or // a 750-token 30 s clip. Queries are short, so the query encoder reserves less and cuts at 512. - let maxTokens = textOnly ? 512 : 2048 + let maxTokens = contextTokens ?? (textOnly ? 512 : 2048) var params = llama_context_default_params() params.n_ctx = UInt32(maxTokens) params.n_batch = UInt32(maxTokens) diff --git a/Sources/digup/IndexService.swift b/Sources/digup/IndexService.swift index db2a807f..17cc8b14 100644 --- a/Sources/digup/IndexService.swift +++ b/Sources/digup/IndexService.swift @@ -12,8 +12,9 @@ import WinSDK /// /// Protocol: one command per line on stdin, one JSON event per line on stdout. /// stdin: sync | sync code | stop | backfill on | backfill off (EOF = stop) -/// stdout: {"event": "missing" | "pruned" | "scanning" | "sync" | "file" | "rest" | "summary" | "waiting" | "idle" | "error" | -/// "exit", ...}; events about the code index say "code": true +/// stdout: {"event": "missing" | "pruned" | "scanning" | "sync" | "file" | "rest" | "summary" | "waiting" | "idle" | +/// "model" (the embedder loaded with the media encoders) | "error" | "exit", ...}; events about the +/// code index say "code": true /// /// The loop: sync, index what the sync just queued (if it's a few files: new files go first, on any power source), /// then the backfill if allowed; then the code index the same way (`--code-root`: its own sync, asked for with @@ -246,6 +247,13 @@ final class IndexService: @unchecked Sendable { private func run(_ part: Part, only ids: Set?) -> Bool { let started = Date() let isNew = ids != nil + // Issue #155: the model takes its shape from the queue before any file is read — the image and + // audio encoders only for kinds that embed media (`FileKind.needsMediaEncoders`). The code index + // is text alone. A queue read that fails is answered conservatively: full, as it always was. + if !part.code { + let kinds = (try? part.store.pendingKinds(only: ids)) ?? Set(FileKind.allCases) + model.note(needsMedia: kinds.contains { $0.needsMediaEncoders }) + } // Cumulative over the pass: what earlier runs did, plus everything pending now (what a sync just added too). var soFar = pass[part.code] ?? (done: 0, total: 0) soFar.total = max(soFar.total, soFar.done + ((try? part.store.pendingCount()) ?? 0)) @@ -362,37 +370,125 @@ final class IndexService: @unchecked Sendable { /// One model for both indexes, started when the first file needs it and closed when the helper is done. The indexers /// get it as an `Embedder` whose `close` leaves it be, since the other may be using it. +/// Issue #155: it starts **text-only** — the image and audio encoders (555 MB of projector on this box's +/// reference model, measured 2026-11-05 on an Apple M1 with 8 GB RAM) stay unloaded until a run says its +/// queue holds files that need them (`note(needsMedia:)`, from `IndexStore.pendingKinds`), or one shows up +/// where the queue said text — a scanned PDF's page, read as a picture, upgrades on the spot through +/// `Kept`. A helper indexing only text and code never pays for the encoders, in the run *or* in the idle +/// linger after it, where the model is kept warm. Once media has been wanted the shape is sticky: the +/// encoders are not dropped and paid for twice on the next image. +/// Embedding runs on the service's one loop thread (both parts), so the lock guards the shape's state. final class SharedEmbedder: @unchecked Sendable { - private let make: () throws -> any Embedder + private let make: (_ textOnly: Bool) throws -> any Embedder + private let lock = NSLock() private var model: (any Embedder)? + private var wantsFull = false // decided: a media embed has been wanted (sticky) + private var full = false // `model` as loaded holds the media encoders + private var lastInfo: WorkerInfo? // the shape info last handed out; `Kept.info` answers with it - init(_ make: @escaping () throws -> any Embedder) { + init(_ make: @escaping (_ textOnly: Bool) throws -> any Embedder) { self.make = make } - var isRunning: Bool { model != nil } + var isRunning: Bool { lock.withLock { model != nil } } + + /// At a run's start: the queue's verdict, and the running model re-shaped to it *before* any file is + /// read. A start that fails here is kept for `embedder()`, where its error marks files the way a model + /// failure always has. + func note(needsMedia: Bool) { + lock.withLock { wantsFull = wantsFull || needsMedia } + lock.lock() + let restart = wantsFull && model != nil && !full + lock.unlock() + if restart { try? swap() } + } func embedder() throws -> any Embedder { - if model == nil { model = try make() } - return Kept(model!) + lock.lock() + let start = model == nil + let textOnly = !wantsFull + let restart = !start && wantsFull && !full + lock.unlock() + if restart { try swap() } + if start { + let started = try make(textOnly) + lock.withLock { model = started; full = !textOnly; lastInfo = started.info } + if !textOnly { emit(["event": "model", "full": true]) } + } + return Kept(self) + } + + /// The model re-shaped to full. Called from an embedding call (`Kept`) or a run's start. + /// A failed load falls back to the text-only shape (the file that asked for media fails as any + /// model failure does); without the fallback a failed swap would strand the session: `wantsFull` + /// sticky, model gone, every later text embed answering "the model is closed". + private func swap() throws { + lock.lock() + let old = model + model = nil + lock.unlock() + old?.close() + do { + let started = try make(false) + lock.withLock { model = started; full = true; wantsFull = true; lastInfo = started.info } + emit(["event": "model", "full": true]) + } catch { + lock.withLock { wantsFull = false } + throw error + } } func close() { + lock.lock() + let model = self.model + self.model = nil + lock.unlock() model?.close() - model = nil } + /// The indexers' view: always the *current* model — media calls upgrade it on the spot (issue #155: + /// a scanned PDF page from a queue that said text), text calls run on either shape. private final class Kept: Embedder { - let model: any Embedder - init(_ model: any Embedder) { self.model = model } - var info: WorkerInfo { model.info } - func embedTexts(_ texts: [String]) throws -> [[Float]] { try model.embedTexts(texts) } + let owner: SharedEmbedder + init(_ owner: SharedEmbedder) { self.owner = owner } + var info: WorkerInfo { owner.infoNow } + func embedTexts(_ texts: [String]) throws -> [[Float]] { + try owner.use { try $0.embedTexts(texts) } + } func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { - try model.embedImages(paths, budget: budget) + try owner.use(media: true) { try $0.embedImages(paths, budget: budget) } + } + func embedAudio(_ paths: [String]) throws -> [Embedding] { + try owner.use(media: true) { try $0.embedAudio(paths) } } - func embedAudio(_ paths: [String]) throws -> [Embedding] { try model.embedAudio(paths) } func close() {} } + + /// `Kept.info`'s answer: the live model's shape, or the last one seen. `embedder()` hands out a + /// `Kept` only once a model is set, and every set updates `lastInfo`, so this never traps on a + /// closed model (the force-unwrap it replaces would, and a trap takes the whole helper down). + fileprivate var infoNow: WorkerInfo { + lock.lock() + defer { lock.unlock() } + if let model { return model.info } + if let lastInfo { return lastInfo } + return WorkerInfo(model: "digup", revision: "", runtime: "", dtype: "", dim: 0, textOnly: true, loadSeconds: nil) + } + + fileprivate func use(media: Bool = false, _ body: (any Embedder) throws -> T) throws -> T { + if media, !lock.withLock({ wantsFull }) { + lock.withLock { wantsFull = true } + try swap() // the first media embed: bring the encoders in now + } + lock.lock() + let restart = model != nil && wantsFull && !full + lock.unlock() + if restart { try swap() } + lock.lock() + defer { lock.unlock() } + guard let model else { throw WorkerError.notFound("the model is closed") } + return try body(model) + } } /// One JSON object per line on stdout, written straight to the file descriptor (print() would buffer in a pipe). diff --git a/Sources/digup/main.swift b/Sources/digup/main.swift index 38c3ba2c..82fe9f53 100644 --- a/Sources/digup/main.swift +++ b/Sources/digup/main.swift @@ -194,13 +194,14 @@ func skippedLine(_ skipped: [String: Int]) -> String { // MARK: Commands /// The embedding runtime: llama.cpp in this process, or the worker in $DIGUP_WORKER_DIR (the mlx-vlm reference). -func makeEmbedder(textOnly: Bool, log: URL, qos: QualityOfService = .default) throws -> any Embedder { +/// `contextTokens`: see `LlamaEmbedder.init` — a llama.cpp-only knob (the worker sets its own width). +func makeEmbedder(textOnly: Bool, log: URL, qos: QualityOfService = .default, contextTokens: Int? = nil) throws -> any Embedder { if var config = try WorkerConfig.fromEnvironment() { config.log = log config.qos = qos return try EmbeddingWorker(config, textOnly: textOnly) } - return try LlamaEmbedder(models: try LlamaModels.locate(), textOnly: textOnly, log: log) + return try LlamaEmbedder(models: try LlamaModels.locate(), textOnly: textOnly, contextTokens: contextTokens, log: log) } /// `makeEmbedder`'s vector source when it's known without starting it: llama.cpp's always is. @@ -415,9 +416,15 @@ func status(_ args: Arguments) throws { /// The app's indexing helper: sync + index on request, JSON events out (IndexService.swift). func indexService(_ args: Arguments) throws -> Int32 { - // One model for both indexes, started when the first file needs it (`SharedEmbedder`). - let model = SharedEmbedder { - try makeEmbedder(textOnly: false, log: args.indexFolder.appendingPathComponent("worker.log"), qos: .utility) + // One model for both indexes, started when the first file needs it (`SharedEmbedder`), text-only + // until a queue needs the media encoders (issue #155). The text-only shape keeps the 2048-token + // document budget (only the media encoders are left out): 512 is the *query* cut, and indexing + // long pages, chunks, code, or multilingual text at 512 would change its vectors (measured with + // the review probe: a 2000-token text embeds identically at text-only-2048 and full, differs at + // 512 — so #155 changes no stored vector, by construction). + let model = SharedEmbedder { textOnly in + try makeEmbedder(textOnly: textOnly, log: args.indexFolder.appendingPathComponent("worker.log"), qos: .utility, + contextTokens: textOnly ? 2048 : nil) } func part(_ file: String, options: IndexOptions, roots: [URL]) throws -> IndexService.Part { let store = try IndexStore(directory: args.indexFolder, file: file) diff --git a/Tests/DigUpKitTests/IndexStoreTests.swift b/Tests/DigUpKitTests/IndexStoreTests.swift index ee3bbc51..fab8a5d8 100644 --- a/Tests/DigUpKitTests/IndexStoreTests.swift +++ b/Tests/DigUpKitTests/IndexStoreTests.swift @@ -174,6 +174,31 @@ import WinSDK #expect(try store.pendingCount() == 2) #expect(try store.prune(keepingUnder: chosen.map(\.path)) == 0) } + + // Issue #155: the queue's look `SharedEmbedder` decides the model's shape from. + + @Test func pendingKindsSeeTheQueuedShapesOnly() throws { + let store = try IndexStore(directory: folder) + func candidate(_ name: String, _ kind: FileKind, _ inode: UInt64) -> Candidate { + Candidate(path: folder.appendingPathComponent(name).path, kind: kind, size: 100, mtime: 1, + inode: inode, device: 1) + } + let image = try store.insert(candidate("a.jpg", .image, 1)) + try store.insert(candidate("b.pdf", .pdf, 2)) + try store.insert(candidate("c.txt", .doc, 3)) + #expect(try store.pendingKinds() == [.image, .pdf, .doc]) + #expect(try store.pendingKinds(only: [image]) == [.image]) + #expect(try store.pendingKinds(only: []) == []) + #expect(try store.markStopped(folder.appendingPathComponent("c.txt").path, why: "stopped")) + #expect(try store.pendingKinds() == [.image, .pdf]) // a marked file isn't queued + } + + @Test func mediaKindsAreTheOnesThatEmbedPicturesAndSound() { + #expect(FileKind.image.needsMediaEncoders && FileKind.screenshot.needsMediaEncoders) + #expect(FileKind.video.needsMediaEncoders && FileKind.audio.needsMediaEncoders) + #expect(!FileKind.pdf.needsMediaEncoders && !FileKind.doc.needsMediaEncoders) + #expect(!FileKind.code.needsMediaEncoders) + } } /// Sets an environment variable of this process, or removes it (nil).