@atlaskit/editor-plugin-autocomplete 0.1.0 → 0.3.0

package/dist/es2019/autocompletePlugin.js

@@ -1,6 +1,20 @@
-// This file will contain the autocomplete plugin implementation.
-// The full implementation will be added in a follow-up PR.
-
-export const autocompletePlugin = () => ({
-  name: 'autocomplete'
-});
+import { autocompletePluginKey, createAutocompletePlugin } from './pm-plugins/autocomplete-plugin';
+export const autocompletePlugin = ({
+  config: options
+}) => {
+  return {
+    name: 'autocomplete',
+    getSharedState(editorState) {
+      if (!editorState) {
+        return undefined;
+      }
+      return autocompletePluginKey.getState(editorState);
+    },
+    pmPlugins() {
+      return [{
+        name: 'autocomplete',
+        plugin: () => createAutocompletePlugin(options)
+      }];
+    }
+  };
+};

package/dist/es2019/autocompletePluginType.js

@@ -0,0 +1 @@
+export {};

package/dist/es2019/pm-plugins/autocomplete-plugin.js

@@ -0,0 +1,368 @@
+// url: prefix is an Atlaspack/Parcel directive that resolves this file as an emitted
+// asset URL (content-hashed). For rspack, this is handled via staticAssetsLoader.
+// eslint-disable-next-line @repo/internal/import/no-unresolved
+import wordVectorsUrl from 'url:./data/word-vectors_10k.bin';
+import { SafePlugin } from '@atlaskit/editor-common/safe-plugin';
+import { keydownHandler } from '@atlaskit/editor-prosemirror/keymap';
+import { PluginKey } from '@atlaskit/editor-prosemirror/state';
+import { DecorationSet } from '@atlaskit/editor-prosemirror/view';
+import { createGhostTextDecorationSet } from './ghost-text-decoration';
+import { createSlowLaneClient, setDefaultSlowLaneClient, isWordBoundary } from './slow-lane-client';
+import { predict, loadDefaultVocabulary, loadVectorsAsync, incrementSessionFreq, ingestDocumentPage } from './text-predictor';
+const SLOW_LANE_ENDPOINT = '/gateway/api/v1/autocomplete/typeahead-encodings';
+export const autocompletePluginKey = new PluginKey('autocomplete');
+const DEBOUNCE_MS = 150;
+const createInitialState = () => ({
+  ghostText: '',
+  ghostPosition: -1,
+  decorationSet: DecorationSet.empty
+});
+
+/**
+ * Extract text content before the cursor from the current document.
+ * Returns the last ~200 characters for context.
+ */
+const getTextBeforeCursor = state => {
+  const {
+    $from
+  } = state.selection;
+  const maxChars = 200;
+
+  // 1. Get the perfectly flattened text of the current block up to the cursor
+  const blockNode = $from.parent;
+  const offsetInBlock = $from.parentOffset;
+  const blockText = blockNode.textContent.slice(0, offsetInBlock);
+  if (blockText.length >= maxChars) {
+    return blockText.slice(-maxChars);
+  }
+  let fullText = blockText;
+
+  // 2. Walk backwards through previous blocks
+  let depth = $from.depth - 1;
+  while (fullText.length < maxChars && depth >= 0) {
+    const parentNode = $from.node(depth);
+    const indexInParent = $from.index(depth);
+    for (let i = indexInParent - 1; i >= 0 && fullText.length < maxChars; i--) {
+      const sibling = parentNode.child(i);
+      const siblingText = sibling.textContent;
+      fullText = siblingText + '\n' + fullText;
+    }
+    depth--;
+  }
+  return fullText.slice(-maxChars);
+};
+
+/**
+ * Set the autocomplete state via a transaction metadata.
+ */
+const setAutocompleteMeta = (tr, meta) => {
+  return tr.setMeta(autocompletePluginKey, meta);
+};
+
+/**
+ * Apply a ghost text suggestion to the editor state.
+ */
+const showGhostText = (view, text, position) => {
+  const {
+    state,
+    dispatch
+  } = view;
+  const decorationSet = createGhostTextDecorationSet(state, position, text);
+  const tr = setAutocompleteMeta(state.tr, {
+    ghostText: text,
+    ghostPosition: position,
+    decorationSet
+  });
+  dispatch(tr);
+};
+
+/**
+ * Clear the current ghost text from the editor.
+ */
+const clearGhostText = (state, dispatch) => {
+  const pluginState = autocompletePluginKey.getState(state);
+  if (!pluginState || !pluginState.ghostText) {
+    return false;
+  }
+  if (dispatch) {
+    const tr = setAutocompleteMeta(state.tr, {
+      ghostText: '',
+      ghostPosition: -1,
+      decorationSet: DecorationSet.empty
+    });
+    dispatch(tr);
+  }
+  return true;
+};
+
+/**
+ * Accept the current ghost text suggestion and insert it into the document.
+ */
+const acceptGhostText = (state, dispatch) => {
+  const pluginState = autocompletePluginKey.getState(state);
+  if (!pluginState || !pluginState.ghostText) {
+    return false;
+  }
+  if (dispatch) {
+    const {
+      ghostText,
+      ghostPosition
+    } = pluginState;
+    let tr = state.tr.insertText(ghostText, ghostPosition);
+    tr = setAutocompleteMeta(tr, {
+      ghostText: '',
+      ghostPosition: -1,
+      decorationSet: DecorationSet.empty
+    });
+    dispatch(tr);
+  }
+  return true;
+};
+
+/**
+ * Context provided to the autocomplete plugin on first editor focus.
+ * Text fields are selectively ingested to boost word-frequency scoring for
+ * predictions, giving words already present in the document/thread an L1
+ * priority boost.
+ */
+
+/**
+ * Build the text payload for the slow-lane request by prepending any available
+ * comment context ahead of the live document text. This gives the backend
+ * model richer context about the thread the user is writing in.
+ */
+const buildSlowLaneText = (docText, context) => {
+  var _context$siblingComme, _context$siblingComme2, _context$siblingComme3;
+  const lines = [];
+  if (context !== null && context !== void 0 && context.parentCommentContent) {
+    lines.push(`comment: ${context.parentCommentContent}`);
+  }
+  context === null || context === void 0 ? void 0 : (_context$siblingComme = context.siblingCommentsContents) === null || _context$siblingComme === void 0 ? void 0 : _context$siblingComme.forEach((sibling, index) => {
+    lines.push(`reply ${index + 1}: ${sibling}`);
+  });
+  const nextReplyNumber = ((_context$siblingComme2 = context === null || context === void 0 ? void 0 : (_context$siblingComme3 = context.siblingCommentsContents) === null || _context$siblingComme3 === void 0 ? void 0 : _context$siblingComme3.length) !== null && _context$siblingComme2 !== void 0 ? _context$siblingComme2 : 0) + 1;
+  lines.push(`reply ${nextReplyNumber}: ${docText}`);
+  return lines.join('\n');
+};
+export const createAutocompletePlugin = options => {
+  let debounceTimer = null;
+  let hasIngestedPage = false;
+  let resolvedContext;
+  /**
+   * Set after accepting a suggestion so the next doc-change update
+   * skips scheduling a new prediction for the just-inserted text.
+   * Scoped to the factory so multiple editor instances don't share state.
+   */
+  let justAccepted = false;
+
+  /**
+   * Stores the text-before-cursor snapshot at the moment the user dismissed
+   * a suggestion via Escape. While the context remains identical, we suppress
+   * re-showing the same suggestion. Resets to null as soon as the text changes.
+   */
+  let dismissedContext = null;
+  const slowLaneClient = createSlowLaneClient({
+    baseUrl: '',
+    endpoint: SLOW_LANE_ENDPOINT
+  });
+  setDefaultSlowLaneClient(slowLaneClient);
+
+  /**
+   * Schedule a prediction after a short debounce.
+   * Tier 1 predictions are synchronous (<0.1ms) but we still debounce
+   * to avoid unnecessary work on rapid keystrokes.
+   */
+  const schedulePrediction = view => {
+    if (debounceTimer) {
+      clearTimeout(debounceTimer);
+    }
+    debounceTimer = setTimeout(() => {
+      const {
+        state
+      } = view;
+      const {
+        selection
+      } = state;
+
+      // Only predict for cursor selections (not range selections)
+      if (!selection.empty) {
+        return;
+      }
+      const textBefore = getTextBeforeCursor(state);
+
+      // Suppress re-showing the same suggestion the user just dismissed.
+      // Once the text context changes (user types or deletes), this clears automatically.
+      if (textBefore === dismissedContext) {
+        return;
+      }
+      dismissedContext = null;
+
+      // Don't predict if there's not enough context
+      if (textBefore.trim().length < 3) {
+        return;
+      }
+
+      // Tier 1 prediction is synchronous -- no async needed
+      const prediction = predict(textBefore);
+      if (prediction && prediction.length > 0) {
+        showGhostText(view, prediction, selection.from);
+      }
+    }, DEBOUNCE_MS);
+  };
+  const maybeUpdateSessionFrequency = (view, prevState) => {
+    const newText = getTextBeforeCursor(view.state);
+    const prevText = getTextBeforeCursor(prevState);
+    if (newText.length <= prevText.length) {
+      return;
+    }
+    const lastChar = newText[newText.length - 1];
+    // eslint-disable-next-line require-unicode-regexp
+    if (!/[\s.,;:!?]/.test(lastChar)) {
+      return;
+    }
+
+    // Only fire if the previous state did not already end on a boundary,
+    // so we don't double-count when multiple boundary chars are inserted.
+    const prevLastChar = prevText[prevText.length - 1];
+    // eslint-disable-next-line require-unicode-regexp
+    if (prevLastChar && /[\s.,;:!?]/.test(prevLastChar)) {
+      return;
+    }
+    const beforeBoundary = newText.slice(0, -1).trimEnd();
+    const lastSpaceIdx = beforeBoundary.lastIndexOf(' ');
+    const completedWord = beforeBoundary.slice(lastSpaceIdx + 1).toLowerCase();
+    if (completedWord.length >= 2) {
+      incrementSessionFreq(completedWord);
+    }
+  };
+  return new SafePlugin({
+    key: autocompletePluginKey,
+    state: {
+      init: () => createInitialState(),
+      apply: (tr, pluginState) => {
+        const meta = tr.getMeta(autocompletePluginKey);
+        if (meta) {
+          return {
+            ...pluginState,
+            ...meta
+          };
+        }
+
+        // If the document changed, clear the ghost text
+        // (new prediction will be scheduled from view.update)
+        if (tr.docChanged) {
+          return {
+            ...pluginState,
+            ghostText: '',
+            ghostPosition: -1,
+            decorationSet: DecorationSet.empty
+          };
+        }
+
+        // If selection changed without doc change, clear ghost text
+        if (tr.selectionSet && pluginState.ghostText) {
+          return {
+            ...pluginState,
+            ghostText: '',
+            ghostPosition: -1,
+            decorationSet: DecorationSet.empty
+          };
+        }
+        return pluginState;
+      }
+    },
+    props: {
+      decorations: state => {
+        var _pluginState$decorati;
+        const pluginState = autocompletePluginKey.getState(state);
+        return (_pluginState$decorati = pluginState === null || pluginState === void 0 ? void 0 : pluginState.decorationSet) !== null && _pluginState$decorati !== void 0 ? _pluginState$decorati : DecorationSet.empty;
+      },
+      handleKeyDown: keydownHandler({
+        Tab: (state, dispatch) => {
+          const accepted = acceptGhostText(state, dispatch);
+          if (accepted) justAccepted = true;
+          return accepted;
+        },
+        ArrowRight: (state, dispatch) => {
+          const accepted = acceptGhostText(state, dispatch);
+          if (accepted) justAccepted = true;
+          return accepted;
+        },
+        Escape: (state, dispatch) => {
+          const didClear = clearGhostText(state, dispatch);
+          if (didClear) {
+            dismissedContext = getTextBeforeCursor(state);
+          }
+          return didClear;
+        }
+      }),
+      handleDOMEvents: {
+        blur: view => {
+          const pluginState = autocompletePluginKey.getState(view.state);
+          if (pluginState !== null && pluginState !== void 0 && pluginState.ghostText) {
+            clearGhostText(view.state, view.dispatch);
+          }
+          return false;
+        },
+        focus: () => {
+          loadDefaultVocabulary();
+          loadVectorsAsync({
+            vectorsUrl: wordVectorsUrl
+          }).catch(() => {});
+          if (!hasIngestedPage) {
+            hasIngestedPage = true;
+            if (options !== null && options !== void 0 && options.getContext) {
+              options.getContext().then(context => {
+                var _context$siblingComme4;
+                if (!context) {
+                  return;
+                }
+                resolvedContext = context;
+                if (context.fullPageContent) {
+                  ingestDocumentPage(context.fullPageContent);
+                }
+                if (context.parentCommentContent) {
+                  ingestDocumentPage(context.parentCommentContent);
+                }
+                (_context$siblingComme4 = context.siblingCommentsContents) === null || _context$siblingComme4 === void 0 ? void 0 : _context$siblingComme4.forEach(ingestDocumentPage);
+              }).catch(() => {});
+            }
+          }
+          return false;
+        }
+      }
+    },
+    view: () => ({
+      update: (view, prevState) => {
+        if (!prevState.doc.eq(view.state.doc)) {
+          if (justAccepted) {
+            justAccepted = false;
+
+            // ✨ THE FIX: Memorize the text state right after acceptance.
+            // Any follow-up transactions will hit the 'dismissedContext'
+            // block and abort until the user actually types a new character!
+            dismissedContext = getTextBeforeCursor(view.state);
+
+            // Also clear any pending debounce timers from before the acceptance
+            if (debounceTimer) {
+              clearTimeout(debounceTimer);
+            }
+            return;
+          }
+          maybeUpdateSessionFrequency(view, prevState);
+          const textBefore = getTextBeforeCursor(view.state);
+          if (isWordBoundary(textBefore)) {
+            slowLaneClient.updateContext(buildSlowLaneText(view.state.doc.textContent, resolvedContext));
+          }
+          schedulePrediction(view);
+        }
+      },
+      destroy: () => {
+        if (debounceTimer) {
+          clearTimeout(debounceTimer);
+        }
+        setDefaultSlowLaneClient(null);
+      }
+    })
+  });
+};

