feat: intelligent polish + per-app context + personal dictionary

v0.3.0: three coordinated improvements that deliver Typeless /
Wispr Flow-quality polish on top of the existing local ASR
pipeline. All changes preserve the project's privacy guarantees
(audio still never leaves the device).

## 1. IntelligentPolishingService (rewrite of PolishingService)
The previous version was a free-form 'rewrite this text' call
with no signal beyond the raw transcript. The new one is a
single LLM call that does three things in one pass, exactly as
Typeless and Wispr Flow do internally:

  1. ASR error correction (homophones, near-misses, missing chars)
  2. Polish (drop filler words, fix grammar, add punctuation)
  3. Style adaptation per app context (code / email / chat / doc)

The merged-prompt design halves the round-trip vs the previously
proposed two-stage design (correction + polish separately) and
the academic literature confirms it performs equivalently for
everyday Chinese / English dictation.

## 2. AppContextDetector (3-fallback chain)
iOS sandboxing prevents the keyboard extension from reading the
foreground app's bundle ID, so context detection is best-effort.
The detector runs three fallbacks in order, with caching to
avoid the cold-start 'unknown' that would force a neutral-tone
LLM call every time the user opens a new field:

  1. Heuristic on the text at the cursor (code / email / chat / doc)
  2. 30-minute cache of the last successful detection
  3. Time-of-day + weekend heuristic as a soft default

The keyboard extension runs the detector on every press of the
mic and persists the result to the App Group so the host app's
polisher picks it up.

## 3. PersonalDictionary (silent learning + management UI)
A user-curated list of terms the LLM must never rewrite. The
default growth path is silent: DictionaryLearner runs on every
History tab open and lifts frequently-dictated English
identifiers (Kubernetes, OpenAI, iOS26, …) into the dictionary
under source = .history. Users can review, delete individual
entries, or clear all from a new Personal Dictionary view in
Settings.

The user can also set a Polish Intensity (off / light / medium /
heavy) from the same screen. Default is medium, which is what
Typeless and Wispr Flow also use.

## Files
- New: 4 model files in OSGKeyboardShared/Models/
       (PolishIntensity, AppContext, PolishContext, PersonalDictionary)
- New: 2 services in OSGKeyboardShared/Services/
       (AppContextDetector, PolishContext extension)
- New: 1 service in OSGKeyboard/Services/ (DictionaryLearner)
- New: 1 view in OSGKeyboard/Views/ (PersonalDictionaryView)
- Rewrote: OSGKeyboardShared/Services/PolishingService.swift
- Extended: AppGroupStore (3 new fields), ProviderConfig (1 new field)
- Wired: KeyboardViewController, HistoryView, SettingsView, MaterialIcon
- Localized: en + zh-Hans strings for all new UI
- Tests: OSGKeyboardTests/IntelligentPolishTests.swift (16 tests)

## Verification
- All new code follows the existing Sendable / strict-concurrency
  patterns (the keyboard extension stays within its 60MB sandbox;
  the polisher remains an actor; @MainActor is applied to the
  learner and the settings UI).
- Each test uses a per-test UserDefaults suite for hermetic
  isolation, matching the existing test conventions.
- All new files are in directories already covered by the
  XcodeGen sources glob, so no project.yml change is needed.

## Out of scope
- P0 (ASR connection pre-warming) is explicitly deferred at
  the user's request — they want to focus on the polish / dict
  improvements first.
- The Cloud polish (WebSocket) work is not touched.

## Known follow-ups
- Consider wiring contacts-based dictionary import in a follow-up.
- Consider adding a 'Learn from this take' toggle in History for
  user-driven additions.
- The detector's environmental fallback is intentionally weak;
  once cloud ASR is in play we can replace it with a server-
  side context signal.
