Indexing Helper: Start Text-Only and Load Media Encoders on Demand (#200)

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<Int64>? = nil) throws -> Set<FileKind> {
+        if let ids, ids.isEmpty { return [] }
+        let filter = ids.map { " AND id IN (\($0.map(String.init).joined(separator: ",")))" } ?? ""
+        var out: Set<FileKind> = []
+        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] = []

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

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)

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<Int64>?) -> 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<T>(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).

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)

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).