From dcdd22020c77cf2420ec3fa7ea5cc4dc768adcf1 Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Mon, 1 Jun 2026 08:30:24 -0700 Subject: [PATCH] Surface per-line OCR confidence so the low-confidence filter runs ScreenTextExtractor discarded Vision's per-observation confidence, so the OCRTextHygiene.dropLowConfidence filter only ever saw a synthesized 1.0 and never dropped anything. Carry each recognized line's confidence through a new ExtractedScreenText.lines field and feed those lines to OCRTextHygiene.clean, so the recognizer's weakest guesses are dropped before they reach the completion prompt. The joined text field is kept for logging and the window-title fallback. OCRLine gains Sendable so it can ride inside the Sendable ExtractedScreenText across the OCR continuation boundary. --- .../Services/Visual/ScreenTextExtractor.swift | 39 +++++++++++--- .../Visual/ScreenshotContextGenerator.swift | 10 ++-- Cotabby/Support/OCRTextHygiene.swift | 2 +- .../ScreenshotContextGeneratorTests.swift | 54 +++++++++++++++++-- 4 files changed, 87 insertions(+), 18 deletions(-) diff --git a/Cotabby/Services/Visual/ScreenTextExtractor.swift b/Cotabby/Services/Visual/ScreenTextExtractor.swift index e9bcddf7..467b725b 100644 --- a/Cotabby/Services/Visual/ScreenTextExtractor.swift +++ b/Cotabby/Services/Visual/ScreenTextExtractor.swift @@ -14,6 +14,17 @@ import Logging struct ExtractedScreenText: Sendable { let text: String let lineCount: Int + /// Per-line OCR text paired with Vision's recognition confidence, in reading order. Carries the + /// confidence that the joined `text` discards, so `OCRTextHygiene.dropLowConfidence` can filter on + /// real values instead of a synthesized constant. Defaults to empty for callers (and tests) that + /// only supply joined text. + let lines: [OCRTextHygiene.OCRLine] + + init(text: String, lineCount: Int, lines: [OCRTextHygiene.OCRLine] = []) { + self.text = text + self.lineCount = lineCount + self.lines = lines + } } /// Test seam for screenshot OCR. @@ -72,7 +83,10 @@ struct ScreenTextExtractor: ScreenTextExtracting { } let observations = (request.results as? [VNRecognizedTextObservation]) ?? [] - let orderedLines = observations + // Keep each line's confidence (from its top candidate) so the hygiene pass can drop + // the recognizer's weakest guesses; the joined `text` below is for logging and the + // window-title fallback only. + let recognizedLines: [OCRTextHygiene.OCRLine] = observations .sorted { if Swift.abs($0.boundingBox.minY - $1.boundingBox.minY) > 0.02 { return $0.boundingBox.minY > $1.boundingBox.minY @@ -80,27 +94,36 @@ struct ScreenTextExtractor: ScreenTextExtracting { return $0.boundingBox.minX < $1.boundingBox.minX } - .compactMap { $0.topCandidates(1).first?.string } - .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } - .filter { !$0.isEmpty } + .compactMap { observation -> OCRTextHygiene.OCRLine? in + guard let candidate = observation.topCandidates(1).first else { return nil } + let trimmed = candidate.string.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + return OCRTextHygiene.OCRLine(text: trimmed, confidence: candidate.confidence) + } - let joinedText = orderedLines.joined(separator: "\n") + let joinedText = recognizedLines.map(\.text).joined(separator: "\n") let cappedText = String(joinedText.prefix(maxRecognizedCharacters)) guard !cappedText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { let elapsedMilliseconds = Int(Date().timeIntervalSince(startedAt) * 1000) - self.log("ocr-empty elapsed_ms=\(elapsedMilliseconds) lines=\(orderedLines.count)") + self.log("ocr-empty elapsed_ms=\(elapsedMilliseconds) lines=\(recognizedLines.count)") continuation.resume(throwing: ScreenTextExtractionError.noRecognizedText) return } let elapsedMilliseconds = Int(Date().timeIntervalSince(startedAt) * 1000) self.log( - "ocr-success elapsed_ms=\(elapsedMilliseconds) lines=\(orderedLines.count) chars=\(cappedText.count) " + + "ocr-success elapsed_ms=\(elapsedMilliseconds) lines=\(recognizedLines.count) chars=\(cappedText.count) " + "preview=\(self.preview(cappedText))" ) - continuation.resume(returning: ExtractedScreenText(text: cappedText, lineCount: orderedLines.count)) + continuation.resume( + returning: ExtractedScreenText( + text: cappedText, + lineCount: recognizedLines.count, + lines: recognizedLines + ) + ) } // Accurate OCR is slower, but visual context is only captured once per focused diff --git a/Cotabby/Services/Visual/ScreenshotContextGenerator.swift b/Cotabby/Services/Visual/ScreenshotContextGenerator.swift index f8daf9f3..fe0f42cf 100644 --- a/Cotabby/Services/Visual/ScreenshotContextGenerator.swift +++ b/Cotabby/Services/Visual/ScreenshotContextGenerator.swift @@ -60,9 +60,9 @@ final class ScreenshotContextGenerator { await onStatusChange?(.extractingText) - let extractedText: String + let extracted: ExtractedScreenText do { - extractedText = try await textExtractor.extractText(from: screenshot.image).text + extracted = try await textExtractor.extractText(from: screenshot.image) } catch ScreenTextExtractionError.noRecognizedText { guard let windowTitle = screenshot.windowTitle else { throw ScreenshotContextGenerationError.unavailable( @@ -94,9 +94,7 @@ final class ScreenshotContextGenerator { // safety. No model summarization: a base model conditions fine on cleaned raw context, and // the old summary step cost an extra generation per refresh and could hallucinate. let cleanedOCR = OCRTextHygiene.clean( - lines: extractedText - .split(separator: "\n", omittingEmptySubsequences: true) - .map { OCRTextHygiene.OCRLine(text: String($0), confidence: 1.0) }, + lines: extracted.lines, fieldText: context.precedingText + " " + context.trailingText, maxChars: configuration.maxRecognizedCharacters ) @@ -105,7 +103,7 @@ final class ScreenshotContextGenerator { if CotabbyDebugOptions.isEnabled { saveDebugScreenshot( screenshot.image, - text: extractedText, + text: extracted.text, name: sanitizedDebugName(from: context.applicationName) ) } diff --git a/Cotabby/Support/OCRTextHygiene.swift b/Cotabby/Support/OCRTextHygiene.swift index 0413248a..20863709 100644 --- a/Cotabby/Support/OCRTextHygiene.swift +++ b/Cotabby/Support/OCRTextHygiene.swift @@ -21,7 +21,7 @@ enum OCRTextHygiene { /// Confidence is carried alongside the text because the cheapest, highest-signal filter /// (`dropLowConfidence`) needs it. The orchestrating extractor currently discards Vision's /// per-candidate confidence; surfacing it into this value type is what lets filter #1 run. - struct OCRLine: Equatable { + struct OCRLine: Equatable, Sendable { let text: String let confidence: Float diff --git a/CotabbyTests/ScreenshotContextGeneratorTests.swift b/CotabbyTests/ScreenshotContextGeneratorTests.swift index 8484d7bd..1fdc6faf 100644 --- a/CotabbyTests/ScreenshotContextGeneratorTests.swift +++ b/CotabbyTests/ScreenshotContextGeneratorTests.swift @@ -44,18 +44,66 @@ final class ScreenshotContextGeneratorTests: XCTestCase { private func makeGenerator( extractedText: String, configuration: VisualContextConfiguration = .default + ) -> ScreenshotContextGenerator { + // Mirror the real extractor: split into per-line OCR with a confidence above the hygiene + // threshold, so these cases keep exercising the non-confidence filters exactly as before. + let lines = extractedText + .split(separator: "\n", omittingEmptySubsequences: true) + .map { OCRTextHygiene.OCRLine(text: String($0), confidence: 0.9) } + return makeGenerator( + extracted: ExtractedScreenText(text: extractedText, lineCount: lines.count, lines: lines), + configuration: configuration + ) + } + + private func makeGenerator( + lines: [OCRTextHygiene.OCRLine], + configuration: VisualContextConfiguration = .default + ) -> ScreenshotContextGenerator { + let joined = lines.map(\.text).joined(separator: "\n") + return makeGenerator( + extracted: ExtractedScreenText(text: joined, lineCount: lines.count, lines: lines), + configuration: configuration + ) + } + + private func makeGenerator( + extracted: ExtractedScreenText, + configuration: VisualContextConfiguration ) -> ScreenshotContextGenerator { ScreenshotContextGenerator( screenshotService: StubScreenshotCapture( screenshot: CapturedWindowScreenshot(image: makeImage(), windowTitle: nil) ), - textExtractor: StubTextExtractor( - result: .success(ExtractedScreenText(text: extractedText, lineCount: 1)) - ), + textExtractor: StubTextExtractor(result: .success(extracted)), configuration: configuration ) } + func test_generateContext_dropsLowConfidenceOCRLines() async throws { + // A clean, plausible sentence at low confidence must be dropped even though no other hygiene + // filter would catch it, proving real per-line Vision confidence now reaches the hygiene pass. + let configuration = VisualContextConfiguration( + snapshotDimension: 700, + maxImageDimension: 1600, + minRecognizedCharacterCount: 12, + maxRecognizedCharacters: 500, + maxSummaryCharacters: 200 + ) + let generator = makeGenerator( + lines: [ + OCRTextHygiene.OCRLine(text: "The quarterly report is due on Friday afternoon.", confidence: 0.2), + OCRTextHygiene.OCRLine(text: "Please review the attached budget spreadsheet carefully.", confidence: 0.95) + ], + configuration: configuration + ) + + let excerpt = try await generator.generateContext(for: makeSnapshot()) + + XCTAssertFalse(excerpt.text.contains("quarterly report")) + XCTAssertTrue(excerpt.text.contains("budget spreadsheet")) + } + private func makeSnapshot() -> FocusedInputSnapshot { FocusedInputSnapshot( applicationName: "Xcode",