package/dist/es2019/pm-plugins/ghost-text-decoration.js

@@ -0,0 +1,33 @@
+import { Decoration, DecorationSet } from '@atlaskit/editor-prosemirror/view';
+const GHOST_TEXT_CLASS = 'autocomplete-ghost-text';
+
+/**
+ * Creates a DecorationSet containing a ghost text widget at the given position.
+ * The ghost text is rendered as a styled <span> that appears after the cursor.
+ */
+export const createGhostTextDecorationSet = (state, position, text) => {
+  if (!text) {
+    return DecorationSet.empty;
+  }
+  const decoration = Decoration.widget(position, () => {
+    const container = document.createElement('span');
+    container.className = GHOST_TEXT_CLASS;
+    container.setAttribute('data-autocomplete-ghost', 'true');
+    container.style.color = '#999';
+    container.style.opacity = '0.6';
+    container.style.pointerEvents = 'none';
+    container.style.userSelect = 'none';
+    container.style.fontStyle = 'italic';
+    // U+200B (Zero Width Space) gives the browser a line-break opportunity
+    // immediately before the ghost text. This ensures the typed text before
+    // the span is never pushed to the next line by the ghost text's width —
+    // only the ghost text itself will wrap if it doesn't fit.
+    container.textContent = '\u200b' + text;
+    return container;
+  }, {
+    side: 1,
+    // Render after content at this position
+    key: 'autocomplete-ghost-text'
+  });
+  return DecorationSet.create(state.doc, [decoration]);
+};

