From 57b58ec80e25741ee441e4afbb0a8a70be0cfd9b Mon Sep 17 00:00:00 2001
From: Beingpax <prakashjoshipaxtek@gmail.com>
Date: Tue, 19 Aug 2025 21:48:22 +0545
Subject: [PATCH] Add dictionary context to AI enhancement

---
 VoiceInk/Services/AIEnhancementService.swift  |  8 +++--
 .../Services/DictionaryContextService.swift   | 33 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 VoiceInk/Services/DictionaryContextService.swift

diff --git a/VoiceInk/Services/AIEnhancementService.swift b/VoiceInk/Services/AIEnhancementService.swift
index 56bc161..96d5a89 100644
--- a/VoiceInk/Services/AIEnhancementService.swift
+++ b/VoiceInk/Services/AIEnhancementService.swift
@@ -61,6 +61,7 @@ class AIEnhancementService: ObservableObject {
     
     private let aiService: AIService
     private let screenCaptureService: ScreenCaptureService
+    private let dictionaryContextService: DictionaryContextService
     private let baseTimeout: TimeInterval = 30
     private let rateLimitInterval: TimeInterval = 1.0
     private var lastRequestTime: Date?
@@ -70,6 +71,7 @@ class AIEnhancementService: ObservableObject {
         self.aiService = aiService
         self.modelContext = modelContext
         self.screenCaptureService = ScreenCaptureService()
+        self.dictionaryContextService = DictionaryContextService.shared
         
         self.isEnhancementEnabled = UserDefaults.standard.bool(forKey: "isAIEnhancementEnabled")
         self.useClipboardContext = UserDefaults.standard.bool(forKey: "useClipboardContext")
@@ -154,8 +156,10 @@ class AIEnhancementService: ObservableObject {
             ""
         }
         
-        let contextSection = if !clipboardContext.isEmpty || !screenCaptureContext.isEmpty {
-            "\n\n<CONTEXT_INFORMATION>\(clipboardContext)\(screenCaptureContext)\n</CONTEXT_INFORMATION>"
+        let dictionaryContext = dictionaryContextService.getDictionaryContext()
+        
+        let contextSection = if !clipboardContext.isEmpty || !screenCaptureContext.isEmpty || !dictionaryContext.isEmpty {
+            "\n\n<CONTEXT_INFORMATION>\(clipboardContext)\(screenCaptureContext)\(dictionaryContext)\n</CONTEXT_INFORMATION>"
         } else {
             ""
         }
diff --git a/VoiceInk/Services/DictionaryContextService.swift b/VoiceInk/Services/DictionaryContextService.swift
new file mode 100644
index 0000000..1025ff5
--- /dev/null
+++ b/VoiceInk/Services/DictionaryContextService.swift
@@ -0,0 +1,33 @@
+import Foundation
+import SwiftUI
+
+class DictionaryContextService {
+    static let shared = DictionaryContextService()
+    
+    private init() {}
+    
+    /// Gets dictionary context information to be included in AI enhancement
+    func getDictionaryContext() -> String {
+        guard let dictionaryWords = getDictionaryWords(), !dictionaryWords.isEmpty else {
+            return ""
+        }
+        
+        let wordsText = dictionaryWords.joined(separator: ", ")
+        return "\n\nImportant Vocabulary: \(wordsText)"
+    }
+    
+    /// Gets enabled custom dictionary words from UserDefaults
+    private func getDictionaryWords() -> [String]? {
+        guard let data = UserDefaults.standard.data(forKey: "CustomDictionaryItems") else {
+            return nil
+        }
+        
+        do {
+            let items = try JSONDecoder().decode([DictionaryItem].self, from: data)
+            let enabledWords = items.filter { $0.isEnabled }.map { $0.word }
+            return enabledWords.isEmpty ? nil : enabledWords
+        } catch {
+            return nil
+        }
+    }
+}