@atlaskit/editor-plugin-autocomplete 0.1.0 → 0.2.0

package/src/pm-plugins/scoring-pipeline.ts

@@ -0,0 +1,297 @@
+/**
+ * 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).
+ */
+
+// resolveJsonModule is disabled for this package (see tsconfig.json) to prevent
+// TypeScript from parsing LFS pointer files during CI typecheck. JSON imports
+// are typed via the '*.json' declaration in typings.d.ts.
+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 ──────────────────────────────────────────────────
+
+export interface ScoringCandidate {
+    authorFreq: number;
+    docFreq: number;
+    sessionFreq: number;
+    tenantFreq: number;
+    word: string;
+}
+
+export interface ScoredCandidate {
+    finalScore: number;
+    freqScore: number;
+    lmScore: number;
+    semanticScore: number;
+    word: string;
+}
+
+/** Metadata returned by the grammar filter for debug logging in the caller. */
+export interface GrammarFilterMeta {
+    after: number;
+    before: number;
+    dropped: string[];
+    prevTags: string[];
+    prevWord: string;
+}
+
+interface PosTransitionRule {
+    allowed: string[];
+}
+
+interface GrammarTransitions {
+    transitions: Record<string, PosTransitionRule>;
+}
+
+// ─── 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: Map<string, string[]> = new Map(
+    Object.entries(posTagsData as Record<string, string[]>),
+);
+
+const grammarTransitions = grammarTransitionsData as GrammarTransitions;
+
+/**
+ * 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: Map<string, Set<string>> = new Map(
+	Object.entries(grammarTransitions.transitions).map(([pos, rule]) => [
+		pos,
+		new Set(rule.allowed),
+	]),
+);
+
+// ─── Math ───────────────────────────────────────────────────
+
+function cosineSimilarity(a: Float32Array, b: Float32Array): number {
+    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: ScoringCandidate,
+    contextVector: Float32Array | null,
+    getWordVector: (word: string) => Float32Array | null,
+    maxTenantFreq: number,
+): { freqScore: number; semanticScore: number; stage1Score: number } {
+
+	// 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: Record<string, string[]> = ghostPosTagsData as Record<string, string[]>;
+
+type FilterEntry = { candidate: ScoringCandidate; freqScore: number; semanticScore: number; stage1Score: number };
+
+function applyGrammarFilter(
+    candidates: FilterEntry[],
+    previousWord: string,
+): { filtered: FilterEntry[]; grammarMeta: GrammarFilterMeta | null } {
+    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: Set<string>;
+    if (prevTags.length === 1) {
+        // Common case: single POS tag — reuse the precomputed Set directly (no allocation)
+        allowedNextTags = precomputedAllowedByPos.get(prevTags[0]) ?? new Set();
+    } else {
+        allowedNextTags = new Set<string>();
+        for (const pt of prevTags) {
+            const allowed = precomputedAllowedByPos.get(pt);
+            if (allowed) allowed.forEach(tag => allowedNextTags.add(tag));
+        }
+    }
+
+    const filtered: FilterEntry[] = [];
+    const dropped: string[] = [];
+
+    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: string,
+    lmLogits: Record<string, number> | null
+): number {
+    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 interface RankCandidatesResult {
+    candidates: ScoredCandidate[];
+    grammarMeta: GrammarFilterMeta | null;
+}
+
+export function rankCandidates(
+    candidates: ScoringCandidate[],
+    contextVector: Float32Array | null,
+    getWordVector: (word: string) => Float32Array | null,
+    lmLogits: Record<string, number> | null,
+    maxTenantFreq: number,
+    previousWord: string,
+): RankCandidatesResult {
+    // 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: ScoredCandidate[] = 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 };
+}

package/src/pm-plugins/slow-lane-client/package.json

@@ -0,0 +1,15 @@
+{
+  "name": "@atlaskit/editor-plugin-autocomplete/src/pm-plugins/slow-lane-client",
+  "main": "../../../dist/cjs/pm-plugins/slow-lane-client.js",
+  "module": "../../../dist/esm/pm-plugins/slow-lane-client.js",
+  "module:es2019": "../../../dist/es2019/pm-plugins/slow-lane-client.js",
+  "sideEffects": false,
+  "types": "../../../dist/types/pm-plugins/slow-lane-client.d.ts",
+  "typesVersions": {
+    ">=4.5 <5.9": {
+      "*": [
+        "../../../dist/types-ts4.5/pm-plugins/slow-lane-client.d.ts"
+      ]
+    }
+  }
+}

package/src/pm-plugins/slow-lane-client.ts

@@ -0,0 +1,220 @@
+/**
+ * Slow Lane Client: Backend context encoding for autocomplete.
+ *
+ * Fires a BE request on word boundaries to encode document context.
+ * Expects the typeahead-encodings API format:
+ *   Request:  { text, session_id }
+ *   Response: { semantic_vector: number[], lm_logits: Record<string, number> }
+ */
+
+// ─── Types ───────────────────────────────────────────────────────────────────
+
+/** Request payload for typeahead-encodings endpoint. */
+export interface TypeaheadEncodingsRequest {
+	session_id: string;
+	text: string;
+}
+
+/** Response from typeahead-encodings endpoint. */
+export interface TypeaheadEncodingsResponse {
+	lm_logits: Record<string, number>;
+	semantic_vector: number[];
+}
+
+// ─── Constants ───────────────────────────────────────────────────────────────
+
+const WORD_BOUNDARY_CHARS = /[\s.,;:!?]/u;
+const DEFAULT_DEBOUNCE_MS = 300;
+
+/**
+ * Check if text ends with a word boundary character (space or punctuation).
+ */
+export const isWordBoundary = (text: string): boolean => {
+	if (!text || text.length === 0) {
+		return false;
+	}
+	const lastChar = text[text.length - 1];
+	return WORD_BOUNDARY_CHARS.test(lastChar);
+};
+
+// ─── Slow Lane Client ────────────────────────────────────────────────────────
+
+export interface SlowLaneClientConfig {
+	baseUrl: string;
+	debounceMs?: number;
+	endpoint?: string;
+	fetchFn?: typeof fetch;
+	onUpdate?: (opts: { hasLmLogits: boolean; hasVector: boolean; textLength: number }) => void;
+	productKey?: string;
+	sessionId?: string;
+}
+
+export const createSlowLaneClient = (
+	config: SlowLaneClientConfig,
+): {
+	getContextVector: () => Float32Array | null;
+	getLmLogits: () => Record<string, number> | null;
+	isWordBoundary: (text: string) => boolean;
+	setContextVector: (vector: Float32Array | null) => void;
+	setLmLogits: (logits: Record<string, number> | null) => void;
+	updateContext: (text: string) => void;
+} => {
+	const {
+		baseUrl,
+		sessionId: configSessionId,
+		productKey = 'confluence',
+		endpoint = '/gateway/api/v1/autocomplete/typeahead-encodings',
+		debounceMs = DEFAULT_DEBOUNCE_MS,
+		fetchFn = fetch,
+		onUpdate,
+	} = config;
+
+	const sessionId = configSessionId ?? crypto.randomUUID();
+
+	const headers: Record<string, string> = {
+		'Content-Type': 'application/json',
+		'x-experience-id': 'confluence-smart-typeahead-encodings',
+		'x-product': productKey,
+	};
+
+	let debounceTimer: ReturnType<typeof setTimeout> | null = null;
+	let lastRequestedText = '';
+	let storedContextVector: Float32Array | null = null;
+	let storedLmLogits: Record<string, number> | null = null;
+
+	const doUpdateContext = async (text: string): Promise<void> => {
+		if (!text || text.trim().length === 0) {
+			return;
+		}
+
+		const url = `${baseUrl.replace(/\/$/u, '')}${endpoint}`;
+		const payload: TypeaheadEncodingsRequest = {
+			text,
+			session_id: sessionId,
+		};
+
+		// Log the exact payload being sent so the comment/reply structure is visible
+		// eslint-disable-next-line no-console
+		console.groupCollapsed(
+			`%c[SlowLane] %c📤 Sending context | ${text.length} chars`,
+			'color: #9c27b0; font-weight: bold;',
+			'color: inherit;',
+		);
+		text.split('\n').forEach((line, i, arr) => {
+			// eslint-disable-next-line no-console
+			console.log(`  ${i === arr.length - 1 ? '▶' : ' '} ${line}`);
+		});
+		// eslint-disable-next-line no-console
+		console.groupEnd();
+
+		try {
+			const res = await fetchFn(url, {
+				method: 'POST',
+				headers,
+				body: JSON.stringify(payload),
+			});
+
+			if (!res.ok) {
+				storedContextVector = null;
+				storedLmLogits = null;
+				// eslint-disable-next-line no-console
+				console.log(
+					`%c[SlowLane] %c❌ Request failed (${res.status})`,
+					'color: #9c27b0; font-weight: bold;',
+					'color: #f44336;',
+				);
+				return;
+			}
+
+			const data = (await res.json()) as TypeaheadEncodingsResponse;
+
+			if (data.semantic_vector && Array.isArray(data.semantic_vector)) {
+				storedContextVector = new Float32Array(data.semantic_vector);
+			} else {
+				storedContextVector = null;
+			}
+
+			if (data.lm_logits && typeof data.lm_logits === 'object') {
+				storedLmLogits = data.lm_logits;
+			} else {
+				storedLmLogits = null;
+			}
+
+			// Log what came back so it can be correlated with the next prediction group
+			// eslint-disable-next-line no-console
+			console.groupCollapsed(
+				`%c[SlowLane] %c📥 Response received`,
+				'color: #9c27b0; font-weight: bold;',
+				'color: inherit;',
+			);
+			// eslint-disable-next-line no-console
+			console.log(
+				storedContextVector
+					? `✅ semantic_vector: ${storedContextVector.length} dims`
+					: '❌ No semantic_vector',
+			);
+			// eslint-disable-next-line no-console
+			console.log(
+				storedLmLogits
+					? `✅ lm_logits: ${Object.keys(storedLmLogits).length} tokens`
+					: '❌ No lm_logits',
+			);
+			// eslint-disable-next-line no-console
+			console.groupEnd();
+
+			onUpdate?.({
+				textLength: text.length,
+				hasVector: storedContextVector !== null,
+				hasLmLogits: storedLmLogits !== null,
+			});
+			// eslint-disable-next-line no-unused-vars
+		} catch (e) {
+			storedContextVector = null;
+			storedLmLogits = null;
+			// eslint-disable-next-line no-console
+			console.log(
+				'%c[SlowLane] %c❌ Network error — context cleared',
+				'color: #9c27b0; font-weight: bold;',
+				'color: #f44336;',
+			);
+		}
+	};
+
+	const updateContextDebounced = (text: string): void => {
+		if (debounceTimer) {
+			clearTimeout(debounceTimer);
+		}
+		lastRequestedText = text;
+		debounceTimer = setTimeout(() => {
+			debounceTimer = null;
+			doUpdateContext(lastRequestedText);
+		}, debounceMs);
+	};
+
+	return {
+		updateContext: updateContextDebounced,
+		getContextVector: () => storedContextVector,
+		getLmLogits: () => storedLmLogits,
+		setContextVector: (vector) => {
+			storedContextVector = vector;
+		},
+		setLmLogits: (logits) => {
+			storedLmLogits = logits;
+		},
+		isWordBoundary,
+	};
+};
+
+let defaultSlowLaneClient: ReturnType<typeof createSlowLaneClient> | null = null;
+
+export const setDefaultSlowLaneClient = (
+	client: ReturnType<typeof createSlowLaneClient> | null,
+): void => {
+	defaultSlowLaneClient = client;
+};
+
+export const getStoredContextVector = (): Float32Array | null =>
+	defaultSlowLaneClient?.getContextVector() ?? null;
+
+export const getStoredLmLogits = (): Record<string, number> | null =>
+	defaultSlowLaneClient?.getLmLogits() ?? null;

package/src/pm-plugins/text-predictor/package.json

@@ -0,0 +1,15 @@
+{
+  "name": "@atlaskit/editor-plugin-autocomplete/src/pm-plugins/text-predictor",
+  "main": "../../../dist/cjs/pm-plugins/text-predictor.js",
+  "module": "../../../dist/esm/pm-plugins/text-predictor.js",
+  "module:es2019": "../../../dist/es2019/pm-plugins/text-predictor.js",
+  "sideEffects": false,
+  "types": "../../../dist/types/pm-plugins/text-predictor.d.ts",
+  "typesVersions": {
+    ">=4.5 <5.9": {
+      "*": [
+        "../../../dist/types-ts4.5/pm-plugins/text-predictor.d.ts"
+      ]
+    }
+  }
+}