package/dist/es2019/pm-plugins/scoring-pipeline.js

@@ -0,0 +1,221 @@
+/**
+ * Scoring Pipeline: Stage 1 (Semantic + Frequency), Grammar Filter, Stage 2 (LM Re-ranking).
+ *
+ * Operates synchronously on pre-loaded data. Each stage gracefully degrades
+ * when its required data isn't available (cold → warm → full warm).
+ */
+
+import posTagsData from './data/combined_l2_l3_pos_tags.json';
+import ghostPosTagsData from './data/ghost_pos_tags.json';
+import grammarTransitionsData from './data/grammar_transitions_10k.json';
+
+// ─── Types ──────────────────────────────────────────────────
+
+/** Metadata returned by the grammar filter for debug logging in the caller. */
+
+// ─── Scoring Constants ──────────────────────────────────────
+
+const ALPHA = 0.5;
+const BETA = 0.5;
+const NEUTRAL_SCORE = 0.5;
+const STAGE1_WEIGHT = 0.6;
+const STAGE2_WEIGHT = 0.4;
+const MIN_STAGE1_SCORE = 0.35;
+const L1_SESSION_CAP = 1.2;
+
+// ─── Grammar Data (loaded once on import) ───────────────────
+
+const posTags = new Map(Object.entries(posTagsData));
+const grammarTransitions = grammarTransitionsData;
+
+/**
+ * Precomputed map from each POS tag to the set of allowed next POS tags.
+ * Built once at module load from grammarTransitions so applyGrammarFilter
+ * never re-iterates the transition rules per call.
+ */
+const precomputedAllowedByPos = new Map(Object.entries(grammarTransitions.transitions).map(([pos, rule]) => [pos, new Set(rule.allowed)]));
+
+// ─── Math ───────────────────────────────────────────────────
+
+function cosineSimilarity(a, b) {
+  let dot = 0;
+  let normA = 0;
+  let normB = 0;
+  for (let i = 0; i < a.length; i++) {
+    dot += a[i] * b[i];
+    normA += a[i] * a[i];
+    normB += b[i] * b[i];
+  }
+  const dNormA = Math.sqrt(normA);
+  const dNormB = Math.sqrt(normB);
+  if (dNormA === 0 || dNormB === 0) {
+    return NEUTRAL_SCORE;
+  }
+  return (1 + dot / (dNormA * dNormB)) / 2;
+}
+
+// ─── Stage 1: Semantic + Frequency ──────────────────────────
+
+function scoreStage1(candidate, contextVector, getWordVector, maxTenantFreq) {
+  // 1. Calculate Base Global Score (Normalized Log)
+  const maxPossibleLog = Math.log10(maxTenantFreq + 1);
+
+  // Diversity Adjustment
+  const diversityRaw = (Math.log10(candidate.tenantFreq + 1) * 0.50 + Math.log10(candidate.docFreq + 1) * 0.25 + Math.log10(candidate.authorFreq + 1) * 0.25) / maxPossibleLog;
+  const sessionMultiplier = candidate.sessionFreq > 0 ? 1 + Math.log10(candidate.sessionFreq + 1) * 2.5 : 1;
+
+  // Apply multiplier; capped at L1_SESSION_CAP (default 1.2) to prevent excessive over-indexing
+  const freqScore = Math.min(diversityRaw * sessionMultiplier, L1_SESSION_CAP);
+
+  // 3. Semantic Scoring
+  let semanticScore = NEUTRAL_SCORE;
+  if (contextVector) {
+    const wordVec = getWordVector(candidate.word);
+    semanticScore = wordVec ? cosineSimilarity(contextVector, wordVec) : NEUTRAL_SCORE;
+  }
+  return {
+    semanticScore,
+    freqScore,
+    stage1Score: ALPHA * semanticScore + BETA * freqScore
+  };
+}
+
+// ─── Grammar Filter ─────────────────────────────────────────
+
+/**
+ * GHOST POS DICTIONARY
+ * A hardcoded mapping of common structural English words that were stripped 
+ * from the main domain vocabulary. This allows the grammar filter to understand 
+ * context without suggesting these words to the user.
+ */
+const ghostPosTags = ghostPosTagsData;
+function applyGrammarFilter(candidates, previousWord) {
+  if (!previousWord) return {
+    filtered: candidates,
+    grammarMeta: null
+  };
+  const lowerPrev = previousWord.toLowerCase();
+  const prevTags = ghostPosTags[lowerPrev] || posTags.get(lowerPrev);
+  if (!prevTags || prevTags.length === 0) {
+    return {
+      filtered: candidates,
+      grammarMeta: null
+    };
+  }
+  let allowedNextTags;
+  if (prevTags.length === 1) {
+    var _precomputedAllowedBy;
+    // Common case: single POS tag — reuse the precomputed Set directly (no allocation)
+    allowedNextTags = (_precomputedAllowedBy = precomputedAllowedByPos.get(prevTags[0])) !== null && _precomputedAllowedBy !== void 0 ? _precomputedAllowedBy : new Set();
+  } else {
+    allowedNextTags = new Set();
+    for (const pt of prevTags) {
+      const allowed = precomputedAllowedByPos.get(pt);
+      if (allowed) allowed.forEach(tag => allowedNextTags.add(tag));
+    }
+  }
+  const filtered = [];
+  const dropped = [];
+  for (const entry of candidates) {
+    const candidateTags = posTags.get(entry.candidate.word.toLowerCase());
+
+    // If candidate has no tags (unknown word), let it pass to be safe
+    if (!candidateTags || candidateTags.length === 0) {
+      filtered.push(entry);
+      continue;
+    }
+    if (candidateTags.some(ct => allowedNextTags.has(ct))) {
+      filtered.push(entry);
+    } else {
+      dropped.push(entry.candidate.word);
+    }
+  }
+  const finalFiltered = filtered.length > 0 ? filtered : candidates;
+  return {
+    filtered: finalFiltered,
+    grammarMeta: {
+      prevWord: lowerPrev,
+      prevTags,
+      before: candidates.length,
+      after: finalFiltered.length,
+      dropped: filtered.length > 0 ? dropped : []
+    }
+  };
+}
+
+// ─── Stage 2: LM Re-ranking ────────────────────────────────
+function getLmScore(word, lmLogits) {
+  if (!lmLogits) return 0;
+
+  // Look up the word directly! No more tokens.
+  const val = lmLogits[word.toLowerCase()];
+  if (typeof val === 'number') {
+    return val;
+  }
+  return 0;
+}
+
+// ─── Public API ─────────────────────────────────────────────
+
+export function rankCandidates(candidates, contextVector, getWordVector, lmLogits, maxTenantFreq, previousWord) {
+  // Stage 1
+  const stage1Results = candidates.map(candidate => {
+    const {
+      semanticScore,
+      freqScore,
+      stage1Score
+    } = scoreStage1(candidate, contextVector, getWordVector, maxTenantFreq);
+    return {
+      candidate,
+      semanticScore,
+      freqScore,
+      stage1Score
+    };
+  });
+  const stage1Survivors = stage1Results.filter(entry => entry.stage1Score >= MIN_STAGE1_SCORE);
+
+  // Grammar Filter
+  const {
+    filtered,
+    grammarMeta
+  } = applyGrammarFilter(stage1Survivors, previousWord);
+
+  // Stage 2 + final assembly
+  let lmMax = 0;
+  if (lmLogits && Object.keys(lmLogits).length > 0) {
+    const values = Object.values(lmLogits);
+    lmMax = Math.max(...values);
+  }
+  const scored = filtered.map(entry => {
+    let lmScore = 0;
+    let finalScore = entry.stage1Score;
+    if (lmLogits && Object.keys(lmLogits).length > 0) {
+      const rawLm = getLmScore(entry.candidate.word, lmLogits);
+      if (rawLm !== 0) {
+        // The word was in the top_k! Score it normally.
+        const logitDiff = Math.log(rawLm) - Math.log(lmMax);
+        lmScore = Math.exp(logitDiff);
+      } else {
+        lmScore = 0.05;
+      }
+      finalScore = STAGE1_WEIGHT * entry.stage1Score + STAGE2_WEIGHT * lmScore;
+    }
+    return {
+      word: entry.candidate.word,
+      freqScore: entry.freqScore,
+      semanticScore: entry.semanticScore,
+      lmScore,
+      finalScore
+    };
+  });
+  scored.sort((a, b) => {
+    if (b.finalScore !== a.finalScore) {
+      return b.finalScore - a.finalScore;
+    }
+    return a.word.length - b.word.length;
+  });
+  return {
+    candidates: scored,
+    grammarMeta
+  };
+}