change(web): rework applied-suggestion tracking + reversions

This commit is contained in:
Joshua Horton 2025-08-01 10:25:23 -05:00
parent 2490ad322d
commit c67ea99484
8 changed files with 120 additions and 106 deletions

View file

@ -14,7 +14,6 @@ import { SearchSpace } from "./distance-modeler.js";
import Distribution = LexicalModelTypes.Distribution;
import LexicalModel = LexicalModelTypes.LexicalModel;
import Suggestion = LexicalModelTypes.Suggestion;
import Transform = LexicalModelTypes.Transform;
/**
@ -50,22 +49,6 @@ export class ContextToken {
*/
readonly searchSpace: SearchSpace;
/* The next two fields will **not land here** in the final version for
epic/autocorrect / 19.0-beta! That said, their future location has
not yet been reworked, so we'll keep them here for now. */
/**
* The set of suggestions generated for the current token
*/
suggestions: Suggestion[];
/**
* The ID of the suggestion applied to the current token, if any.
*
* Should be set to undefined when no such suggestion exists.
*/
appliedSuggestionId?: number;
/**
* Constructs a new, empty instance for use with the specified LexicalModel.
* @param model
@ -92,12 +75,6 @@ export class ContextToken {
// In case we are unable to perfectly track context (say, due to multitaps)
// we need to ensure that only fully-utilized keystrokes are considered.
this.searchSpace = new SearchSpace(priorToken.searchSpace);
this.suggestions = priorToken.suggestions.slice();
// because of unit tests.
if(priorToken.appliedSuggestionId !== undefined) {
this.appliedSuggestionId = priorToken.appliedSuggestionId;
}
} else {
const model = param;
@ -112,8 +89,6 @@ export class ContextToken {
return [{sample: transform, p: 1.0}];
});
rawTransformDistributions.forEach((entry) => this.searchSpace.addInput(entry));
this.suggestions = [];
}
}

View file

@ -300,12 +300,19 @@ export class ContextTracker {
let result = ContextTracker.attemptMatchContext(context, model, priorMatchState.final, transformDistribution);
if(result?.final) {
if(priorMatchState.transitionId !== undefined) {
// Already has a taggedContext.
this.cache.get(priorMatchState.transitionId);
}
if(transitionId !== undefined) {
if(priorMatchState.transitionId != transitionId) {
// Already has a taggedContext.
this.cache.get(transformDistribution.find((entry) => entry.sample.id !== undefined).sample.id);
// Special case: if base and final match, we should use the old Transition instance.
// This is currently used in some unit tests.
if(result.final != result.base) {
this.cache.add(transitionId, result);
} else {
return this.cache.peek(priorMatchState.transitionId);
}
this.cache.add(transitionId, result);
}
return result;

View file

@ -12,7 +12,10 @@ import { LexicalModelTypes } from '@keymanapp/common-types';
import { ContextState } from './context-state.js';
import Distribution = LexicalModelTypes.Distribution;
import Suggestion = LexicalModelTypes.Suggestion;
import Transform = LexicalModelTypes.Transform;
import { buildMergedTransform } from '@keymanapp/models-templates';
import { ContextTracker } from './context-tracker.js';
/**
* Represents the transition between two context states as triggered
@ -112,4 +115,44 @@ export class ContextTransition {
this._transitionId = inputDistribution?.find((entry) => entry.sample.id !== undefined)?.sample.id;
this.preservationTransform = preservationTransform;
}
/**
* Applies a suggestion generated from this context transition on top of the transition itself,
* replacing its final context state. This does _not_, however, replace the original fat-finger
* distribution or other intermediate data regarding associated keystrokes.
* @param suggestion
* @returns
*/
applySuggestion(suggestion: Suggestion) {
const fullTransform = suggestion.appendedTransform
? buildMergedTransform(suggestion.transform, suggestion.appendedTransform)
: suggestion.transform;
// An applied suggestion should replace the original Transition's effects, though keeping
// the original input around.
const appliedState = ContextTracker.attemptMatchContext(
this.base.context,
this.base.model,
this.base,
[{sample: fullTransform, p: 1}]
).final;
const preAppliedState = this.final;
if(!preAppliedState.suggestions.find((s) => s.id == suggestion?.id)) {
throw new Error("Could not find matching suggestion to apply");
}
// Start from a deep copy, then replace as needed to overwrite with the context
// state resulting from the suggestion while preserving suggestion + primary
// keystroke data.
const resultTransition = new ContextTransition(this);
resultTransition._final = appliedState;
resultTransition._transitionId = suggestion.transformId;
appliedState.appliedSuggestionId = suggestion.id;
appliedState.appliedInput = preAppliedState.appliedInput;
appliedState.suggestions = preAppliedState.suggestions;
return resultTransition;
}
}

