| 1 | import Foundation |
| 2 | import Vision |
| 3 | |
| 4 | // 🤖 Photos-style search over a chat's shared pictures — native-first, and |
| 5 | // deliberately so: Apple's Vision framework runs ON DEVICE, so the decrypted |
| 6 | // bytes never leave the machine that already holds them, and the E2EE rule |
| 7 | // holds exactly as it does for the feed rendering the same photos. (Web has |
| 8 | // no free on-device classifier — a model download is the price of parity |
| 9 | // there; scoped "native-only for v1", docs/superpowers/plans/ |
| 10 | // 2026-08-28-media-links-docs-native.md § Image search.) |
| 11 | // |
| 12 | // Two passes over each still (videos → poster), persisted per blob in |
| 13 | // MessageStore's image_index (derived plaintext beside the decrypted payloads |
| 14 | // it came from — never on a server), each row stamped with how far it got: |
| 15 | // |
| 16 | // pass 1, ver 1 — scene/object classification (Vision's ~1,300 labels: |
| 17 | // "dog", "beach", "receipt"). ~0.1s a photo, so "dog" works within |
| 18 | // seconds of opening. |
| 19 | // pass 2, ver 2 — accurate text recognition, the sleeper hit in a chat |
| 20 | // (screenshots, whiteboards, menus, boarding passes). ~0.2s a photo once |
| 21 | // warm — but Vision COMPILES the text model on its first use per machine |
| 22 | // (measured: 58s on an M-series Mac, 0.2s ever after), so the pass runs |
| 23 | // second and the compile is warmed in parallel with pass 1, never in |
| 24 | // front of the counter. (Seen as "Indexing photos… 0 of 251" for a |
| 25 | // minute, 2026-08-29, before this split.) |
| 26 | // |
| 27 | // Proportional to the thread's photos, narrated inline with Stop, resumable |
| 28 | // (every row lands as it's done), one image at a time off the main actor. |
| 29 | @MainActor |
| 30 | public final class ImageSearchIndex: ObservableObject { |
| 31 | /// A row at this version is complete; bump when the recipe changes so |
| 32 | /// older rows re-index on next open. |
| 33 | public static let version = 2 |
| 34 | static let labelsVersion = 1 |
| 35 | |
| 36 | @Published public private(set) var running = false |
| 37 | @Published public private(set) var progress: String? |
| 38 | /// Bumps as each photo's row lands — search UIs re-filter on it. |
| 39 | @Published public private(set) var tick = 0 |
| 40 | |
| 41 | private var instance: String? |
| 42 | private var rows: [String: MessageStore.ImageIndexRow] = [:] // blob id → what's in it |
| 43 | private var searchText: [String: String] = [:] // blob id → lowercased haystack |
| 44 | private var labelQueue: [Attachment] = [] |
| 45 | private var textQueue: [Attachment] = [] |
| 46 | private var queued: Set<String> = [] |
| 47 | private var loaded = false |
| 48 | private var task: Task<Void, Never>? |
| 49 | private static var ocrWarmed = false |
| 50 | |
| 51 | public init() {} |
| 52 | |
| 53 | /// The indexed haystack for a still (labels + recognised text so far, |
| 54 | /// lowercased), nil until that photo has been looked at. |
| 55 | public func textFor(_ blobID: String) -> String? { searchText[blobID] } |
| 56 | |
| 57 | /// Everything known about a still, for UI that wants to show labels. |
| 58 | public func row(_ blobID: String) -> MessageStore.ImageIndexRow? { rows[blobID] } |
| 59 | |
| 60 | /// Queue every media entry's still that isn't fully indexed yet and start |
| 61 | /// (or keep) the worker draining. Re-entrant: each backfill page calls |
| 62 | /// this with the grown list and only the new stills enqueue. |
| 63 | public func ensure(instance inst: String, entries: [ThreadMedia.Entry], |
| 64 | load: @escaping (Attachment) async -> Data?) { |
| 65 | if instance != inst { |
| 66 | instance = inst |
| 67 | rows = [:] |
| 68 | searchText = [:] |
| 69 | labelQueue = [] |
| 70 | textQueue = [] |
| 71 | queued = [] |
| 72 | loaded = false |
| 73 | } |
| 74 | let stills = entries.compactMap { $0.stillAttachment } |
| 75 | Task { |
| 76 | // One read of the instance's persisted index on first use; the |
| 77 | // in-memory map is the working copy from then on. |
| 78 | if !loaded { |
| 79 | let known = await MessageStore.shared.imageIndex(instance: inst, minVersion: Self.labelsVersion) |
| 80 | guard instance == inst else { return } |
| 81 | for (blob, row) in known { adopt(blob, row) } |
| 82 | loaded = true |
| 83 | } |
| 84 | var fresh = 0 |
| 85 | for still in stills where !queued.contains(still.id) { |
| 86 | let ver = rows[still.id]?.ver ?? 0 |
| 87 | if ver >= Self.version { continue } |
| 88 | queued.insert(still.id) |
| 89 | if ver >= Self.labelsVersion { textQueue.append(still) } else { labelQueue.append(still) } |
| 90 | fresh += 1 |
| 91 | } |
| 92 | guard fresh > 0 else { return } |
| 93 | startIfIdle(load: load) |
| 94 | } |
| 95 | } |
| 96 | |
| 97 | public func cancel() { |
| 98 | task?.cancel() |
| 99 | } |
| 100 | |
| 101 | private func adopt(_ blob: String, _ row: MessageStore.ImageIndexRow) { |
| 102 | rows[blob] = row |
| 103 | searchText[blob] = Self.haystack(labels: row.labels, text: row.text) |
| 104 | } |
| 105 | |
| 106 | private func startIfIdle(load: @escaping (Attachment) async -> Data?) { |
| 107 | guard task == nil, let inst = instance else { return } |
| 108 | running = true |
| 109 | progress = Self.progressLabel(reading: labelQueue.isEmpty, done: 0, |
| 110 | total: labelQueue.isEmpty ? textQueue.count : labelQueue.count) |
| 111 | // Warm the text model's one-off compile while pass 1 runs, so pass 2 |
| 112 | // starts at full speed instead of parking on its first photo. |
| 113 | if !Self.ocrWarmed { |
| 114 | Self.ocrWarmed = true |
| 115 | Task.detached(priority: .utility) { Self.warmOCR() } |
| 116 | } |
| 117 | task = Task { [weak self] in |
| 118 | await self?.drain(instance: inst, load: load) |
| 119 | await MainActor.run { |
| 120 | guard let self else { return } |
| 121 | self.running = false |
| 122 | self.progress = nil |
| 123 | self.task = nil |
| 124 | // A cancelled pass leaves its unvisited stills queued; the |
| 125 | // next ensure() picks them straight back up. |
| 126 | } |
| 127 | } |
| 128 | } |
| 129 | |
| 130 | private func drain(instance inst: String, load: @escaping (Attachment) async -> Data?) async { |
| 131 | // Pass 1 for everything waiting, then pass 2 — a still that arrives |
| 132 | // mid-pass-2 (a backfill page) gets its labels first, then its text. |
| 133 | while !Task.isCancelled, !labelQueue.isEmpty || !textQueue.isEmpty { |
| 134 | let reading = labelQueue.isEmpty |
| 135 | let total = reading ? textQueue.count : labelQueue.count |
| 136 | var done = 0 |
| 137 | while !Task.isCancelled, let still = reading ? textQueue.first : labelQueue.first { |
| 138 | if reading { textQueue.removeFirst() } else { labelQueue.removeFirst() } |
| 139 | progress = Self.progressLabel(reading: reading, done: done, total: total) |
| 140 | if let data = await load(still) { |
| 141 | // Vision is CPU/ANE-bound and synchronous — never on the |
| 142 | // main actor. One image at a time keeps memory flat. |
| 143 | let prior = rows[still.id] |
| 144 | let row = await Task.detached(priority: .utility) { |
| 145 | reading ? Self.readText(data, labels: prior?.labels ?? []) : Self.classify(data) |
| 146 | }.value |
| 147 | if Task.isCancelled { break } |
| 148 | await MessageStore.shared.setImageIndex(instance: inst, blob: still.id, row: row) |
| 149 | guard instance == inst else { return } |
| 150 | adopt(still.id, row) |
| 151 | if !reading { textQueue.append(still) } else { queued.remove(still.id) } |
| 152 | tick += 1 |
| 153 | } else { |
| 154 | queued.remove(still.id) |
| 155 | } |
| 156 | done += 1 |
| 157 | progress = Self.progressLabel(reading: reading, done: done, total: total) |
| 158 | if !reading, labelQueue.isEmpty { break } // pass 1 drained: on to the text |
| 159 | } |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | // ---- the recipe (pure, off-actor) ------------------------------------ |
| 164 | |
| 165 | /// Pass 1: what's in the picture. A failure (undecodable image, an |
| 166 | /// unsupported format) indexes as empty rather than retrying forever. |
| 167 | nonisolated static func classify(_ data: Data) -> MessageStore.ImageIndexRow { |
| 168 | let req = VNClassifyImageRequest() |
| 169 | try? VNImageRequestHandler(data: data, options: [:]).perform([req]) |
| 170 | // Apple's own idiom for "labels a person would agree with": keep an |
| 171 | // observation only where the model is ≥90% precise at ≥1% recall. |
| 172 | let labels = (req.results ?? []) |
| 173 | .filter { $0.hasMinimumRecall(0.01, forPrecision: 0.9) } |
| 174 | .map { normaliseLabel($0.identifier) } |
| 175 | return MessageStore.ImageIndexRow(labels: labels, text: "", ver: labelsVersion) |
| 176 | } |
| 177 | |
| 178 | /// Pass 2: the words in the picture, merged onto pass 1's labels. |
| 179 | nonisolated static func readText(_ data: Data, labels: [String]) -> MessageStore.ImageIndexRow { |
| 180 | let req = VNRecognizeTextRequest() |
| 181 | req.recognitionLevel = .accurate |
| 182 | req.usesLanguageCorrection = true |
| 183 | try? VNImageRequestHandler(data: data, options: [:]).perform([req]) |
| 184 | let text = (req.results ?? []) |
| 185 | .compactMap { $0.topCandidates(1).first?.string } |
| 186 | .joined(separator: "\n") |
| 187 | return MessageStore.ImageIndexRow(labels: labels, text: text, ver: version) |
| 188 | } |
| 189 | |
| 190 | /// Trip the text model's compile on a 1×1 image so the real pass never |
| 191 | /// waits for it. Vision caches the compiled model on disk, so this is |
| 192 | /// cheap on every launch but the first on a machine. |
| 193 | nonisolated static func warmOCR() { |
| 194 | var px: [UInt8] = [255, 255, 255, 255] |
| 195 | let ctx = px.withUnsafeMutableBytes { buf in |
| 196 | CGContext(data: buf.baseAddress, width: 1, height: 1, bitsPerComponent: 8, bytesPerRow: 4, |
| 197 | space: CGColorSpaceCreateDeviceRGB(), |
| 198 | bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) |
| 199 | } |
| 200 | guard let cg = ctx?.makeImage() else { return } |
| 201 | let req = VNRecognizeTextRequest() |
| 202 | req.recognitionLevel = .accurate |
| 203 | try? VNImageRequestHandler(cgImage: cg, options: [:]).perform([req]) |
| 204 | } |
| 205 | |
| 206 | /// Vision identifiers read "birthday_cake"; people type "birthday cake". |
| 207 | nonisolated static func normaliseLabel(_ id: String) -> String { |
| 208 | id.replacingOccurrences(of: "_", with: " ").lowercased() |
| 209 | } |
| 210 | |
| 211 | /// The lowercased haystack the media search substring-matches. |
| 212 | nonisolated static func haystack(labels: [String], text: String) -> String { |
| 213 | (labels.joined(separator: " ") + "\n" + text).lowercased() |
| 214 | } |
| 215 | |
| 216 | /// "Indexing photos… 12 of 340" / "Reading text in photos… 12 of 340" — |
| 217 | /// counts localised. |
| 218 | nonisolated static func progressLabel(reading: Bool, done: Int, total: Int) -> String { |
| 219 | let verb = reading ? "Reading text in photos…" : "Indexing photos…" |
| 220 | return total > 0 ? "\(verb) \(done.formatted()) of \(total.formatted())" : verb |
| 221 | } |
| 222 | } |
| 223 | |
amadan