diff --git a/Makefile b/Makefile index e512fa5..133e9e5 100644 --- a/Makefile +++ b/Makefile @@ -104,14 +104,14 @@ endif test: $(TEST_RUNNER) @$(TEST_RUNNER) -$(TEST_RUNNER): Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift Sources/ShortcutCore/ShortcutModels.swift Tests/ShortcutCancelBindingTests.swift Sources/ShortcutCore/MouseDictationButton.swift Tests/MouseDictationButtonTests.swift Tests/SmartCleanupValidationTests.swift +$(TEST_RUNNER): Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift Sources/ShortcutCore/ShortcutModels.swift Tests/ShortcutCancelBindingTests.swift Sources/ShortcutCore/MouseDictationButton.swift Tests/MouseDictationButtonTests.swift Tests/SmartCleanupValidationTests.swift Tests/StructuredOutputUnwrapTests.swift @mkdir -p "$(BUILD_DIR)" swiftc \ -parse-as-library \ -o "$(TEST_RUNNER)" \ -sdk $(shell xcrun --show-sdk-path) \ -target $(ARCH)-apple-macosx26.0 \ - Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift Sources/ShortcutCore/ShortcutModels.swift Tests/ShortcutCancelBindingTests.swift Sources/ShortcutCore/MouseDictationButton.swift Tests/MouseDictationButtonTests.swift Tests/SmartCleanupValidationTests.swift + Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift Sources/ShortcutCore/ShortcutModels.swift Tests/ShortcutCancelBindingTests.swift Sources/ShortcutCore/MouseDictationButton.swift Tests/MouseDictationButtonTests.swift Tests/SmartCleanupValidationTests.swift Tests/StructuredOutputUnwrapTests.swift icon: $(ICON_ICNS) diff --git a/Sources/AppleFoundationModelsPostProcessor.swift b/Sources/AppleFoundationModelsPostProcessor.swift index 375898a..a1d5006 100644 --- a/Sources/AppleFoundationModelsPostProcessor.swift +++ b/Sources/AppleFoundationModelsPostProcessor.swift @@ -595,7 +595,8 @@ actor AppleFoundationModelsPostProcessor { /// (e.g. "wrap this in a div") is never stripped. private static let wrapperTags = [ "response", "result", "output", "answer", "reply", "message", - "bulleted_list", "numbered_list", "list", "rewritten_text" + "bulleted_list", "numbered_list", "list", "rewritten_text", + "cleaned_text", "clean_text" ] /// Prompt sections the model sometimes replays before its actual answer. private static let echoedPromptTags = [ @@ -607,6 +608,12 @@ actor AppleFoundationModelsPostProcessor { var value = raw.trimmingCharacters(in: .whitespacesAndNewlines) let options: String.CompareOptions = [.regularExpression, .caseInsensitive] + // The small model sometimes returns its answer wrapped in a markdown + // code fence and/or a JSON object like {"cleaned_text": "…"} despite the + // instructions. Peel those structured wrappers before the tag loop so + // the text they contain — not the scaffolding — reaches the user. + value = Self.unwrapStructuredOutput(value) + while !value.isEmpty { let before = value for tag in Self.echoedPromptTags { @@ -640,6 +647,68 @@ actor AppleFoundationModelsPostProcessor { return value } + /// JSON keys the model uses when it wraps a plain answer in an object, e.g. + /// `{"cleaned_text": "…"}`. Ordered by preference for extraction. + private static let jsonTextKeys = [ + "cleaned_text", "clean_text", "cleaned", "corrected_text", + "rewritten_text", "text", "output", "result", "response", "answer" + ] + + /// Peels a single-purpose JSON object the model emits around its answer, so + /// `{"cleaned_text": "Hi."}` — bare or wrapped in a ```` ```json ```` fence — + /// collapses to `Hi.`. A code fence is removed *only* when it wraps such an + /// object, so a code block the user genuinely dictated survives untouched. + private static func unwrapStructuredOutput(_ input: String) -> String { + let value = input.trimmingCharacters(in: .whitespacesAndNewlines) + if let unwrapped = Self.jsonTextValue(in: value) { + return unwrapped.trimmingCharacters(in: .whitespacesAndNewlines) + } + let defenced = Self.stripCodeFence(value) + if defenced != value, let unwrapped = Self.jsonTextValue(in: defenced) { + return unwrapped.trimmingCharacters(in: .whitespacesAndNewlines) + } + return value + } + + /// Removes a surrounding ```` ``` ```` fence (with optional language tag) + /// only when the whole string is one fenced block, so inline prose that + /// merely mentions backticks is untouched. + private static func stripCodeFence(_ text: String) -> String { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard trimmed.hasPrefix("```") else { return trimmed } + var lines = trimmed.components(separatedBy: .newlines) + guard lines.count >= 2, + let opener = lines.first?.trimmingCharacters(in: .whitespaces), + opener.range(of: "^```[a-zA-Z0-9+#-]*$", options: .regularExpression) != nil, + lines.last?.trimmingCharacters(in: .whitespaces) == "```" else { + return trimmed + } + lines.removeFirst() + lines.removeLast() + return lines.joined(separator: "\n").trimmingCharacters(in: .whitespacesAndNewlines) + } + + /// Returns the wrapped text when `input` is a JSON object whose keys are all + /// known answer-wrapper keys. Requiring *every* key to be known leaves a + /// genuinely dictated object such as `{"name": "Ada"}` untouched. + private static func jsonTextValue(in input: String) -> String? { + let trimmed = input.trimmingCharacters(in: .whitespacesAndNewlines) + guard trimmed.hasPrefix("{"), trimmed.hasSuffix("}"), + let data = trimmed.data(using: .utf8), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + !object.isEmpty else { + return nil + } + var lowered: [String: Any] = [:] + for (key, value) in object { lowered[key.lowercased()] = value } + let known = Set(Self.jsonTextKeys) + guard lowered.keys.allSatisfy({ known.contains($0) }) else { return nil } + for key in Self.jsonTextKeys { + if let value = lowered[key] as? String { return value } + } + return nil + } + private func makeSession(instructions: String) -> LanguageModelSession { LanguageModelSession(model: model, tools: [], instructions: instructions) } diff --git a/Tests/AppContextServiceTests.swift b/Tests/AppContextServiceTests.swift index dccebef..74e56c6 100644 --- a/Tests/AppContextServiceTests.swift +++ b/Tests/AppContextServiceTests.swift @@ -29,6 +29,7 @@ struct AppContextServiceTests { ShortcutCancelBindingTests.run() MouseDictationButtonTests.run() SmartCleanupValidationTests.run() + StructuredOutputUnwrapTests.run() print("MegaphoneTests passed") } diff --git a/Tests/StructuredOutputUnwrapTests.swift b/Tests/StructuredOutputUnwrapTests.swift new file mode 100644 index 0000000..3395a99 --- /dev/null +++ b/Tests/StructuredOutputUnwrapTests.swift @@ -0,0 +1,66 @@ +import Foundation + +/// Covers issue #14: the on-device model occasionally wraps its answer in a JSON +/// object or `` tags, which `normalizeCommandOutput` must peel so +/// the user gets their text and not the scaffolding. +enum StructuredOutputUnwrapTests { + static func run() { + testBareJSONObjectIsUnwrapped() + testFencedJSONObjectIsUnwrapped() + testCleanedTextTagsAreStripped() + testGenuineDataObjectIsUntouched() + testDictatedCodeBlockSurvives() + testOrdinaryProseIsUnchanged() + } + + private static func testBareJSONObjectIsUnwrapped() { + expect("{\n \"cleaned_text\": \"Testing, testing.\"\n}", becomes: "Testing, testing.") + expect("{\"text\": \"- YOLO.\"}", becomes: "- YOLO.") + expect("{\"clean_text\":\"What's up?\"}", becomes: "What's up?") + } + + private static func testFencedJSONObjectIsUnwrapped() { + expect("```json\n{\n \"cleaned_text\": \"Testing, testing.\"\n}\n```", becomes: "Testing, testing.") + expect("```\n{\"text\": \"- YOLO.\"}\n```", becomes: "- YOLO.") + } + + private static func testCleanedTextTagsAreStripped() { + expect("\nWhat's up, my fellow humans?\n", becomes: "What's up, my fellow humans?") + expect("How do you do, fellow kids?", becomes: "How do you do, fellow kids?") + } + + /// An object whose keys are not all answer-wrappers is real dictated content, + /// not scaffolding, so it must pass through verbatim. + private static func testGenuineDataObjectIsUntouched() { + let object = "{\"name\": \"Ada\", \"text\": \"hello\"}" + expect(object, becomes: object) + } + + /// De-fencing is scoped to JSON wrappers, so a code block the user actually + /// dictated (e.g. via a wake command) keeps its fence. + private static func testDictatedCodeBlockSurvives() { + let block = "```swift\nprint(\"hi\")\n```" + expect(block, becomes: block) + } + + private static func testOrdinaryProseIsUnchanged() { + expect("Let's ship the release on Wednesday.", becomes: "Let's ship the release on Wednesday.") + // A stray brace in prose is not a JSON object and must not be touched. + expect("Use the {placeholder} token here.", becomes: "Use the {placeholder} token here.") + } + + private static func expect( + _ input: String, + becomes expected: String, + file: StaticString = #file, + line: UInt = #line + ) { + let actual = AppleFoundationModelsPostProcessor.normalizeCommandOutput(input) + if actual != expected { + fatalError( + "\(file):\(line): normalizeCommandOutput(\(input.debugDescription)) == " + + "\(actual.debugDescription), expected \(expected.debugDescription)" + ) + } + } +}