View file

@ -191,7 +191,7 @@ export class ModelCompositor {
// Store the suggestions on the final token of the current context state (if it exists).
// Or, once phrase-level suggestions are possible, on whichever token serves as each prediction's root.
if(postContextState) {
postContextState.tokenization.tail.suggestions = suggestions;
postContextState.suggestions = suggestions;
}
return suggestions;
@ -253,17 +253,13 @@ export class ModelCompositor {
// Step 3: if we track Contexts, update the tracking data as appropriate.
if(this.contextTracker) {
let contextState = this.contextTracker.newest.final;
if(!contextState) {
contextState = this.contextTracker.analyzeState(this.lexicalModel, context).final;
let originalTransition = this.contextTracker.newest;
if(!originalTransition) {
originalTransition = this.contextTracker.analyzeState(this.lexicalModel, context);
}
contextState.tokenization.tail.appliedSuggestionId = suggestion.id;
let acceptedContext = models.applyTransform(suggestion.transform, context);
if(suggestion.appendedTransform) {
acceptedContext = models.applyTransform(suggestion.appendedTransform, acceptedContext);
}
this.contextTracker.analyzeState(this.lexicalModel, acceptedContext);
const appliedTransition = originalTransition.applySuggestion(suggestion);
this.contextTracker.cache.add(suggestion.transformId, appliedTransition);
}
return reversion;
@ -273,6 +269,7 @@ export class ModelCompositor {
// If we are unable to track context (because the model does not support LexiconTraversal),
// we need a "fallback" strategy.
let compositor = this;
let suggestions: Promise<Suggestion[]>;
let fallbackSuggestions = async function() {
let revertedContext = models.applyTransform(reversion.transform, context);
const suggestions = await compositor.predict({insert: '', deleteLeft: 0}, revertedContext);
@ -293,27 +290,42 @@ export class ModelCompositor {
}
// When the context is tracked, we prefer the tracked information.
let contextMatchFound = this.contextTracker.cache.peek(-reversion.transformId);
let originalTransition = this.contextTracker.cache.peek(-reversion.transformId);
if(!contextMatchFound) {
return fallbackSuggestions();
if(!originalTransition) {
suggestions = fallbackSuggestions();
}
// Remove all contexts more recent than the one we're reverting to.
this.contextTracker.cache.rewindTo(-reversion.transformId);
this.contextTracker.newest.final.tokenization.tail.appliedSuggestionId = undefined;
// Will need to be modified a bit if/when phrase-level suggestions are implemented.
// Those will be tracked on the first token of the phrase, which won't be the tail
// if they cover multiple tokens.
let suggestions = this.contextTracker.newest.final.tokenization.tail.suggestions;
if(!suggestions) {
// Will need to be modified a bit if/when phrase-level suggestions are implemented.
// Those will be tracked on the first token of the phrase, which won't be the tail
// if they cover multiple tokens.
let suggests = this.contextTracker.newest.final.suggestions;
suggests.forEach(function(suggestion) {
// A reversion's transform ID is the additive inverse of its original suggestion;
// we revert to the state of said original suggestion.
suggestion.transformId = -reversion.transformId;
suggestion.autoAccept = false;
});
suggestions = Promise.resolve(suggests);
}
// An applied reversion should replace the original Transition's effects.
const revertedTransition = correction.ContextTracker.attemptMatchContext(
models.applyTransform(originalTransition.inputDistribution[0].sample, originalTransition.base.context),
this.lexicalModel,
originalTransition.base,
originalTransition.inputDistribution
);
revertedTransition.final.suggestions = await suggestions;
this.contextTracker.cache.add(-reversion.transformId, revertedTransition);
suggestions.forEach(function(suggestion) {
// A reversion's transform ID is the additive inverse of its original suggestion;
// we revert to the state of said original suggestion.
suggestion.transformId = -reversion.transformId;
suggestion.autoAccept = false;
});
return suggestions;
}

View file

@ -29,8 +29,6 @@ describe('ContextToken', function() {
assert.isEmpty(token.searchSpace.inputSequence);
assert.isEmpty(token.exampleInput);
assert.isFalse(token.isWhitespace);
assert.isEmpty(token.suggestions);
assert.isUndefined(token.appliedSuggestionId);
// While searchSpace has no inputs, it _can_ match lexicon entries (via insertions).
let searchIterator = token.searchSpace.getBestMatches(new ExecutionTimer(Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY));
@ -57,44 +55,13 @@ describe('ContextToken', function() {
assert.equal(token.exampleInput, 'and');
assert.isFalse(token.isWhitespace);
// Is only set with a different value later, outside the constructor.
assert.isEmpty(token.suggestions);
assert.isUndefined(token.appliedSuggestionId);
});
it("(token: ContextToken", () => {
// Same as in a test above, since we verified that it works correctly.
let baseToken = new ContextToken(plainModel, "and");
baseToken.suggestions = [
{
transform: {
insert: 'd ',
deleteLeft: 0
},
id: 37,
transformId: 1,
displayAs: '"and"',
tag: 'keep',
autoAccept: true
},
{
transform: {
insert: 'Andes ',
deleteLeft: 2
},
id: 38,
transformId: 1,
displayAs: 'Andes',
}
]
baseToken.appliedSuggestionId = 37;
let clonedToken = new ContextToken(baseToken);
assert.notEqual(clonedToken.suggestions, baseToken.suggestions);
assert.deepEqual(clonedToken.suggestions, baseToken.suggestions);
assert.notEqual(clonedToken.searchSpace, baseToken.searchSpace);
// Deep equality on .searchSpace can't be directly checked due to the internal complexities involved.
// We CAN check for the most important members, though.

View file

@ -42,7 +42,6 @@ describe('ContextTokenization', function() {
assert.isNotOk(tokenization.alignment);
assert.equal(tokenization.tail.exampleInput, 'day');
assert.isFalse(tokenization.tail.isWhitespace);
assert.isUndefined(tokenization.tail.appliedSuggestionId);
});
it("constructs from a token array + alignment data", () => {
@ -63,7 +62,6 @@ describe('ContextTokenization', function() {
assert.deepEqual(tokenization.alignment, alignment);
assert.equal(tokenization.tail.exampleInput, 'day');
assert.isFalse(tokenization.tail.isWhitespace);
assert.isUndefined(tokenization.tail.appliedSuggestionId);
});
it('clones', () => {

View file

@ -23,6 +23,8 @@ const plainModel = new TrieModel(jsonFixture('models/tries/english-1000'),
{wordBreaker: defaultBreaker}
);
var emptyInput = (id: number) => [{sample: {insert: '', deleteLeft: 0, id: id}, p: 1}];
describe('ContextTracker', function() {
function toWrapperDistribution(transform: Transform) {
return [{
@ -388,11 +390,15 @@ describe('ContextTracker', function() {
it('tracks an accepted suggestion', function() {
let baseSuggestion: Suggestion = {
transform: {
insert: 'world ',
insert: 'world',
deleteLeft: 3,
id: 1
id: 2
},
transformId: 0,
appendedTransform: {
insert: ' ',
deleteLeft: 0
},
transformId: 2,
id: 1,
displayAs: 'world'
};
@ -415,20 +421,20 @@ describe('ContextTracker', function() {
let model = new models.TrieModel(jsonFixture('models/tries/english-1000'), options);
let compositor = new ModelCompositor(model);
let baseContextMatch = compositor.contextTracker.analyzeState(model, baseContext, [{sample: {insert: '', deleteLeft: 0, id: 0}, p: 1}]);
baseContextMatch.final.tokenization.tail.suggestions = [ baseSuggestion ];
let preAppliedTransition = compositor.contextTracker.analyzeState(model, baseContext, emptyInput(0));
// Mocks a prior prediction request without having done it.
preAppliedTransition.final.suggestions = [baseSuggestion];
let reversion = compositor.acceptSuggestion(baseSuggestion, baseContext, postTransform);
// Actual test assertion - was the replacement tracked?
assert.equal(baseContextMatch.final.tokenization.tail.appliedSuggestionId, baseSuggestion.id);
assert.isUndefined(preAppliedTransition.final.appliedSuggestionId);
assert.equal(reversion.id, -baseSuggestion.id);
compositor.contextTracker.cache.keys().forEach((key) => assert.isDefined(key));
// Next step - on the followup context, is the replacement still active?
let postContext = models.applyTransform(baseSuggestion.transform, baseContext);
let postContextMatch = compositor.contextTracker.analyzeState(model, postContext);
let postContext = models.applyTransform(baseSuggestion.appendedTransform, models.applyTransform(baseSuggestion.transform, baseContext));
let postContextMatch = compositor.contextTracker.analyzeState(model, postContext, emptyInput(2));
assert.equal(postContextMatch.final.appliedSuggestionId, baseSuggestion.id);
// Penultimate token corresponds to whitespace, which does not have a 'raw' representation.
assert.equal(postContextMatch.final.tokenization.tokens[postContextMatch.final.tokenization.tokens.length - 2].exampleInput, ' ');

View file

@ -19,6 +19,8 @@ import Suggestion = LexicalModelTypes.Suggestion;
import Transform = LexicalModelTypes.Transform;
import TrieModel = models.TrieModel;
var emptyInput = (id: number) => [{sample: {insert: '', deleteLeft: 0, id: id}, p: 1}];
describe('ModelCompositor', function() {
describe('Prediction with 14.0+ models', function() {
describe('Basic suggestion generation', function() {
@ -879,6 +881,8 @@ describe('ModelCompositor', function() {
let model = new models.TrieModel(jsonFixture('models/tries/english-1000'), {punctuation: englishPunctuation});
let compositor = new ModelCompositor(model, true);
compositor.contextTracker.analyzeState(model, baseContext, emptyInput(0));
let initialSuggestions = await compositor.predict(postTransform, baseContext);
const suggestionContextState = compositor.contextTracker.newest;
@ -890,26 +894,28 @@ describe('ModelCompositor', function() {
assert.equal(compositor.contextTracker.cache.size, 2);
let contextIds = compositor.contextTracker.cache.keys();
let transitionInstances = compositor.contextTracker.cache.keys().map((key) => compositor.contextTracker.cache.get(key));
let transitionInstances = compositor.contextTracker.cache.keys().map((key) => compositor.contextTracker.cache.peek(key));
let baseSuggestion = initialSuggestions[1];
let reversion = compositor.acceptSuggestion(baseSuggestion, baseContext, postTransform);
assert.equal(reversion.transformId, -baseSuggestion.transformId);
assert.equal(reversion.id, -baseSuggestion.id);
let postContext = models.applyTransform(baseSuggestion.appendedTransform, models.applyTransform(baseSuggestion.transform, baseContext));
const appliedContextState = compositor.contextTracker.analyzeState(model, postContext, emptyInput(13));
// Accepting the suggestion rewrites the latest context transition.
assert.equal(compositor.contextTracker.cache.size, 2);
assert.sameMembers(compositor.contextTracker.cache.keys(), contextIds);
assert.notSameDeepMembers(compositor.contextTracker.cache.keys().map((key) => compositor.contextTracker.cache.get(key)), transitionInstances);
assert.notSameDeepMembers(compositor.contextTracker.cache.keys().map((key) => compositor.contextTracker.cache.peek(key)), transitionInstances);
// The replacement should be marked on the context-tracking token.
assert.isAtLeast(suggestionContextState.final.tokenization.tail.appliedSuggestionId, 0);
// The replacement should be marked on the context-tracking token for the applied version of the results.
assert.equal(suggestionContextState.final.appliedSuggestionId, undefined);
assert.isAtLeast(appliedContextState.final.appliedSuggestionId, 0);
let appliedContext = models.applyTransform(baseSuggestion.transform, baseContext);
compositor.applyReversion(reversion, appliedContext);
// The replacement should no longer be marked for the context-tracking token.
assert.isNotOk(suggestionContextState.final.tokenization.tail.appliedSuggestionId);
await compositor.applyReversion(reversion, appliedContext);
assert.isUndefined(compositor.contextTracker.cache.peek(13).final.appliedSuggestionId);
});
});
});