feat: AI next-word suggestions (on-device first, default off)

app/src/main/java/helium314/keyboard/latin/ClipboardHistoryManager.kt

@@ -521,6 +521,11 @@ class ClipboardHistoryManager(
     fun getClipboardSuggestionView(editorInfo: EditorInfo?, parent: ViewGroup?): View? {
         clipboardSuggestionView = null
 
+        // Keyboard (and its icons set) may not be initialised yet (e.g. right after app start),
+        // and both the clipboard-content and screenshot views dereference it. Bail early to avoid
+        // an NPE instead of showing the suggestion.
+        if (latinIME.mKeyboardSwitcher.keyboard?.mIconsSet == null) return null
+
         // check for screenshot first if enabled
         if (latinIME.mSettings.current.mSuggestScreenshots) {
             val screenshotView = getScreenshotSuggestionView(parent)
@@ -583,6 +588,9 @@ class ClipboardHistoryManager(
 
     private fun getScreenshotSuggestionView(parent: ViewGroup?): View? {
         if (parent == null || dontShowCurrentSuggestion) return null
+        // Keyboard may not be initialised yet (e.g. right after app start); the view's close/paste
+        // icons come from its icons set, so bail out early rather than NPE.
+        if (latinIME.mKeyboardSwitcher.keyboard?.mIconsSet == null) return null
 
         val permission = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU) {
             android.Manifest.permission.READ_MEDIA_IMAGES

app/src/main/java/helium314/keyboard/latin/DictionaryFacilitator.java

@@ -170,4 +170,30 @@ default Map<String, Integer> getAllMainDictionaryWordsWithFrequency() {
 
     default void forEachMainDictionaryWord(java.util.function.BiConsumer<String, Integer> consumer) {
     }
+
+    /**
+     * Registers a callback that should be run when the AI next-word dictionary has produced new
+     * LLM candidates in the background. The implementer should trigger a suggestion-strip refresh
+     * (e.g. LatinIME handler {@code postUpdateSuggestionStrip}) so the freshly cached candidates
+     * surface without an extra keystroke. Default no-op; only next-word-capable locales override.
+     */
+    default void setNextWordRefreshListener(@Nullable Runnable listener) {
+    }
+
+    /**
+     * Supplies the latest bounded text before the cursor, captured by the IME just before a
+     * next-word suggestion fetch. Used by the AI next-word dictionary to build a richer prompt
+     * (the actual sentence/message context) than the last few n-gram words alone. Default no-op.
+     */
+    default void setAINextWordContextText(@Nullable String text) {
+    }
+
+    /**
+     * Supplies per-app metadata for the AI next-word dictionary: the current app package name
+     * (context buffer key), the app's preferred keyboard language code (e.g. "en", "pl"), whether
+     * the current field is private/incognito (noLearning), and the persisted per-app context text.
+     */
+    default void setAINextWordAppInfo(@Nullable String packageName, @Nullable String language,
+                                      boolean noLearning, @Nullable String persistedContext) {
+    }
 }

app/src/main/java/helium314/keyboard/latin/DictionaryFacilitatorImpl.kt

@@ -21,6 +21,7 @@ import helium314.keyboard.latin.common.StringUtils
 import helium314.keyboard.latin.common.decapitalize
 import helium314.keyboard.latin.common.mightBeEmoji
 import helium314.keyboard.latin.common.splitOnWhitespace
+import helium314.keyboard.latin.dictionary.AINextWordDictionary
 import helium314.keyboard.latin.dictionary.AppsBinaryDictionary
 import helium314.keyboard.latin.dictionary.ContactsBinaryDictionary
 import helium314.keyboard.latin.dictionary.Dictionary
@@ -33,6 +34,7 @@ import helium314.keyboard.latin.personalization.SessionWordBoost
 import helium314.keyboard.latin.personalization.UserHistoryDictionary
 import helium314.keyboard.latin.settings.Settings
 import helium314.keyboard.latin.settings.SettingsValuesForSuggestion
+import helium314.keyboard.latin.utils.AINextWordEngineFactory
 import helium314.keyboard.latin.utils.Log
 import helium314.keyboard.latin.utils.SubtypeSettings
 import helium314.keyboard.latin.utils.SuggestionResults
@@ -66,6 +68,7 @@ class DictionaryFacilitatorImpl : DictionaryFacilitator {
     private var mContext: Context? = null
     private var mEnabledDictionariesState: Map<String, Boolean> = emptyMap()
     private var mLoadedDownloadPrefs: Map<String, Any?> = emptyMap()
+    private var mLoadedAINextWord: Boolean? = null
     private var dictionaryGroups = listOf(DictionaryGroup())
 
     @Volatile
@@ -99,6 +102,14 @@ class DictionaryFacilitatorImpl : DictionaryFacilitator {
     // Limit parallelism to prevent excessive CPU usage during dictionary operations
     private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default.limitedParallelism(2))
 
+    // Optional AI next-word dictionary (LeanType). Null when the feature is off or the engine
+    // cannot be built (factory gates on the pref / model availability).
+    private var aiNextWordDict: AINextWordDictionary? = null
+
+    // Runnable that LatinIME registers so we can ask it to redraw the suggestion strip once the
+    // AI next-word dictionary has cached fresh candidates (see setNextWordRefreshListener).
+    private var nextWordRefreshListener: (() -> Unit)? = null
+
     override fun setValidSpellingWordReadCache(cache: LruCache<String, Boolean>) {
         mValidSpellingWordReadCache = cache
     }
@@ -151,6 +162,10 @@ class DictionaryFacilitatorImpl : DictionaryFacilitator {
             if (currentPrefs != mEnabledDictionariesState || currentDownloadPrefs != mLoadedDownloadPrefs) {
                 return false
             }
+            val aiNextWord = prefs.getBoolean(Settings.PREF_AI_NEXT_WORD, false)
+            if (aiNextWord != mLoadedAINextWord) {
+                return false
+            }
         }
         val ctx = mContext ?: return false
         val currentSuggestEmojis = Settings.getValues().mSuggestEmojis
@@ -186,6 +201,9 @@ class DictionaryFacilitatorImpl : DictionaryFacilitator {
             .mapValues { it.value as? Boolean ?: true }
         mLoadedDownloadPrefs = prefs.all.filterKeys { it.startsWith("pref_dict_download_link_") }
 
+        // Track the AI next-word pref so the dictionary group is rebuilt when it is toggled.
+        mLoadedAINextWord = prefs.getBoolean(Settings.PREF_AI_NEXT_WORD, false)
+
         // Initialize session word boost with context if not yet done
         if (sessionWordBoost == null) {
             sessionWordBoost = SessionWordBoost.getInstance(context)
@@ -228,6 +246,42 @@ class DictionaryFacilitatorImpl : DictionaryFacilitator {
 
         mValidSpellingWordWriteCache?.evictAll()
         mValidSpellingWordReadCache?.evictAll()
+
+        // AI next-word dictionary: rebuilt on every dictionary reset. Factory returns null when
+        // the feature is disabled or the engine cannot be built, making this a no-op by default.
+        aiNextWordDict = AINextWordEngineFactory.create(context)?.let {
+            AINextWordDictionary(it, scope).also { dict ->
+                // Fresh LLM candidates cached -> ask LatinIME to redraw the suggestion strip so
+                // they surface without another keystroke.
+                dict.onCandidatesReady = { nextWordRefreshListener?.invoke() }
+            }
+        }
+        Log.i(TAG, "resetDictionaries: AINextWord dict ${if (aiNextWordDict == null) "NOT created" else "created"} (pref=${prefs.getBoolean(Settings.PREF_AI_NEXT_WORD, false)})")
+    }
+
+    override fun setNextWordRefreshListener(listener: Runnable?) {
+        nextWordRefreshListener = listener?.let { { it.run() } }
+    }
+
+    @Volatile
+    private var aiNextWordContextText: String? = null
+
+    override fun setAINextWordContextText(text: String?) {
+        aiNextWordContextText = text
+    }
+
+    @Volatile private var aiNextWordAppPackage: String? = null
+    @Volatile private var aiNextWordAppLanguage: String? = null
+    @Volatile private var aiNextWordNoLearning: Boolean = false
+    @Volatile private var aiNextWordPersistedContext: String? = null
+
+    override fun setAINextWordAppInfo(
+        packageName: String?, language: String?, noLearning: Boolean, persistedContext: String?
+    ) {
+        aiNextWordAppPackage = packageName
+        aiNextWordAppLanguage = language
+        aiNextWordNoLearning = noLearning
+        aiNextWordPersistedContext = persistedContext
     }
 
     /** creates dictionaryGroups for [newLocales] with given [newSubDictTypes], trying to re-use existing dictionaries.
@@ -714,6 +768,42 @@ class DictionaryFacilitatorImpl : DictionaryFacilitator {
                 }
             }
         }
+        // AI next-word candidates, only in prediction/next-word mode. The dictionary returns
+        // null (or filtered-away empties) unless it has cached LLM candidates, so it never
+        // blocks this pass. AOSP suggestions always come first (added above).
+        if (composedData.mTypedWord.isEmpty()) {
+            aiNextWordDict?.let { dict ->
+                dict.prevTextForPrompt = aiNextWordContextText
+                dict.appPackageName = aiNextWordAppPackage
+                dict.appLanguageHint = aiNextWordAppLanguage
+                dict.noLearning = aiNextWordNoLearning
+                dict.appContext = if (aiNextWordNoLearning) null else aiNextWordPersistedContext
+                val aiInfos = dict.getSuggestions(
+                    composedData, ngramContext, proximityInfoHandle, settingsValuesForSuggestion,
+                    sessionId, weightForLocale, weightOfLangModelVsSpatialModel
+                )?.filter { it.word.isNotEmpty() }
+                if (!aiInfos.isNullOrEmpty()) {
+                    // Rank AI words BELOW every AOSP suggestion in this list. AOSP next-word scores
+                    // vary wildly (history ~700-800, user dict can be millions), so a fixed base is
+                    // unreliable — instead key off the actual AOSP scores present so AI words ALWAYS
+                    // append as extra choices and can never displace or reorder the AOSP suggestions
+                    // the user is about to tap. mScore is final, so rebuild the infos with lower
+                    // (strictly decreasing, hence unique within the TreeSet) scores.
+                    val floor = ((suggestions.minOfOrNull { it.mScore } ?: AI_NEXT_WORD_BASE_SCORE) - 1)
+                    aiInfos.forEachIndexed { i, info ->
+                        val score = floor - i
+                        suggestions.add(
+                            if (info.mScore == score) info
+                            else SuggestedWordInfo(
+                                info.mWord, info.mPrevWordsContext, score, info.mKindAndFlags,
+                                info.mSourceDict, info.mIndexOfTouchPointOfSecondWord,
+                                info.mAutoCommitFirstWordConfidence
+                            )
+                        )
+                    }
+                }
+            }
+        }
         return suggestions
     }
 
@@ -812,6 +902,8 @@ class DictionaryFacilitatorImpl : DictionaryFacilitator {
 
     companion object {
         private val TAG = DictionaryFacilitatorImpl::class.java.simpleName
+        // Fallback base when there are no AOSP suggestions to rank AI words below.
+        private const val AI_NEXT_WORD_BASE_SCORE = 1000
 
         // HACK: This threshold is being used when adding a capitalized entry in the User History dictionary.
         private const val CAPITALIZED_FORM_MAX_PROBABILITY_FOR_INSERT = 140

app/src/main/java/helium314/keyboard/latin/LatinIME.java

@@ -586,6 +586,17 @@ public void onCreate() {
         mClipboardHistoryManager.onCreate();
         mHandler.onCreate();
 
+        // When the AI next-word dictionary caches fresh candidates, ask the suggestion strip to
+        // redraw so they surface without another keystroke (see AINextWordDictionary.onCandidatesReady).
+        // Also evict the next-word suggestions cache first: it may hold the AOSP-only list built
+        // before the async AI response landed, so getNextWordSuggestions() would otherwise return
+        // that stale list and the fresh AI words would never reach the strip.
+        mDictionaryFacilitator.setNextWordRefreshListener(
+                () -> {
+                    mInputLogic.getSuggest().clearNextWordSuggestionsCache();
+                    mHandler.postUpdateSuggestionStrip(SuggestedWords.INPUT_STYLE_TYPING);
+                });
+
         // Register to receive ringer mode change.
         final IntentFilter filter = new IntentFilter();
         filter.addAction(AudioManager.RINGER_MODE_CHANGED_ACTION);
@@ -1206,6 +1217,23 @@ void onFinishInputInternal() {
         super.onFinishInput();
         Log.i(TAG, "onFinishInput");
 
+        // Persist per-app context for the AI next-word feature. NEVER records in private/incognito
+        // or anonymous fields (noLearning / no-suggestions), covering password and incognito input.
+        try {
+            final InputAttributes attrs = mSettings.getCurrent().mInputAttributes;
+            if (attrs.mShouldShowSuggestions && !attrs.mNoLearning) {
+                final String pkg = attrs.mTargetApplicationPackageName;
+                if (pkg != null) {
+                    final android.view.inputmethod.InputConnection ic = getCurrentInputConnection();
+                    if (ic != null) {
+                        final CharSequence text = ic.getTextBeforeCursor(200, 0);
+                        mSettings.appendAppContext(pkg, text);
+                    }
+                }
+            }
+        } catch (Exception ignored) {
+        }
+
         mDictionaryFacilitator.onFinishInput();
         final MainKeyboardView mainKeyboardView = mKeyboardSwitcher.getMainKeyboardView();
         if (mainKeyboardView != null) {