mirror of
https://github.com/BigMarc/grammargem.git
synced 2026-10-13 17:01:55 +02:00
feat: backend-neutral AI seam, Ollama backend, extension bridge + trust/correctness fixes
Wave 1 (correctness/trust): - Fix paid large-model dead path: engine reads the persisted selected repo via a thread-safe ModelManager.activeRepo(), so download/load/isReady now agree. - Lemon Squeezy store/product check fails closed in release builds — the placeholder-ID 'any key unlocks paid' bypass can no longer ship. - Replace force-unwrapped Hugging Face URLs with safe URLComponents construction. - Never read/replace/scan secure (password) fields via AX.isSecureTextField. Wave 2 (compounding seam): - Extract AIPromptKit (prompts/temperature/maxTokens/cleanup) out of MLXEngine so every backend produces identical prompts and output. - AIEngine gains protectedTerms (personal dictionary now actually injected into the LLM), warmup(), and runStreaming(); add makeAIEngine() factory. - Honor WritingMode.lengthCap for maxTokens; warm the model on launch/after download. - Guard custom modes against empty prompts and duplicate names. Wave 3: OllamaEngine — a URLSession-only AIEngine backend (cross-platform, opt-in via GRAMMARGEM_AI_BACKEND=ollama). Wave 4: LocalEngineServer — loopback-only, token-authed HTTP bridge (Network.framework, no new deps), off by default, exposing Harper + the LLM + the entitlement gate to a browser extension. Tests: 6 new unit tests for the AIPrompts seam and custom-mode guard (19 pass). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
e5b0977abf
commit
201fe04d8a
@ -10,11 +10,52 @@ protocol AIEngine: AnyObject {
|
|||||||
var isReady: Bool { get }
|
var isReady: Bool { get }
|
||||||
|
|
||||||
/// Run an action against `text` and return the rewritten result.
|
/// Run an action against `text` and return the rewritten result.
|
||||||
/// Throws on model-not-ready or generation failure. Never performs network I/O.
|
/// `protectedTerms` are vocabulary the model must reproduce verbatim (the
|
||||||
func run(_ action: AIAction, on text: String) async throws -> String
|
/// personal dictionary / style-guide terms). Throws on model-not-ready or
|
||||||
|
/// generation failure. Never performs network I/O on the on-device backend.
|
||||||
|
func run(_ action: AIAction, on text: String, protectedTerms: [String]) async throws -> String
|
||||||
|
|
||||||
/// Best-effort local tone classification (used by tone-detection UI).
|
/// Best-effort local tone classification (used by tone-detection UI).
|
||||||
func detectTone(_ text: String) async -> Tone
|
func detectTone(_ text: String) async -> Tone
|
||||||
|
|
||||||
|
/// Optionally preload the model so the first real action is instant.
|
||||||
|
func warmup() async
|
||||||
|
|
||||||
|
/// Streaming variant: yields incremental output as it generates. The default
|
||||||
|
/// implementation wraps `run` as a single chunk, so a backend that doesn't
|
||||||
|
/// natively stream still conforms.
|
||||||
|
func runStreaming(_ action: AIAction, on text: String, protectedTerms: [String])
|
||||||
|
-> AsyncThrowingStream<String, Error>
|
||||||
|
}
|
||||||
|
|
||||||
|
extension AIEngine {
|
||||||
|
/// Convenience: run with no protected terms.
|
||||||
|
func run(_ action: AIAction, on text: String) async throws -> String {
|
||||||
|
try await run(action, on: text, protectedTerms: [])
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Default: no preload.
|
||||||
|
func warmup() async {}
|
||||||
|
|
||||||
|
/// Default streaming: produce the full result, then emit it as one chunk.
|
||||||
|
func runStreaming(_ action: AIAction, on text: String, protectedTerms: [String])
|
||||||
|
-> AsyncThrowingStream<String, Error> {
|
||||||
|
AsyncThrowingStream { continuation in
|
||||||
|
Task {
|
||||||
|
do {
|
||||||
|
continuation.yield(try await run(action, on: text, protectedTerms: protectedTerms))
|
||||||
|
continuation.finish()
|
||||||
|
} catch {
|
||||||
|
continuation.finish(throwing: error)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Convenience streaming overload with no protected terms.
|
||||||
|
func runStreaming(_ action: AIAction, on text: String) -> AsyncThrowingStream<String, Error> {
|
||||||
|
runStreaming(action, on: text, protectedTerms: [])
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
enum AIEngineError: LocalizedError {
|
enum AIEngineError: LocalizedError {
|
||||||
|
|||||||
111
Sources/GrammarGem/AI/AIPromptKit.swift
Normal file
111
Sources/GrammarGem/AI/AIPromptKit.swift
Normal file
@ -0,0 +1,111 @@
|
|||||||
|
import Foundation
|
||||||
|
|
||||||
|
/// Backend-neutral prompt construction, generation parameters, and output
|
||||||
|
/// cleanup for Layer-2 actions. Extracted from `MLXEngine` so EVERY `AIEngine`
|
||||||
|
/// backend (MLX today, Ollama/llama.cpp next, the localhost extension server) is
|
||||||
|
/// guaranteed to produce identical prompts and identical output trimming — no
|
||||||
|
/// behavioral drift between runtimes — and so the personal dictionary / a future
|
||||||
|
/// style guide can inject protected vocabulary in exactly ONE place.
|
||||||
|
///
|
||||||
|
/// Everything here is a pure function of its inputs (no MLX, no I/O), which also
|
||||||
|
/// makes per-action behavior unit-testable without a model present.
|
||||||
|
enum AIPrompts {
|
||||||
|
|
||||||
|
/// The system prompt for an action. `protectedTerms` (personal-dictionary
|
||||||
|
/// entries + future style-guide vocabulary) are terms the model must never
|
||||||
|
/// alter, so AI rewrites stop "correcting" brand names and jargon.
|
||||||
|
static func systemPrompt(for action: AIAction, protectedTerms: [String] = []) -> String {
|
||||||
|
base(for: action) + protectedTermsClause(protectedTerms)
|
||||||
|
}
|
||||||
|
|
||||||
|
private static func base(for action: AIAction) -> String {
|
||||||
|
switch action {
|
||||||
|
case .rewriteClarity:
|
||||||
|
return "Rewrite the user's text to be clearer and more concise. Preserve the "
|
||||||
|
+ "original meaning and the author's voice. Output only the rewritten text."
|
||||||
|
case .rewrite:
|
||||||
|
return "Paraphrase the user's text. Keep the meaning but vary the wording and "
|
||||||
|
+ "structure. Output only the rewritten text."
|
||||||
|
case .adjustTone(let tone):
|
||||||
|
return "Rewrite the user's text in a \(tone.rawValue) tone. Preserve the meaning. "
|
||||||
|
+ "Output only the rewritten text."
|
||||||
|
case .ask(let instruction):
|
||||||
|
return "Apply this instruction to the user's text: \(instruction). "
|
||||||
|
+ "Output only the resulting text."
|
||||||
|
case .translate(let language):
|
||||||
|
return "Translate the user's text into \(language). Output only the translation."
|
||||||
|
case .applyMode(let mode):
|
||||||
|
return mode.systemPrompt
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Appended to every system prompt: terms the model must reproduce verbatim
|
||||||
|
/// (brand names, jargon the user added to their personal dictionary). Capped
|
||||||
|
/// so a large dictionary can't blow the context window.
|
||||||
|
private static func protectedTermsClause(_ terms: [String]) -> String {
|
||||||
|
let cleaned = terms
|
||||||
|
.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
|
||||||
|
.filter { !$0.isEmpty }
|
||||||
|
guard !cleaned.isEmpty else { return "" }
|
||||||
|
let list = cleaned.prefix(50).map { "\"\($0)\"" }.joined(separator: ", ")
|
||||||
|
return " Never change, translate, or \"correct\" these exact terms — reproduce them "
|
||||||
|
+ "verbatim, including capitalization: \(list)."
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Low temperature for faithful rewrites; a touch higher for free paraphrase.
|
||||||
|
static func temperature(for action: AIAction) -> Float {
|
||||||
|
switch action {
|
||||||
|
case .rewrite: return 0.5
|
||||||
|
case .ask: return 0.4
|
||||||
|
default: return 0.2
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Generation token budget. Honors a Mode's `lengthCap` (in characters) for
|
||||||
|
/// short-form modes (Post = 280, Team Chat = 60) so we don't spend the full
|
||||||
|
/// 512-token default on a one-line message; falls back to the default otherwise.
|
||||||
|
static func maxTokens(for action: AIAction) -> Int {
|
||||||
|
let fallback = 512
|
||||||
|
guard case .applyMode(let mode) = action, let cap = mode.lengthCap else { return fallback }
|
||||||
|
// ~1 token per 3 characters of English, plus headroom; clamp to a floor.
|
||||||
|
return min(fallback, max(64, cap / 3 + 48))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Strip a leading "Sure, here's …:" preamble and surrounding quotes that small
|
||||||
|
/// instruct models sometimes add, so the result drops cleanly into the field.
|
||||||
|
static func clean(_ s: String) -> String {
|
||||||
|
var out = s.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
|
let lower = out.lowercased()
|
||||||
|
for prefix in ["sure, here", "sure! here", "here is", "here's", "certainly"] {
|
||||||
|
if lower.hasPrefix(prefix), let colon = out.firstIndex(of: ":") {
|
||||||
|
let after = out[out.index(after: colon)...]
|
||||||
|
.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
|
if !after.isEmpty { out = after }
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if out.count >= 2, out.first == "\"", out.last == "\"" {
|
||||||
|
out = String(out.dropFirst().dropLast())
|
||||||
|
}
|
||||||
|
return out.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The tone-classifier system prompt (kept here so any backend reuses it).
|
||||||
|
static func toneClassifierPrompt() -> String {
|
||||||
|
let options = Tone.allCases.map(\.rawValue).joined(separator: ", ")
|
||||||
|
return """
|
||||||
|
You are a tone classifier. Read the user's text and respond with exactly ONE \
|
||||||
|
word — the single best-matching tone from this list: \(options). \
|
||||||
|
Output only that word, lowercase, no punctuation or explanation.
|
||||||
|
"""
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse the one-word tone classification output into a `Tone`.
|
||||||
|
static func parseTone(_ raw: String) -> Tone {
|
||||||
|
let token = raw.lowercased().trimmingCharacters(in: CharacterSet.alphanumerics.inverted)
|
||||||
|
if let exact = Tone(rawValue: token) { return exact }
|
||||||
|
let lower = raw.lowercased()
|
||||||
|
for tone in Tone.allCases where lower.contains(tone.rawValue) { return tone }
|
||||||
|
return .professional
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -25,19 +25,27 @@ final class CustomModesStore: ObservableObject {
|
|||||||
/// Built-in + custom modes (custom ones are always paid).
|
/// Built-in + custom modes (custom ones are always paid).
|
||||||
func allModes() -> [WritingMode] { ModeRegistry.all + customModes }
|
func allModes() -> [WritingMode] { ModeRegistry.all + customModes }
|
||||||
|
|
||||||
func addCustom(name: String, prompt: String) {
|
/// Returns false (no-op) if the name or prompt is empty, or the name
|
||||||
|
/// collides with an existing built-in/custom mode — so an empty-prompt mode
|
||||||
|
/// can't slip through to the LLM as an empty system prompt.
|
||||||
|
@discardableResult
|
||||||
|
func addCustom(name: String, prompt: String) -> Bool {
|
||||||
let trimmed = name.trimmingCharacters(in: .whitespacesAndNewlines)
|
let trimmed = name.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
guard !trimmed.isEmpty else { return }
|
let trimmedPrompt = prompt.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
|
guard !trimmed.isEmpty, !trimmedPrompt.isEmpty else { return false }
|
||||||
|
guard !allModes().contains(where: { $0.name.caseInsensitiveCompare(trimmed) == .orderedSame })
|
||||||
|
else { return false }
|
||||||
let mode = WritingMode(
|
let mode = WritingMode(
|
||||||
id: "custom-\(UUID().uuidString.prefix(8))",
|
id: "custom-\(UUID().uuidString.prefix(8))",
|
||||||
name: trimmed,
|
name: trimmed,
|
||||||
systemPrompt: prompt.trimmingCharacters(in: .whitespacesAndNewlines),
|
systemPrompt: trimmedPrompt,
|
||||||
isPaid: true,
|
isPaid: true,
|
||||||
lengthCap: nil,
|
lengthCap: nil,
|
||||||
autoFormat: "Custom"
|
autoFormat: "Custom"
|
||||||
)
|
)
|
||||||
customModes.append(mode)
|
customModes.append(mode)
|
||||||
persistModes()
|
persistModes()
|
||||||
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
func removeCustom(id: String) {
|
func removeCustom(id: String) {
|
||||||
|
|||||||
@ -33,33 +33,63 @@ final class MLXEngine: AIEngine {
|
|||||||
|
|
||||||
// MARK: - AIEngine
|
// MARK: - AIEngine
|
||||||
|
|
||||||
func run(_ action: AIAction, on text: String) async throws -> String {
|
func run(_ action: AIAction, on text: String, protectedTerms: [String]) async throws -> String {
|
||||||
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
guard !trimmed.isEmpty else { return text }
|
guard !trimmed.isEmpty else { return text }
|
||||||
|
|
||||||
let container = try await container()
|
let container = try await container()
|
||||||
let params = GenerateParameters(temperature: temperature(for: action))
|
let params = GenerateParameters(temperature: AIPrompts.temperature(for: action))
|
||||||
let raw = try await complete(
|
let raw = try await complete(
|
||||||
system: systemPrompt(for: action), user: trimmed,
|
system: AIPrompts.systemPrompt(for: action, protectedTerms: protectedTerms),
|
||||||
maxTokens: 512, parameters: params, container: container)
|
user: trimmed,
|
||||||
return clean(raw)
|
maxTokens: AIPrompts.maxTokens(for: action),
|
||||||
|
parameters: params, container: container)
|
||||||
|
return AIPrompts.clean(raw)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Preload the model container so the first real action doesn't pay the
|
||||||
|
/// multi-second load. Safe to call repeatedly (the loader caches).
|
||||||
|
func warmup() async {
|
||||||
|
_ = try? await container()
|
||||||
|
}
|
||||||
|
|
||||||
|
func runStreaming(_ action: AIAction, on text: String, protectedTerms: [String])
|
||||||
|
-> AsyncThrowingStream<String, Error> {
|
||||||
|
AsyncThrowingStream { continuation in
|
||||||
|
Task {
|
||||||
|
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
|
guard !trimmed.isEmpty else { continuation.yield(text); continuation.finish(); return }
|
||||||
|
do {
|
||||||
|
let container = try await container()
|
||||||
|
var params = GenerateParameters(temperature: AIPrompts.temperature(for: action))
|
||||||
|
params.maxTokens = AIPrompts.maxTokens(for: action)
|
||||||
|
let system = AIPrompts.systemPrompt(for: action, protectedTerms: protectedTerms)
|
||||||
|
try await container.perform { context in
|
||||||
|
let input = try await context.processor.prepare(
|
||||||
|
input: UserInput(chat: [.system(system), .user(trimmed)]))
|
||||||
|
let stream = try MLXLMCommon.generate(
|
||||||
|
input: input, parameters: params, context: context)
|
||||||
|
for await item in stream {
|
||||||
|
if let chunk = item.chunk { continuation.yield(chunk) }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continuation.finish()
|
||||||
|
} catch {
|
||||||
|
continuation.finish(throwing: error)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func detectTone(_ text: String) async -> Tone {
|
func detectTone(_ text: String) async -> Tone {
|
||||||
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
guard !trimmed.isEmpty else { return .professional }
|
guard !trimmed.isEmpty else { return .professional }
|
||||||
let options = Tone.allCases.map(\.rawValue).joined(separator: ", ")
|
|
||||||
let system = """
|
|
||||||
You are a tone classifier. Read the user's text and respond with exactly ONE \
|
|
||||||
word — the single best-matching tone from this list: \(options). \
|
|
||||||
Output only that word, lowercase, no punctuation or explanation.
|
|
||||||
"""
|
|
||||||
do {
|
do {
|
||||||
let container = try await container()
|
let container = try await container()
|
||||||
let raw = try await complete(
|
let raw = try await complete(
|
||||||
system: system, user: trimmed, maxTokens: 8,
|
system: AIPrompts.toneClassifierPrompt(), user: trimmed, maxTokens: 8,
|
||||||
parameters: GenerateParameters(temperature: 0.0), container: container)
|
parameters: GenerateParameters(temperature: 0.0), container: container)
|
||||||
return parseTone(raw)
|
return AIPrompts.parseTone(raw)
|
||||||
} catch {
|
} catch {
|
||||||
return .professional
|
return .professional
|
||||||
}
|
}
|
||||||
@ -98,64 +128,9 @@ final class MLXEngine: AIEngine {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// MARK: - Prompting
|
// Prompt construction, temperature, token budget, and output cleanup now live
|
||||||
|
// in `AIPrompts` (backend-neutral; see AIPromptKit.swift) so every AIEngine
|
||||||
private func systemPrompt(for action: AIAction) -> String {
|
// backend produces identical prompts and identical output trimming.
|
||||||
switch action {
|
|
||||||
case .rewriteClarity:
|
|
||||||
return "Rewrite the user's text to be clearer and more concise. Preserve the "
|
|
||||||
+ "original meaning and the author's voice. Output only the rewritten text."
|
|
||||||
case .rewrite:
|
|
||||||
return "Paraphrase the user's text. Keep the meaning but vary the wording and "
|
|
||||||
+ "structure. Output only the rewritten text."
|
|
||||||
case .adjustTone(let tone):
|
|
||||||
return "Rewrite the user's text in a \(tone.rawValue) tone. Preserve the meaning. "
|
|
||||||
+ "Output only the rewritten text."
|
|
||||||
case .ask(let instruction):
|
|
||||||
return "Apply this instruction to the user's text: \(instruction). "
|
|
||||||
+ "Output only the resulting text."
|
|
||||||
case .translate(let language):
|
|
||||||
return "Translate the user's text into \(language). Output only the translation."
|
|
||||||
case .applyMode(let mode):
|
|
||||||
return mode.systemPrompt
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Low temperature for faithful rewrites; a touch higher for free paraphrase.
|
|
||||||
private func temperature(for action: AIAction) -> Float {
|
|
||||||
switch action {
|
|
||||||
case .rewrite: return 0.5
|
|
||||||
case .ask: return 0.4
|
|
||||||
default: return 0.2
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// MARK: - Output cleanup
|
|
||||||
|
|
||||||
private func clean(_ s: String) -> String {
|
|
||||||
var out = s.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
||||||
let lower = out.lowercased()
|
|
||||||
for prefix in ["sure, here", "sure! here", "here is", "here's", "certainly"] {
|
|
||||||
if lower.hasPrefix(prefix), let colon = out.firstIndex(of: ":") {
|
|
||||||
let after = out[out.index(after: colon)...]
|
|
||||||
.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
||||||
if !after.isEmpty { out = after }
|
|
||||||
break
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if out.count >= 2, out.first == "\"", out.last == "\"" {
|
|
||||||
out = String(out.dropFirst().dropLast())
|
|
||||||
}
|
|
||||||
return out.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
||||||
}
|
|
||||||
|
|
||||||
private func parseTone(_ raw: String) -> Tone {
|
|
||||||
let token = raw.lowercased().trimmingCharacters(in: CharacterSet.alphanumerics.inverted)
|
|
||||||
if let exact = Tone(rawValue: token) { return exact }
|
|
||||||
let lower = raw.lowercased()
|
|
||||||
for tone in Tone.allCases where lower.contains(tone.rawValue) { return tone }
|
|
||||||
return .professional
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Caches one loaded `ModelContainer` (keyed by directory) and coalesces
|
/// Caches one loaded `ModelContainer` (keyed by directory) and coalesces
|
||||||
|
|||||||
@ -14,7 +14,9 @@ final class ModelManager: ObservableObject {
|
|||||||
|
|
||||||
@Published private(set) var state: State = .notDownloaded
|
@Published private(set) var state: State = .notDownloaded
|
||||||
@Published private(set) var statusText: String = ""
|
@Published private(set) var statusText: String = ""
|
||||||
@Published var selectedRepo: String = AppConfig.Model.defaultRepo
|
@Published var selectedRepo: String = ModelManager.activeRepo() {
|
||||||
|
didSet { UserDefaults.standard.set(selectedRepo, forKey: Self.selectedRepoKey) }
|
||||||
|
}
|
||||||
|
|
||||||
private var task: Task<Void, Never>?
|
private var task: Task<Void, Never>?
|
||||||
private var generation = 0
|
private var generation = 0
|
||||||
@ -55,6 +57,33 @@ final class ModelManager: ObservableObject {
|
|||||||
isModelPresent(repo: repo) ? modelDir(repo) : nil
|
isModelPresent(repo: repo) ? modelDir(repo) : nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// UserDefaults key backing the user's chosen model repo.
|
||||||
|
nonisolated static let selectedRepoKey = "GrammarGem.selectedModelRepo"
|
||||||
|
|
||||||
|
/// The user's currently-selected repo, readable from ANY thread (it lives in
|
||||||
|
/// UserDefaults, which is thread-safe). The AI engine's model-directory
|
||||||
|
/// provider reads this so download, load, and `isReady` all agree on the SAME
|
||||||
|
/// repo — including the paid large model, which the old hardcoded-default
|
||||||
|
/// provider silently ignored. Falls back to the default repo when unset.
|
||||||
|
nonisolated static func activeRepo() -> String {
|
||||||
|
UserDefaults.standard.string(forKey: selectedRepoKey) ?? AppConfig.Model.defaultRepo
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Safe URL construction for the two Hugging Face endpoints. A malformed
|
||||||
|
/// repo/path yields nil → a thrown `ModelError`, never a force-unwrap crash.
|
||||||
|
nonisolated static func fileURL(repo: String, path: String) -> URL? {
|
||||||
|
var c = URLComponents(string: AppConfig.Model.huggingFaceBase)
|
||||||
|
c?.path = "/\(repo)/resolve/main/\(path)"
|
||||||
|
return c?.url
|
||||||
|
}
|
||||||
|
|
||||||
|
nonisolated static func treeURL(repo: String) -> URL? {
|
||||||
|
var c = URLComponents(string: AppConfig.Model.huggingFaceBase)
|
||||||
|
c?.path = "/api/models/\(repo)/tree/main"
|
||||||
|
c?.queryItems = [URLQueryItem(name: "recursive", value: "1")]
|
||||||
|
return c?.url
|
||||||
|
}
|
||||||
|
|
||||||
/// Kick off (or resume to) a real download. Safe to call from the UI.
|
/// Kick off (or resume to) a real download. Safe to call from the UI.
|
||||||
func startDownload(repo: String? = nil) {
|
func startDownload(repo: String? = nil) {
|
||||||
let target = repo ?? selectedRepo
|
let target = repo ?? selectedRepo
|
||||||
@ -92,7 +121,9 @@ final class ModelManager: ObservableObject {
|
|||||||
try Task.checkCancellation()
|
try Task.checkCancellation()
|
||||||
set(gen, .downloading(progress: min(1, Double(done) / Double(total))),
|
set(gen, .downloading(progress: min(1, Double(done) / Double(total))),
|
||||||
"Downloading \(file.path) (\(byteString(file.size)))…")
|
"Downloading \(file.path) (\(byteString(file.size)))…")
|
||||||
let url = URL(string: "\(AppConfig.Model.huggingFaceBase)/\(repo)/resolve/main/\(file.path)")!
|
guard let url = Self.fileURL(repo: repo, path: file.path) else {
|
||||||
|
throw ModelError.invalidURL(file.path)
|
||||||
|
}
|
||||||
let (tmp, response) = try await URLSession.shared.download(from: url)
|
let (tmp, response) = try await URLSession.shared.download(from: url)
|
||||||
if let http = response as? HTTPURLResponse, http.statusCode != 200 {
|
if let http = response as? HTTPURLResponse, http.statusCode != 200 {
|
||||||
throw ModelError.http(http.statusCode, file.path)
|
throw ModelError.http(http.statusCode, file.path)
|
||||||
@ -148,7 +179,7 @@ final class ModelManager: ObservableObject {
|
|||||||
|
|
||||||
/// List the repo's files (+ sizes) from the Hugging Face tree API.
|
/// List the repo's files (+ sizes) from the Hugging Face tree API.
|
||||||
private func fetchFileList(repo: String) async throws -> [(path: String, size: Int64)] {
|
private func fetchFileList(repo: String) async throws -> [(path: String, size: Int64)] {
|
||||||
let url = URL(string: "\(AppConfig.Model.huggingFaceBase)/api/models/\(repo)/tree/main?recursive=1")!
|
guard let url = Self.treeURL(repo: repo) else { throw ModelError.invalidURL("file list") }
|
||||||
let (data, response) = try await URLSession.shared.data(from: url)
|
let (data, response) = try await URLSession.shared.data(from: url)
|
||||||
if let http = response as? HTTPURLResponse, http.statusCode != 200 {
|
if let http = response as? HTTPURLResponse, http.statusCode != 200 {
|
||||||
throw ModelError.http(http.statusCode, "file list")
|
throw ModelError.http(http.statusCode, "file list")
|
||||||
@ -166,11 +197,13 @@ final class ModelManager: ObservableObject {
|
|||||||
case empty
|
case empty
|
||||||
case http(Int, String)
|
case http(Int, String)
|
||||||
case corrupt(String)
|
case corrupt(String)
|
||||||
|
case invalidURL(String)
|
||||||
var errorDescription: String? {
|
var errorDescription: String? {
|
||||||
switch self {
|
switch self {
|
||||||
case .empty: return "No model files were found for this repository."
|
case .empty: return "No model files were found for this repository."
|
||||||
case .http(let code, let what): return "Download failed (\(code)) while fetching \(what)."
|
case .http(let code, let what): return "Download failed (\(code)) while fetching \(what)."
|
||||||
case .corrupt(let what): return "Downloaded file looked wrong (\(what)); please retry."
|
case .corrupt(let what): return "Downloaded file looked wrong (\(what)); please retry."
|
||||||
|
case .invalidURL(let what): return "Couldn't build a valid download URL for \(what)."
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
176
Sources/GrammarGem/AI/OllamaEngine.swift
Normal file
176
Sources/GrammarGem/AI/OllamaEngine.swift
Normal file
@ -0,0 +1,176 @@
|
|||||||
|
import Foundation
|
||||||
|
|
||||||
|
/// A cross-platform Layer-2 backend that talks to a locally-running **Ollama**
|
||||||
|
/// server over HTTP. Conforms to the same `AIEngine` seam as `MLXEngine`, reuses
|
||||||
|
/// the backend-neutral `AIPrompts` (so prompts/output are identical across
|
||||||
|
/// runtimes), and depends only on `URLSession` — no MLX, no Metal — so it compiles
|
||||||
|
/// and runs anywhere, including Windows/Linux. This is what converts the AIEngine
|
||||||
|
/// abstraction from nominal to real and opens the non-Apple-Silicon market.
|
||||||
|
///
|
||||||
|
/// Privacy is preserved: Ollama runs the model on the user's own machine on
|
||||||
|
/// loopback (127.0.0.1); no text leaves the device.
|
||||||
|
final class OllamaEngine: AIEngine {
|
||||||
|
private let base: URL
|
||||||
|
private let model: String
|
||||||
|
private let session: URLSession
|
||||||
|
|
||||||
|
/// `isReady` is read synchronously by the coordinator but readiness for an HTTP
|
||||||
|
/// backend is an async probe, so we cache the last-known result behind a lock.
|
||||||
|
private let lock = NSLock()
|
||||||
|
private var cachedReady = false
|
||||||
|
|
||||||
|
init(base: URL = AppConfig.Ollama.baseURL,
|
||||||
|
model: String = AppConfig.Ollama.defaultModel,
|
||||||
|
session: URLSession = .shared) {
|
||||||
|
self.base = base
|
||||||
|
self.model = model
|
||||||
|
self.session = session
|
||||||
|
}
|
||||||
|
|
||||||
|
var isReady: Bool {
|
||||||
|
lock.lock(); defer { lock.unlock() }
|
||||||
|
return cachedReady
|
||||||
|
}
|
||||||
|
|
||||||
|
private func setReady(_ value: Bool) {
|
||||||
|
lock.lock(); cachedReady = value; lock.unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - AIEngine
|
||||||
|
|
||||||
|
/// Probe the server and confirm the configured model is pulled. Readiness for
|
||||||
|
/// Ollama means "server up + model present" (vs. MLX's "weights on disk").
|
||||||
|
func warmup() async {
|
||||||
|
setReady(await probeModelPresent())
|
||||||
|
}
|
||||||
|
|
||||||
|
func run(_ action: AIAction, on text: String, protectedTerms: [String]) async throws -> String {
|
||||||
|
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
|
guard !trimmed.isEmpty else { return text }
|
||||||
|
let raw = try await chat(
|
||||||
|
system: AIPrompts.systemPrompt(for: action, protectedTerms: protectedTerms),
|
||||||
|
user: trimmed,
|
||||||
|
temperature: AIPrompts.temperature(for: action),
|
||||||
|
maxTokens: AIPrompts.maxTokens(for: action))
|
||||||
|
setReady(true)
|
||||||
|
return AIPrompts.clean(raw)
|
||||||
|
}
|
||||||
|
|
||||||
|
func detectTone(_ text: String) async -> Tone {
|
||||||
|
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
|
guard !trimmed.isEmpty else { return .professional }
|
||||||
|
do {
|
||||||
|
let raw = try await chat(
|
||||||
|
system: AIPrompts.toneClassifierPrompt(), user: trimmed,
|
||||||
|
temperature: 0.0, maxTokens: 8)
|
||||||
|
return AIPrompts.parseTone(raw)
|
||||||
|
} catch {
|
||||||
|
return .professional
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func runStreaming(_ action: AIAction, on text: String, protectedTerms: [String])
|
||||||
|
-> AsyncThrowingStream<String, Error> {
|
||||||
|
AsyncThrowingStream { continuation in
|
||||||
|
Task {
|
||||||
|
let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||||
|
guard !trimmed.isEmpty else { continuation.yield(text); continuation.finish(); return }
|
||||||
|
do {
|
||||||
|
try await streamChat(
|
||||||
|
system: AIPrompts.systemPrompt(for: action, protectedTerms: protectedTerms),
|
||||||
|
user: trimmed,
|
||||||
|
temperature: AIPrompts.temperature(for: action),
|
||||||
|
maxTokens: AIPrompts.maxTokens(for: action),
|
||||||
|
onChunk: { continuation.yield($0) })
|
||||||
|
setReady(true)
|
||||||
|
continuation.finish()
|
||||||
|
} catch {
|
||||||
|
continuation.finish(throwing: error)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - HTTP
|
||||||
|
|
||||||
|
private func probeModelPresent() async -> Bool {
|
||||||
|
do {
|
||||||
|
let (data, response) = try await session.data(from: base.appendingPathComponent("api/tags"))
|
||||||
|
guard (response as? HTTPURLResponse)?.statusCode == 200 else { return false }
|
||||||
|
let names = (try JSONDecoder().decode(TagsResponse.self, from: data)).models.map(\.name)
|
||||||
|
// Match exact or family (e.g. configured "qwen2.5:3b" vs a pulled "qwen2.5:3b-instruct").
|
||||||
|
return names.contains(model)
|
||||||
|
|| names.contains { $0.hasPrefix(model) || model.hasPrefix($0) }
|
||||||
|
} catch {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private func chat(system: String, user: String,
|
||||||
|
temperature: Float, maxTokens: Int) async throws -> String {
|
||||||
|
let request = try makeRequest(path: "api/chat", body: ChatRequest(
|
||||||
|
model: model,
|
||||||
|
messages: [.init(role: "system", content: system), .init(role: "user", content: user)],
|
||||||
|
stream: false,
|
||||||
|
options: .init(temperature: temperature, num_predict: maxTokens)))
|
||||||
|
let (data, response) = try await session.data(for: request)
|
||||||
|
guard let http = response as? HTTPURLResponse else {
|
||||||
|
throw AIEngineError.generationFailed("No HTTP response from Ollama at \(base.absoluteString)")
|
||||||
|
}
|
||||||
|
guard http.statusCode == 200 else {
|
||||||
|
throw AIEngineError.generationFailed("Ollama returned HTTP \(http.statusCode)")
|
||||||
|
}
|
||||||
|
return (try JSONDecoder().decode(ChatResponse.self, from: data)).message?.content ?? ""
|
||||||
|
}
|
||||||
|
|
||||||
|
private func streamChat(system: String, user: String,
|
||||||
|
temperature: Float, maxTokens: Int,
|
||||||
|
onChunk: (String) -> Void) async throws {
|
||||||
|
let request = try makeRequest(path: "api/chat", body: ChatRequest(
|
||||||
|
model: model,
|
||||||
|
messages: [.init(role: "system", content: system), .init(role: "user", content: user)],
|
||||||
|
stream: true,
|
||||||
|
options: .init(temperature: temperature, num_predict: maxTokens)))
|
||||||
|
let (bytes, response) = try await session.bytes(for: request)
|
||||||
|
guard (response as? HTTPURLResponse)?.statusCode == 200 else {
|
||||||
|
throw AIEngineError.generationFailed("Ollama streaming request failed")
|
||||||
|
}
|
||||||
|
// Ollama streams newline-delimited JSON frames, one per token-ish chunk.
|
||||||
|
for try await line in bytes.lines {
|
||||||
|
guard let data = line.data(using: .utf8),
|
||||||
|
let frame = try? JSONDecoder().decode(ChatResponse.self, from: data),
|
||||||
|
let chunk = frame.message?.content, !chunk.isEmpty
|
||||||
|
else { continue }
|
||||||
|
onChunk(chunk)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private func makeRequest<T: Encodable>(path: String, body: T) throws -> URLRequest {
|
||||||
|
var request = URLRequest(url: base.appendingPathComponent(path))
|
||||||
|
request.httpMethod = "POST"
|
||||||
|
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
||||||
|
request.httpBody = try JSONEncoder().encode(body)
|
||||||
|
return request
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - Wire types
|
||||||
|
|
||||||
|
private struct ChatRequest: Encodable {
|
||||||
|
let model: String
|
||||||
|
let messages: [Message]
|
||||||
|
let stream: Bool
|
||||||
|
let options: Options
|
||||||
|
struct Message: Encodable { let role: String; let content: String }
|
||||||
|
struct Options: Encodable { let temperature: Float; let num_predict: Int }
|
||||||
|
}
|
||||||
|
|
||||||
|
private struct ChatResponse: Decodable {
|
||||||
|
let message: Message?
|
||||||
|
struct Message: Decodable { let content: String }
|
||||||
|
}
|
||||||
|
|
||||||
|
private struct TagsResponse: Decodable {
|
||||||
|
let models: [Model]
|
||||||
|
struct Model: Decodable { let name: String }
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -30,6 +30,8 @@ final class AppState: ObservableObject {
|
|||||||
let updater: UpdaterManager
|
let updater: UpdaterManager
|
||||||
private let capture: TextCapture
|
private let capture: TextCapture
|
||||||
let coordinator: TextReplacementCoordinator
|
let coordinator: TextReplacementCoordinator
|
||||||
|
/// Loopback bridge to the browser extension (opt-in; nil unless enabled).
|
||||||
|
private var localServer: LocalEngineServer?
|
||||||
|
|
||||||
// UI state
|
// UI state
|
||||||
@Published var lastStatus: String = "Ready"
|
@Published var lastStatus: String = "Ready"
|
||||||
@ -48,9 +50,9 @@ final class AppState: ObservableObject {
|
|||||||
grammar = engine
|
grammar = engine
|
||||||
// The on-device LLM loads weights from ModelManager's local snapshot (if a
|
// The on-device LLM loads weights from ModelManager's local snapshot (if a
|
||||||
// complete download is present) — never the network. Resolved lazily on use.
|
// complete download is present) — never the network. Resolved lazily on use.
|
||||||
ai = MLXEngine(modelDirectoryProvider: {
|
// Reads the user's *selected* repo (default or paid large model) so the
|
||||||
ModelManager.completedModelDirectory(repo: AppConfig.Model.defaultRepo)
|
// model that gets loaded matches the one that gets downloaded.
|
||||||
})
|
ai = AppState.makeAIEngine()
|
||||||
permissions = Permissions()
|
permissions = Permissions()
|
||||||
model = ModelManager()
|
model = ModelManager()
|
||||||
license = LicenseManager()
|
license = LicenseManager()
|
||||||
@ -67,8 +69,29 @@ final class AppState: ObservableObject {
|
|||||||
exclusions = excl
|
exclusions = excl
|
||||||
liveMonitor = LiveMonitor(grammar: engine, detector: detector, exclusions: excl, capture: capture)
|
liveMonitor = LiveMonitor(grammar: engine, detector: detector, exclusions: excl, capture: capture)
|
||||||
updater = UpdaterManager()
|
updater = UpdaterManager()
|
||||||
|
let dict = dictionary
|
||||||
coordinator = TextReplacementCoordinator(
|
coordinator = TextReplacementCoordinator(
|
||||||
capture: capture, grammar: engine, ai: ai, gate: gate, detector: detector)
|
capture: capture, grammar: engine, ai: ai, gate: gate, detector: detector,
|
||||||
|
protectedTerms: { dict.entries })
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - Backend factory
|
||||||
|
|
||||||
|
/// The single place the Layer-2 backend is chosen. Defaults to the MLX
|
||||||
|
/// on-device engine (Apple Silicon); a user/CI override selects the
|
||||||
|
/// cross-platform `OllamaEngine`. Adding a backend is a one-line branch here
|
||||||
|
/// instead of edits scattered across construction sites. The MLX provider
|
||||||
|
/// reads the user's *selected* repo so the loaded model matches the downloaded
|
||||||
|
/// one (incl. the paid large model).
|
||||||
|
static func makeAIEngine() -> AIEngine {
|
||||||
|
let override = ProcessInfo.processInfo.environment["GRAMMARGEM_AI_BACKEND"]
|
||||||
|
?? UserDefaults.standard.string(forKey: "GrammarGem.aiBackend")
|
||||||
|
if override == "ollama" {
|
||||||
|
return OllamaEngine()
|
||||||
|
}
|
||||||
|
return MLXEngine(modelDirectoryProvider: {
|
||||||
|
ModelManager.completedModelDirectory(repo: ModelManager.activeRepo())
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Called once at launch (from the AppDelegate).
|
/// Called once at launch (from the AppDelegate).
|
||||||
@ -101,10 +124,19 @@ final class AppState: ObservableObject {
|
|||||||
|
|
||||||
// Auto-install the on-device AI model in the background on first use, so
|
// Auto-install the on-device AI model in the background on first use, so
|
||||||
// the AI features are ready without the user hunting for a button. The
|
// the AI features are ready without the user hunting for a button. The
|
||||||
// download is resumable/cancelable from the AI Model screen.
|
// download is resumable/cancelable from the AI Model screen. If a model is
|
||||||
if !model.isModelPresent(repo: model.selectedRepo) {
|
// already present, warm it up now so the first AI action is instant.
|
||||||
|
if model.isModelPresent(repo: model.selectedRepo) {
|
||||||
|
Task { [weak self] in await self?.ai.warmup() }
|
||||||
|
} else {
|
||||||
model.startDownload()
|
model.startDownload()
|
||||||
}
|
}
|
||||||
|
// Warm the model up as soon as a (first-run) download completes.
|
||||||
|
model.$state.dropFirst()
|
||||||
|
.sink { [weak self] state in
|
||||||
|
if case .ready = state { Task { [weak self] in await self?.ai.warmup() } }
|
||||||
|
}
|
||||||
|
.store(in: &cancellables)
|
||||||
|
|
||||||
permissions.refresh()
|
permissions.refresh()
|
||||||
appAwareEnabled = gate.appAwareEnabled
|
appAwareEnabled = gate.appAwareEnabled
|
||||||
@ -115,9 +147,57 @@ final class AppState: ObservableObject {
|
|||||||
}
|
}
|
||||||
|
|
||||||
Task { await license.validateIfNeeded() }
|
Task { await license.validateIfNeeded() }
|
||||||
|
startLocalServerIfEnabled()
|
||||||
Log.app.info("GrammarGem started. Licensed: \(self.license.isLicensed, privacy: .public)")
|
Log.app.info("GrammarGem started. Licensed: \(self.license.isLicensed, privacy: .public)")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// MARK: - Browser-extension bridge (opt-in)
|
||||||
|
|
||||||
|
/// Start the loopback engine server only when the user has opted in, so the app
|
||||||
|
/// ships no listening socket by default. The server reuses the SAME engines,
|
||||||
|
/// prompts, and entitlement gate as the menu-bar app.
|
||||||
|
private func startLocalServerIfEnabled() {
|
||||||
|
let enabled = ProcessInfo.processInfo.environment["GRAMMARGEM_LOCAL_SERVER"] == "1"
|
||||||
|
|| UserDefaults.standard.bool(forKey: AppConfig.LocalServer.enableKey)
|
||||||
|
guard enabled, localServer == nil else { return }
|
||||||
|
|
||||||
|
let server = LocalEngineServer(
|
||||||
|
port: AppConfig.LocalServer.port,
|
||||||
|
token: LocalEngineServer.loadOrCreateToken(),
|
||||||
|
grammarHandler: { [weak self] text in
|
||||||
|
await MainActor.run {
|
||||||
|
guard let self else { return (text, []) }
|
||||||
|
let corrected = self.grammar.correct(text)
|
||||||
|
let suggestions = self.grammar.check(text).map {
|
||||||
|
LocalEngineServer.WireSuggestion(
|
||||||
|
location: $0.location, length: $0.length,
|
||||||
|
original: $0.original, replacement: $0.replacement,
|
||||||
|
kind: $0.kind.rawValue, message: $0.message)
|
||||||
|
}
|
||||||
|
return (corrected, suggestions)
|
||||||
|
}
|
||||||
|
},
|
||||||
|
aiHandler: { [weak self] (req: LocalEngineServer.AIRequest) async -> LocalEngineServer.AIOutcome in
|
||||||
|
guard let self else { return .error("GrammarGem is shutting down.") }
|
||||||
|
return await self.handleServerAI(req)
|
||||||
|
})
|
||||||
|
server.start()
|
||||||
|
localServer = server
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Run an AI action requested by the extension — same gate + model as the app.
|
||||||
|
private func handleServerAI(_ req: LocalEngineServer.AIRequest) async -> LocalEngineServer.AIOutcome {
|
||||||
|
if case .denied(let reason) = gate.authorizeAI(req.action) { return .error(reason) }
|
||||||
|
guard ai.isReady else { return .error(AIEngineError.modelNotReady.localizedDescription) }
|
||||||
|
do {
|
||||||
|
let out = try await ai.run(req.action, on: req.text, protectedTerms: dictionary.entries)
|
||||||
|
gate.recordAIActionUsed()
|
||||||
|
return .ok(out)
|
||||||
|
} catch {
|
||||||
|
return .error(error.localizedDescription)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// MARK: - Hotkey handling
|
// MARK: - Hotkey handling
|
||||||
|
|
||||||
/// (Re)register the global hotkeys from the current preferences.
|
/// (Re)register the global hotkeys from the current preferences.
|
||||||
|
|||||||
@ -25,7 +25,7 @@ enum MLXSelfTest {
|
|||||||
static func run() {
|
static func run() {
|
||||||
let sample = "me and him was going to the store for to buy some milks"
|
let sample = "me and him was going to the store for to buy some milks"
|
||||||
let engine = MLXEngine(modelDirectoryProvider: {
|
let engine = MLXEngine(modelDirectoryProvider: {
|
||||||
ModelManager.completedModelDirectory(repo: AppConfig.Model.defaultRepo)
|
ModelManager.completedModelDirectory(repo: ModelManager.activeRepo())
|
||||||
})
|
})
|
||||||
guard engine.isReady else {
|
guard engine.isReady else {
|
||||||
FileHandle.standardError.write(Data("SELFTEST_FAIL: model not present\n".utf8))
|
FileHandle.standardError.write(Data("SELFTEST_FAIL: model not present\n".utf8))
|
||||||
|
|||||||
@ -37,6 +37,23 @@ enum AppConfig {
|
|||||||
static let huggingFaceBase = "https://huggingface.co"
|
static let huggingFaceBase = "https://huggingface.co"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The cross-platform Layer-2 backend: a locally-running Ollama server. Used
|
||||||
|
/// when MLX isn't available (Windows/Linux) or when the user opts in. Inference
|
||||||
|
/// stays on-device — Ollama runs the model on the user's own machine.
|
||||||
|
enum Ollama {
|
||||||
|
static let baseURL = URL(string: "http://127.0.0.1:11434")!
|
||||||
|
static let defaultModel = "qwen2.5:3b" // matches the MLX default model family
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Loopback engine server for the browser extension (off by default; see
|
||||||
|
/// `LocalEngineServer`). Exposes Harper + the local LLM to the extension over
|
||||||
|
/// 127.0.0.1 with a per-install token — no cloud round-trip.
|
||||||
|
enum LocalServer {
|
||||||
|
static let port: UInt16 = 47800
|
||||||
|
static let tokenKey = "GrammarGem.localServerToken"
|
||||||
|
static let enableKey = "GrammarGem.enableLocalServer"
|
||||||
|
}
|
||||||
|
|
||||||
/// Free-tier limits + license validation cadence (see spec §5/§6).
|
/// Free-tier limits + license validation cadence (see spec §5/§6).
|
||||||
enum Limits {
|
enum Limits {
|
||||||
static let freeDailyAIActions = 10
|
static let freeDailyAIActions = 10
|
||||||
|
|||||||
@ -58,16 +58,24 @@ final class LemonSqueezyClient {
|
|||||||
private func verifyAndResolveTier(meta: LSMeta?, licenseKey: LSLicenseKey?) throws -> Tier {
|
private func verifyAndResolveTier(meta: LSMeta?, licenseKey: LSLicenseKey?) throws -> Tier {
|
||||||
guard let meta else { throw LicenseError.wrongProduct }
|
guard let meta else { throw LicenseError.wrongProduct }
|
||||||
|
|
||||||
// Store + product check. If AppConfig still holds placeholder (non-numeric)
|
// Store + product check. With real numeric IDs configured this is a hard
|
||||||
// IDs, we log and skip the check so the stub flow runs — TODO(real-integration):
|
// gate that prevents a key from another Lemon Squeezy product unlocking
|
||||||
// once real numeric IDs are set, this becomes a hard gate.
|
// GrammarGem. If AppConfig still holds placeholder (non-numeric) IDs the
|
||||||
|
// gate can't run — allow the bypass ONLY in DEBUG so local stub flows work;
|
||||||
|
// in a release build FAIL CLOSED so the "any key unlocks paid" bypass can
|
||||||
|
// never ship. TODO(real-integration): set real IDs in AppConfig.LemonSqueezy.
|
||||||
if let expectedStore = Int(AppConfig.LemonSqueezy.expectedStoreID),
|
if let expectedStore = Int(AppConfig.LemonSqueezy.expectedStoreID),
|
||||||
let expectedProduct = Int(AppConfig.LemonSqueezy.expectedProductID) {
|
let expectedProduct = Int(AppConfig.LemonSqueezy.expectedProductID) {
|
||||||
guard meta.store_id == expectedStore, meta.product_id == expectedProduct else {
|
guard meta.store_id == expectedStore, meta.product_id == expectedProduct else {
|
||||||
throw LicenseError.wrongProduct
|
throw LicenseError.wrongProduct
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
Log.licensing.warning("Skipping store/product hard-verify — placeholder IDs in AppConfig.")
|
#if DEBUG
|
||||||
|
Log.licensing.warning("Skipping store/product hard-verify — placeholder IDs in AppConfig (DEBUG only).")
|
||||||
|
#else
|
||||||
|
Log.licensing.error("Refusing to verify license: placeholder Lemon Squeezy IDs in a release build.")
|
||||||
|
throw LicenseError.wrongProduct
|
||||||
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
// Resolve tier from the variant id (reverse-lookup the configured map).
|
// Resolve tier from the variant id (reverse-lookup the configured map).
|
||||||
|
|||||||
@ -38,6 +38,17 @@ enum AX {
|
|||||||
AXUIElementSetMessagingTimeout(element, seconds)
|
AXUIElementSetMessagingTimeout(element, seconds)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Whether an element is a secure text field (a password input). GrammarGem
|
||||||
|
/// must NEVER read or rewrite these — so a user who forgot to exclude a
|
||||||
|
/// password manager or login form can't have a decrypted password captured,
|
||||||
|
/// sent through the grammar/AI engines, and written back. This is enforced in
|
||||||
|
/// code rather than relying on the user's exclusion blocklist.
|
||||||
|
static func isSecureTextField(_ element: AXUIElement) -> Bool {
|
||||||
|
// AXSecureTextField is the subrole AppKit/most Electron inputs report for
|
||||||
|
// password fields. String-compared (the constant isn't always imported).
|
||||||
|
copyString(element, kAXSubroleAttribute) == "AXSecureTextField"
|
||||||
|
}
|
||||||
|
|
||||||
/// The currently-focused UI element across the system — robust to the
|
/// The currently-focused UI element across the system — robust to the
|
||||||
/// system-wide element failing.
|
/// system-wide element failing.
|
||||||
///
|
///
|
||||||
|
|||||||
266
Sources/GrammarGem/System/LocalEngineServer.swift
Normal file
266
Sources/GrammarGem/System/LocalEngineServer.swift
Normal file
@ -0,0 +1,266 @@
|
|||||||
|
import Foundation
|
||||||
|
import Network
|
||||||
|
|
||||||
|
/// A minimal, loopback-only HTTP server that exposes GrammarGem's local engines
|
||||||
|
/// (Harper grammar + the on-device LLM) to a **browser extension** — the
|
||||||
|
/// highest-volume surface for a grammar tool — with NO cloud round-trip. Built on
|
||||||
|
/// Network.framework so it adds no third-party HTTP dependency, keeping the
|
||||||
|
/// privacy/trust surface minimal.
|
||||||
|
///
|
||||||
|
/// Security posture (deliberate):
|
||||||
|
/// - Binds to 127.0.0.1 only (never a routable interface), via `requiredLocalEndpoint`.
|
||||||
|
/// - Every engine endpoint requires a per-install bearer token (`X-GrammarGem-Token`);
|
||||||
|
/// a web page that doesn't know the token cannot drive the engines.
|
||||||
|
/// - OFF by default — only starts when the user opts in, so the app ships no
|
||||||
|
/// listening socket unless explicitly enabled.
|
||||||
|
///
|
||||||
|
/// It reuses the SAME engines as the menu-bar app via injected async handlers, so
|
||||||
|
/// there is one engine, one set of prompts (`AIPrompts`), and one entitlement gate
|
||||||
|
/// across the app, the extension, and (later) other backends.
|
||||||
|
final class LocalEngineServer {
|
||||||
|
|
||||||
|
struct AIRequest { let action: AIAction; let text: String }
|
||||||
|
|
||||||
|
struct WireSuggestion: Encodable {
|
||||||
|
let location: Int
|
||||||
|
let length: Int
|
||||||
|
let original: String
|
||||||
|
let replacement: String
|
||||||
|
let kind: String
|
||||||
|
let message: String
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Grammar pass: corrected text + suggestion spans for the extension to underline.
|
||||||
|
typealias GrammarHandler = (String) async -> (corrected: String, suggestions: [WireSuggestion])
|
||||||
|
/// AI action outcome (a plain enum so a user-facing message needn't be an `Error`).
|
||||||
|
enum AIOutcome { case ok(String); case error(String) }
|
||||||
|
/// AI action: runs the gate + model, returns the result or a user-facing error.
|
||||||
|
typealias AIHandler = (AIRequest) async -> AIOutcome
|
||||||
|
|
||||||
|
private let port: UInt16
|
||||||
|
private let token: String
|
||||||
|
private let grammarHandler: GrammarHandler
|
||||||
|
private let aiHandler: AIHandler
|
||||||
|
|
||||||
|
private var listener: NWListener?
|
||||||
|
private let queue = DispatchQueue(label: "com.grammargem.localserver")
|
||||||
|
|
||||||
|
init(port: UInt16, token: String,
|
||||||
|
grammarHandler: @escaping GrammarHandler,
|
||||||
|
aiHandler: @escaping AIHandler) {
|
||||||
|
self.port = port
|
||||||
|
self.token = token
|
||||||
|
self.grammarHandler = grammarHandler
|
||||||
|
self.aiHandler = aiHandler
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A stable per-install token, persisted so the extension can be paired once.
|
||||||
|
static func loadOrCreateToken() -> String {
|
||||||
|
let d = UserDefaults.standard
|
||||||
|
if let existing = d.string(forKey: AppConfig.LocalServer.tokenKey), !existing.isEmpty {
|
||||||
|
return existing
|
||||||
|
}
|
||||||
|
let token = UUID().uuidString + UUID().uuidString.replacingOccurrences(of: "-", with: "")
|
||||||
|
d.set(token, forKey: AppConfig.LocalServer.tokenKey)
|
||||||
|
return token
|
||||||
|
}
|
||||||
|
|
||||||
|
func start() {
|
||||||
|
guard listener == nil else { return }
|
||||||
|
do {
|
||||||
|
let params = NWParameters.tcp
|
||||||
|
guard let nwPort = NWEndpoint.Port(rawValue: port) else { return }
|
||||||
|
// Loopback-only: refuse to bind to any routable interface.
|
||||||
|
params.requiredLocalEndpoint = .hostPort(host: "127.0.0.1", port: nwPort)
|
||||||
|
let listener = try NWListener(using: params)
|
||||||
|
listener.newConnectionHandler = { [weak self] conn in self?.accept(conn) }
|
||||||
|
listener.start(queue: queue)
|
||||||
|
self.listener = listener
|
||||||
|
Log.app.info("Local engine server on 127.0.0.1:\(self.port, privacy: .public) (extension bridge)")
|
||||||
|
} catch {
|
||||||
|
Log.app.error("Local engine server failed to start: \(error.localizedDescription, privacy: .public)")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func stop() {
|
||||||
|
listener?.cancel()
|
||||||
|
listener = nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - Connection lifecycle
|
||||||
|
|
||||||
|
private func accept(_ conn: NWConnection) {
|
||||||
|
conn.start(queue: queue)
|
||||||
|
receive(conn, buffer: Data())
|
||||||
|
}
|
||||||
|
|
||||||
|
private func receive(_ conn: NWConnection, buffer: Data) {
|
||||||
|
conn.receive(minimumIncompleteLength: 1, maximumLength: 1 << 16) { [weak self] data, _, isComplete, error in
|
||||||
|
guard let self else { conn.cancel(); return }
|
||||||
|
var buffer = buffer
|
||||||
|
if let data { buffer.append(data) }
|
||||||
|
if let request = HTTPRequest(buffer), request.isComplete {
|
||||||
|
Task { await self.route(request, on: conn) }
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if isComplete || error != nil { conn.cancel(); return }
|
||||||
|
self.receive(conn, buffer: buffer) // headers/body not fully arrived yet
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - Routing
|
||||||
|
|
||||||
|
private func route(_ request: HTTPRequest, on conn: NWConnection) async {
|
||||||
|
// CORS preflight — extensions issue this before a cross-origin POST.
|
||||||
|
if request.method == "OPTIONS" {
|
||||||
|
send(conn, status: 204, body: Data())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if request.path == "/health" || request.path == "/v1/health" {
|
||||||
|
sendJSON(conn, status: 200, HealthResponse(ok: true, app: AppConfig.appName, version: AppConfig.appVersion))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Token gate for every engine endpoint.
|
||||||
|
guard request.headers["x-grammargem-token"] == token else {
|
||||||
|
sendJSON(conn, status: 401, ErrorResponse(error: "unauthorized"))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
switch (request.method, request.path) {
|
||||||
|
case ("POST", "/v1/grammar"):
|
||||||
|
guard let text = decodeText(request.body) else {
|
||||||
|
sendJSON(conn, status: 400, ErrorResponse(error: "missing 'text'")); return
|
||||||
|
}
|
||||||
|
let result = await grammarHandler(text)
|
||||||
|
sendJSON(conn, status: 200, GrammarResponse(corrected: result.corrected, suggestions: result.suggestions))
|
||||||
|
|
||||||
|
case ("POST", "/v1/ai"):
|
||||||
|
guard let req = decodeAI(request.body) else {
|
||||||
|
sendJSON(conn, status: 400, ErrorResponse(error: "bad request")); return
|
||||||
|
}
|
||||||
|
switch await aiHandler(req) {
|
||||||
|
case .ok(let out): sendJSON(conn, status: 200, AIResponse(result: out))
|
||||||
|
case .error(let msg): sendJSON(conn, status: 422, ErrorResponse(error: msg))
|
||||||
|
}
|
||||||
|
|
||||||
|
default:
|
||||||
|
sendJSON(conn, status: 404, ErrorResponse(error: "not found"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - Body decoding
|
||||||
|
|
||||||
|
private func decodeText(_ body: Data) -> String? {
|
||||||
|
struct P: Decodable { let text: String }
|
||||||
|
guard let p = try? JSONDecoder().decode(P.self, from: body), !p.text.isEmpty else { return nil }
|
||||||
|
return p.text
|
||||||
|
}
|
||||||
|
|
||||||
|
private func decodeAI(_ body: Data) -> AIRequest? {
|
||||||
|
struct P: Decodable {
|
||||||
|
let action: String
|
||||||
|
let text: String
|
||||||
|
let instruction: String?
|
||||||
|
let language: String?
|
||||||
|
let mode: String?
|
||||||
|
let tone: String?
|
||||||
|
}
|
||||||
|
guard let p = try? JSONDecoder().decode(P.self, from: body), !p.text.isEmpty else { return nil }
|
||||||
|
let action: AIAction
|
||||||
|
switch p.action {
|
||||||
|
case "rewriteClarity": action = .rewriteClarity
|
||||||
|
case "rewrite": action = .rewrite
|
||||||
|
case "ask":
|
||||||
|
guard let i = p.instruction, !i.trimmingCharacters(in: .whitespaces).isEmpty else { return nil }
|
||||||
|
action = .ask(i)
|
||||||
|
case "translate":
|
||||||
|
guard let l = p.language, !l.isEmpty else { return nil }
|
||||||
|
action = .translate(language: l)
|
||||||
|
case "tone":
|
||||||
|
action = .adjustTone(Tone(rawValue: p.tone ?? "") ?? .professional)
|
||||||
|
case "mode":
|
||||||
|
guard let id = p.mode, let m = ModeRegistry.mode(id: id) else { return nil }
|
||||||
|
action = .applyMode(m)
|
||||||
|
default:
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return AIRequest(action: action, text: p.text)
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - Responses
|
||||||
|
|
||||||
|
private struct HealthResponse: Encodable { let ok: Bool; let app: String; let version: String }
|
||||||
|
private struct GrammarResponse: Encodable { let corrected: String; let suggestions: [WireSuggestion] }
|
||||||
|
private struct AIResponse: Encodable { let result: String }
|
||||||
|
private struct ErrorResponse: Encodable { let error: String }
|
||||||
|
|
||||||
|
private func sendJSON<T: Encodable>(_ conn: NWConnection, status: Int, _ value: T) {
|
||||||
|
let body = (try? JSONEncoder().encode(value)) ?? Data("{}".utf8)
|
||||||
|
send(conn, status: status, body: body)
|
||||||
|
}
|
||||||
|
|
||||||
|
private func send(_ conn: NWConnection, status: Int, body: Data) {
|
||||||
|
var head = "HTTP/1.1 \(status) \(Self.reason(status))\r\n"
|
||||||
|
head += "Content-Type: application/json\r\n"
|
||||||
|
head += "Content-Length: \(body.count)\r\n"
|
||||||
|
// The extension calls cross-origin; a secret token (not the origin) is the gate.
|
||||||
|
head += "Access-Control-Allow-Origin: *\r\n"
|
||||||
|
head += "Access-Control-Allow-Headers: Content-Type, X-GrammarGem-Token\r\n"
|
||||||
|
head += "Access-Control-Allow-Methods: POST, GET, OPTIONS\r\n"
|
||||||
|
head += "Connection: close\r\n\r\n"
|
||||||
|
var out = Data(head.utf8)
|
||||||
|
out.append(body)
|
||||||
|
conn.send(content: out, completion: .contentProcessed { _ in conn.cancel() })
|
||||||
|
}
|
||||||
|
|
||||||
|
private static func reason(_ status: Int) -> String {
|
||||||
|
switch status {
|
||||||
|
case 200: return "OK"
|
||||||
|
case 204: return "No Content"
|
||||||
|
case 400: return "Bad Request"
|
||||||
|
case 401: return "Unauthorized"
|
||||||
|
case 404: return "Not Found"
|
||||||
|
case 422: return "Unprocessable Entity"
|
||||||
|
default: return "Error"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A just-enough HTTP/1.1 request parser for small JSON POSTs. Returns nil until
|
||||||
|
/// the headers terminator (`\r\n\r\n`) has arrived; `isComplete` is false until the
|
||||||
|
/// full `Content-Length` body is buffered, so the receive loop knows to read more.
|
||||||
|
private struct HTTPRequest {
|
||||||
|
let method: String
|
||||||
|
let path: String
|
||||||
|
let headers: [String: String]
|
||||||
|
let body: Data
|
||||||
|
let isComplete: Bool
|
||||||
|
|
||||||
|
init?(_ raw: Data) {
|
||||||
|
guard let headerEnd = raw.range(of: Data("\r\n\r\n".utf8)) else { return nil }
|
||||||
|
guard let headerString = String(data: raw[..<headerEnd.lowerBound], encoding: .utf8) else { return nil }
|
||||||
|
var lines = headerString.components(separatedBy: "\r\n")
|
||||||
|
guard let requestLine = lines.first else { return nil }
|
||||||
|
let parts = requestLine.split(separator: " ")
|
||||||
|
guard parts.count >= 2 else { return nil }
|
||||||
|
method = String(parts[0])
|
||||||
|
path = String(parts[1])
|
||||||
|
lines.removeFirst()
|
||||||
|
|
||||||
|
var hdrs: [String: String] = [:]
|
||||||
|
for line in lines {
|
||||||
|
guard let colon = line.firstIndex(of: ":") else { continue }
|
||||||
|
let key = line[..<colon].trimmingCharacters(in: .whitespaces).lowercased()
|
||||||
|
let value = line[line.index(after: colon)...].trimmingCharacters(in: .whitespaces)
|
||||||
|
hdrs[key] = value
|
||||||
|
}
|
||||||
|
headers = hdrs
|
||||||
|
|
||||||
|
let contentLength = Int(hdrs["content-length"] ?? "0") ?? 0
|
||||||
|
let available = raw[headerEnd.upperBound...]
|
||||||
|
body = Data(available.prefix(contentLength))
|
||||||
|
isComplete = body.count >= contentLength
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -24,10 +24,12 @@ final class TextCapture {
|
|||||||
enum CaptureError: LocalizedError {
|
enum CaptureError: LocalizedError {
|
||||||
case noSelection
|
case noSelection
|
||||||
case notPermitted
|
case notPermitted
|
||||||
|
case secureField
|
||||||
var errorDescription: String? {
|
var errorDescription: String? {
|
||||||
switch self {
|
switch self {
|
||||||
case .noSelection: return "Select some text first, then press the hotkey."
|
case .noSelection: return "Select some text first, then press the hotkey."
|
||||||
case .notPermitted: return "GrammarGem needs Accessibility permission to read and replace text."
|
case .notPermitted: return "GrammarGem needs Accessibility permission to read and replace text."
|
||||||
|
case .secureField: return "GrammarGem never reads password fields."
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -41,8 +43,15 @@ final class TextCapture {
|
|||||||
// MARK: - Capture
|
// MARK: - Capture
|
||||||
|
|
||||||
/// Try AX first; on empty/unavailable, fall back to the clipboard path.
|
/// Try AX first; on empty/unavailable, fall back to the clipboard path.
|
||||||
|
/// Refuses outright if the focused element is a password field.
|
||||||
func capture() throws -> Capture {
|
func capture() throws -> Capture {
|
||||||
if let ax = captureViaAccessibility() {
|
let focused = AXIsProcessTrusted() ? AX.focusedElement() : nil
|
||||||
|
if let focused, AX.isSecureTextField(focused) {
|
||||||
|
// Bail before the clipboard fallback too — never synthesize ⌘C into a
|
||||||
|
// password field.
|
||||||
|
throw CaptureError.secureField
|
||||||
|
}
|
||||||
|
if let ax = captureViaAccessibility(focused) {
|
||||||
return ax
|
return ax
|
||||||
}
|
}
|
||||||
if let clip = try captureViaClipboard() {
|
if let clip = try captureViaClipboard() {
|
||||||
@ -51,9 +60,8 @@ final class TextCapture {
|
|||||||
throw CaptureError.noSelection
|
throw CaptureError.noSelection
|
||||||
}
|
}
|
||||||
|
|
||||||
private func captureViaAccessibility() -> Capture? {
|
private func captureViaAccessibility(_ focused: AXUIElement?) -> Capture? {
|
||||||
guard AXIsProcessTrusted() else { return nil }
|
guard let focused else { return nil }
|
||||||
guard let focused = AX.focusedElement() else { return nil }
|
|
||||||
|
|
||||||
var valueRef: CFTypeRef?
|
var valueRef: CFTypeRef?
|
||||||
guard AXUIElementCopyAttributeValue(
|
guard AXUIElementCopyAttributeValue(
|
||||||
@ -92,6 +100,7 @@ final class TextCapture {
|
|||||||
func focusedFieldText() -> (text: String, element: AXUIElement)? {
|
func focusedFieldText() -> (text: String, element: AXUIElement)? {
|
||||||
guard AXIsProcessTrusted() else { return nil }
|
guard AXIsProcessTrusted() else { return nil }
|
||||||
guard let focused = AX.focusedElement() else { return nil }
|
guard let focused = AX.focusedElement() else { return nil }
|
||||||
|
guard !AX.isSecureTextField(focused) else { return nil } // never fix a password field
|
||||||
|
|
||||||
// Only treat genuine editable text roles as fixable, so we don't report
|
// Only treat genuine editable text roles as fixable, so we don't report
|
||||||
// "nothing selected" failures on sliders/steppers/custom controls.
|
// "nothing selected" failures on sliders/steppers/custom controls.
|
||||||
|
|||||||
@ -47,7 +47,9 @@ enum TextFieldScanner {
|
|||||||
if role == kAXTextFieldRole as String
|
if role == kAXTextFieldRole as String
|
||||||
|| role == kAXTextAreaRole as String
|
|| role == kAXTextAreaRole as String
|
||||||
|| role == kAXComboBoxRole as String {
|
|| role == kAXComboBoxRole as String {
|
||||||
if let value = AX.copyString(el, kAXValueAttribute),
|
// Never surface password fields in the live monitor.
|
||||||
|
if !AX.isSecureTextField(el),
|
||||||
|
let value = AX.copyString(el, kAXValueAttribute),
|
||||||
value.trimmingCharacters(in: .whitespacesAndNewlines).count >= 3 {
|
value.trimmingCharacters(in: .whitespacesAndNewlines).count >= 3 {
|
||||||
let label = AX.copyString(el, kAXTitleAttribute)
|
let label = AX.copyString(el, kAXTitleAttribute)
|
||||||
?? AX.copyString(el, kAXDescriptionAttribute)
|
?? AX.copyString(el, kAXDescriptionAttribute)
|
||||||
|
|||||||
@ -12,14 +12,19 @@ final class TextReplacementCoordinator {
|
|||||||
private let ai: AIEngine
|
private let ai: AIEngine
|
||||||
private let gate: FeatureGate
|
private let gate: FeatureGate
|
||||||
private let detector: AppDetector
|
private let detector: AppDetector
|
||||||
|
/// Vocabulary the LLM must reproduce verbatim (personal dictionary / style
|
||||||
|
/// guide). Read fresh per action so dictionary edits take effect immediately.
|
||||||
|
private let protectedTerms: () -> [String]
|
||||||
|
|
||||||
init(capture: TextCapture, grammar: GrammarEngine, ai: AIEngine,
|
init(capture: TextCapture, grammar: GrammarEngine, ai: AIEngine,
|
||||||
gate: FeatureGate, detector: AppDetector) {
|
gate: FeatureGate, detector: AppDetector,
|
||||||
|
protectedTerms: @escaping () -> [String] = { [] }) {
|
||||||
self.capture = capture
|
self.capture = capture
|
||||||
self.grammar = grammar
|
self.grammar = grammar
|
||||||
self.ai = ai
|
self.ai = ai
|
||||||
self.gate = gate
|
self.gate = gate
|
||||||
self.detector = detector
|
self.detector = detector
|
||||||
|
self.protectedTerms = protectedTerms
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The "fix" hotkey — instant grammar correction (Layer 1, always free).
|
/// The "fix" hotkey — instant grammar correction (Layer 1, always free).
|
||||||
@ -52,7 +57,7 @@ final class TextReplacementCoordinator {
|
|||||||
do { cap = try await captureOffMain() } catch { return .failed(error.localizedDescription) }
|
do { cap = try await captureOffMain() } catch { return .failed(error.localizedDescription) }
|
||||||
|
|
||||||
do {
|
do {
|
||||||
let result = try await ai.run(action, on: cap.text)
|
let result = try await ai.run(action, on: cap.text, protectedTerms: protectedTerms())
|
||||||
gate.recordAIActionUsed()
|
gate.recordAIActionUsed()
|
||||||
let ok = await replaceOffMain(result, cap)
|
let ok = await replaceOffMain(result, cap)
|
||||||
return ok ? .replaced(result) : .failed("Couldn't write the result back.")
|
return ok ? .replaced(result) : .failed("Couldn't write the result back.")
|
||||||
@ -73,7 +78,7 @@ final class TextReplacementCoordinator {
|
|||||||
do { cap = try await captureOffMain() } catch { return .failed(error.localizedDescription) }
|
do { cap = try await captureOffMain() } catch { return .failed(error.localizedDescription) }
|
||||||
|
|
||||||
do {
|
do {
|
||||||
let result = try await ai.run(action, on: cap.text)
|
let result = try await ai.run(action, on: cap.text, protectedTerms: protectedTerms())
|
||||||
gate.recordAIActionUsed()
|
gate.recordAIActionUsed()
|
||||||
let ok = await replaceOffMain(result, cap)
|
let ok = await replaceOffMain(result, cap)
|
||||||
return ok ? .replaced(result) : .failed("Couldn't write the result back.")
|
return ok ? .replaced(result) : .failed("Couldn't write the result back.")
|
||||||
|
|||||||
@ -149,6 +149,66 @@ final class GrammarGemTests: XCTestCase {
|
|||||||
XCTAssertNil(ModeRegistry.mode(forBundleID: "com.unknown.app"))
|
XCTAssertNil(ModeRegistry.mode(forBundleID: "com.unknown.app"))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// MARK: - AIPrompts (backend-neutral seam)
|
||||||
|
|
||||||
|
func testProtectedTermsInjectedIntoSystemPrompt() {
|
||||||
|
let base = AIPrompts.systemPrompt(for: .rewriteClarity)
|
||||||
|
let withTerms = AIPrompts.systemPrompt(for: .rewriteClarity, protectedTerms: ["GrammarGem", "Harper"])
|
||||||
|
XCTAssertFalse(base.contains("GrammarGem"))
|
||||||
|
XCTAssertTrue(withTerms.contains("\"GrammarGem\""))
|
||||||
|
XCTAssertTrue(withTerms.contains("\"Harper\""))
|
||||||
|
XCTAssertTrue(withTerms.contains("verbatim"))
|
||||||
|
}
|
||||||
|
|
||||||
|
func testProtectedTermsClauseSkipsWhenEmpty() {
|
||||||
|
let blank = AIPrompts.systemPrompt(for: .rewrite, protectedTerms: ["", " "])
|
||||||
|
XCTAssertEqual(blank, AIPrompts.systemPrompt(for: .rewrite))
|
||||||
|
}
|
||||||
|
|
||||||
|
func testMaxTokensHonorsModeLengthCap() {
|
||||||
|
// Default actions get the full budget…
|
||||||
|
XCTAssertEqual(AIPrompts.maxTokens(for: .rewriteClarity), 512)
|
||||||
|
// …but a short-form Mode (Team Chat, lengthCap 60) gets far fewer.
|
||||||
|
let slack = ModeRegistry.slack
|
||||||
|
XCTAssertNotNil(slack.lengthCap)
|
||||||
|
let capped = AIPrompts.maxTokens(for: .applyMode(slack))
|
||||||
|
XCTAssertLessThan(capped, 512)
|
||||||
|
XCTAssertGreaterThanOrEqual(capped, 64)
|
||||||
|
// A mode with no cap (Polish) falls back to the default budget.
|
||||||
|
XCTAssertEqual(AIPrompts.maxTokens(for: .applyMode(ModeRegistry.polish)), 512)
|
||||||
|
}
|
||||||
|
|
||||||
|
func testCleanStripsPreambleAndQuotes() {
|
||||||
|
XCTAssertEqual(AIPrompts.clean("Sure, here's the rewrite: Hello world"), "Hello world")
|
||||||
|
XCTAssertEqual(AIPrompts.clean("\"Quoted result\""), "Quoted result")
|
||||||
|
XCTAssertEqual(AIPrompts.clean(" already clean "), "already clean")
|
||||||
|
}
|
||||||
|
|
||||||
|
func testParseToneIsRobust() {
|
||||||
|
XCTAssertEqual(AIPrompts.parseTone("professional"), .professional)
|
||||||
|
XCTAssertEqual(AIPrompts.parseTone(" Friendly.\n"), .friendly)
|
||||||
|
XCTAssertEqual(AIPrompts.parseTone("the tone is academic here"), .academic)
|
||||||
|
XCTAssertEqual(AIPrompts.parseTone("???"), .professional) // safe fallback
|
||||||
|
}
|
||||||
|
|
||||||
|
// MARK: - Custom modes
|
||||||
|
|
||||||
|
@MainActor
|
||||||
|
func testCustomModeRejectsEmptyPromptAndDuplicates() {
|
||||||
|
// Deterministic regardless of prior runs (the store persists to UserDefaults).
|
||||||
|
UserDefaults.standard.removeObject(forKey: "GrammarGem.customModes")
|
||||||
|
defer { UserDefaults.standard.removeObject(forKey: "GrammarGem.customModes") }
|
||||||
|
|
||||||
|
let store = CustomModesStore()
|
||||||
|
XCTAssertFalse(store.addCustom(name: "Brand Voice", prompt: " "),
|
||||||
|
"empty prompt must be rejected")
|
||||||
|
XCTAssertTrue(store.addCustom(name: "Brand Voice", prompt: "Write in Acme's warm voice."))
|
||||||
|
XCTAssertFalse(store.addCustom(name: "brand voice", prompt: "different prompt"),
|
||||||
|
"duplicate name (case-insensitive) must be rejected")
|
||||||
|
XCTAssertFalse(store.addCustom(name: "Polish", prompt: "collides with a built-in"),
|
||||||
|
"name colliding with a built-in mode must be rejected")
|
||||||
|
}
|
||||||
|
|
||||||
// MARK: - Device fingerprint
|
// MARK: - Device fingerprint
|
||||||
|
|
||||||
func testDeviceFingerprintIsStableAndHashed() {
|
func testDeviceFingerprintIsStableAndHashed() {
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user