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:
@@ -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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user