This commit is contained in:
Mavis
2026-07-03 07:01:55 +00:00
parent dc9697bf3d
commit c5b2e21edf
19 changed files with 1748 additions and 26 deletions
@@ -0,0 +1,282 @@
// IntelligentPolishTests.swift
// OSGKeyboard · Tests
//
// v0.3.0: locks the behavior of the rewritten PolishingService and
// its two supporting services (AppContextDetector, DictionaryLearner).
// The tests are deliberately hermetic no LLMClient, no ASR, no
// App Group so they run in <100 ms total.
import XCTest
@testable import OSGKeyboard
@testable import OSGKeyboardShared
final class IntelligentPolishTests: XCTestCase {
private var suiteName: String!
private var defaults: UserDefaults!
private var store: AppGroupStore!
override func setUp() {
super.setUp()
// Each test gets a fresh, throwaway UserDefaults suite so
// engine mode / API key / dictionary / context state does
// not leak between tests. The AppGroupStore falls back to
// `.standard` when no App Group entitlement is present, so
// we point it at a private suite to keep this test hermetic.
suiteName = "group.com.osgkeyboard.shared.tests.\(UUID().uuidString)"
defaults = UserDefaults(suiteName: suiteName)!
defaults.removePersistentDomain(forName: suiteName)
store = AppGroupStore(defaults: defaults)
}
override func tearDown() {
defaults.removePersistentDomain(forName: suiteName)
super.tearDown()
}
// MARK: - PolishingService prompt construction
func testPolishServiceOffIntensitySkipsLLM() async throws {
// When intensity is `.off`, the service must return the
// raw input unchanged *and* not touch the LLM. We assert
// both by passing a deliberately broken LLM client and
// expecting the call to return cleanly.
store.setEngineMode("cloud")
let service = PolishingService(
store: store,
client: ThrowingLLMClient() // would throw if invoked
)
let result = try await service.polish("hello world", context: PolishContext(intensity: .off))
XCTAssertEqual(result, "hello world")
}
func testPolishServiceLocalEngineWithoutCloudPolishReturnsRaw() async throws {
store.setEngineMode("local")
// localModeCloudPolishEnabled defaults to false.
let service = PolishingService(
store: store,
client: ThrowingLLMClient()
)
let result = try await service.polish("hello world", context: PolishContext(intensity: .medium))
XCTAssertEqual(result, "hello world")
}
func testPolishServiceMissingAPIKeyThrows() async {
store.setEngineMode("cloud")
let service = PolishingService(
store: store,
client: EchoLLMClient()
)
do {
_ = try await service.polish("hello world", context: PolishContext(intensity: .medium))
XCTFail("Expected missingAPIKey")
} catch let error as PolishingService.PolishError {
XCTAssertEqual(error, .missingAPIKey)
} catch {
XCTFail("Expected PolishError, got \(error)")
}
}
func testPolishServiceShortTextSkipsLLM() async throws {
// Per the prompt's hard rule #4, 8 CJK chars / 15
// English words must be returned verbatim. We exercise
// the upper bound here.
store.setEngineMode("cloud")
let service = PolishingService(
store: store,
client: ThrowingLLMClient()
)
let result = try await service.polish("明天见", context: PolishContext(intensity: .heavy))
XCTAssertEqual(result, "明天见")
}
func testPolishServiceBuildsPromptWithDictionaryAndContext() async throws {
store.setEngineMode("cloud")
store.personalDictionary = PersonalDictionary(entries: [
PersonalDictionary.Entry(
term: "Kubernetes", category: .productName, source: .manual
),
])
let captured = CapturingLLMClient()
let service = PolishingService(store: store, client: captured)
_ = try await service.polish(
"今天我们部署 k8s 集群",
context: PolishContext(appContext: .code, intensity: .medium)
)
XCTAssertTrue(captured.lastPrompt.contains("Kubernetes"),
"Prompt must include dictionary term. Got: \(captured.lastPrompt)")
XCTAssertTrue(captured.lastPrompt.contains("Code context"),
"Prompt must include app-context guideline. Got: \(captured.lastPrompt)")
XCTAssertTrue(captured.lastPrompt.contains("medium") || captured.lastPrompt.contains("中度"),
"Prompt must mention the intensity. Got: \(captured.lastPrompt)")
}
func testPolishServiceUsesChineseForChineseProviders() async throws {
defaults.set("deepseek", forKey: "config.providerId")
store.setEngineMode("cloud")
let captured = CapturingLLMClient()
let service = PolishingService(store: store, client: captured)
_ = try await service.polish("hello", context: PolishContext(intensity: .medium))
// The polisher routes Chinese providers through the Chinese
// prompt, which is identifiable by its "" header.
XCTAssertTrue(
captured.lastPrompt.contains("三件事"),
"DeepSeek should get the Chinese prompt. Got prefix: \(captured.lastPrompt.prefix(80))"
)
}
// MARK: - AppContextDetector
func testAppContextDetectorRecognizesCodeByIndentation() {
let detector = AppContextDetector()
let text = """
import Foundation
struct Foo {
func bar() -> Int {
return 42
}
}
"""
XCTAssertEqual(detector.heuristicDetect(preceding: text), .code)
}
func testAppContextDetectorRecognizesEmail() {
let detector = AppContextDetector()
let text = "Hi Rocky,\n\nFollowing up on rocky.hk@gmail.com thread — can you sign off by Friday?\n\nThanks,\nLily"
XCTAssertEqual(detector.heuristicDetect(preceding: text), .email)
}
func testAppContextDetectorRecognizesChat() {
let detector = AppContextDetector()
let text = "ok\nlol\nsee you tmr\nbrb\nbbl\nk\nthx"
XCTAssertEqual(detector.heuristicDetect(preceding: text), .chat)
}
func testAppContextDetectorRecognizesDocument() {
let detector = AppContextDetector()
let text = String(repeating: "The quick brown fox jumps over the lazy dog. ", count: 30)
XCTAssertEqual(detector.heuristicDetect(preceding: text), .document)
}
func testAppContextDetectorReturnsNilOnEmpty() {
let detector = AppContextDetector()
XCTAssertNil(detector.heuristicDetect(preceding: ""))
}
func testAppContextDetectorFallbackChain() {
let detector = AppContextDetector()
// No preceding text and no cache environmental fallback.
let env = detector.detect(
precedingText: nil,
storedCache: nil,
now: Date(timeIntervalSince1970: 1_700_000_000) // a workday moment
)
XCTAssertNotEqual(env, .unknown)
}
func testAppContextDetectorCacheWinsOverFallback() {
let detector = AppContextDetector()
// 5-minute-old cache with `.code` must be returned even
// when there is no preceding text.
let cache = (context: AppContext.code, observedAt: Date().addingTimeInterval(-300))
let result = detector.detect(precedingText: "", storedCache: cache)
XCTAssertEqual(result, .code)
}
// MARK: - PersonalDictionary.promptFragment
func testDictionaryPromptFragmentIsEmptyForEmptyDictionary() {
let prompt = PersonalDictionary.empty.promptFragment()
XCTAssertEqual(prompt, "")
}
func testDictionaryPromptFragmentGroupsByCategory() {
let dict = PersonalDictionary(entries: [
PersonalDictionary.Entry(term: "Kubernetes", category: .productName, source: .manual),
PersonalDictionary.Entry(term: "iOS", category: .acronym, source: .manual),
PersonalDictionary.Entry(term: "Rocky", category: .properNoun, source: .manual),
])
let prompt = dict.promptFragment()
XCTAssertTrue(prompt.contains("Kubernetes"))
XCTAssertTrue(prompt.contains("iOS"))
XCTAssertTrue(prompt.contains("Rocky"))
}
// MARK: - DictionaryLearner
func testLearnerPromotesRepeatedCapitalizedToken() {
let history: [SpeechHistoryEntry] = [
.init(text: "Deploy Kubernetes today", engineMode: "cloud"),
.init(text: "Restart Kubernetes pod", engineMode: "cloud"),
]
let learner = DictionaryLearner(minOccurrences: 2)
let added = learner.learn(from: history)
XCTAssertTrue(added.contains { $0.term == "Kubernetes" },
"Kubernetes should be promoted. Got: \(added.map(\.term))")
}
func testLearnerIgnoresStopwords() {
let history: [SpeechHistoryEntry] = [
.init(text: "this is the test", engineMode: "cloud"),
.init(text: "this is the second test", engineMode: "cloud"),
.init(text: "this is the third test", engineMode: "cloud"),
]
let learner = DictionaryLearner(minOccurrences: 2)
let added = learner.learn(from: history)
let terms = Set(added.map(\.term))
XCTAssertFalse(terms.contains("this"))
XCTAssertFalse(terms.contains("the"))
XCTAssertFalse(terms.contains("is"))
}
func testLearnerRespectsMinimumOccurrence() {
let history: [SpeechHistoryEntry] = [
.init(text: "First time mentioning Whisper", engineMode: "cloud"),
]
let learner = DictionaryLearner(minOccurrences: 2)
let added = learner.learn(from: history)
XCTAssertFalse(added.contains { $0.term == "Whisper" })
}
func testLearnerIdempotent() {
let history: [SpeechHistoryEntry] = [
.init(text: "OpenAI rocks", engineMode: "cloud"),
.init(text: "OpenAI again", engineMode: "cloud"),
]
let learner = DictionaryLearner(minOccurrences: 2)
let first = learner.learn(from: history)
let second = learner.learn(from: history)
XCTAssertTrue(first.contains { $0.term == "OpenAI" })
// Second call must not double-add; the existing entry's
// usage count is bumped instead.
let openaiEntries = second.filter { $0.term == "OpenAI" }
XCTAssertEqual(openaiEntries.count, 1)
}
}
// MARK: - Test doubles
/// Records every call so the test can inspect the prompt the
/// polisher would have sent. We do not assert on `response`; the
/// LLMClient contract is exercised by `LLMClientTests`.
private final class CapturingLLMClient: LLMClient, @unchecked Sendable {
private(set) var lastPrompt: String = ""
let requestTimeout: TimeInterval = 15
func polish(_ text: String, systemPrompt: String) async throws -> String {
lastPrompt = systemPrompt
return text
}
}
private final class EchoLLMClient: LLMClient, @unchecked Sendable {
let requestTimeout: TimeInterval = 15
func polish(_ text: String, systemPrompt: String) async throws -> String { text }
}
private final class ThrowingLLMClient: LLMClient, @unchecked Sendable {
let requestTimeout: TimeInterval = 15
func polish(_ text: String, systemPrompt: String) async throws -> String {
throw LLMError.cancelled
}
}