ML Suggested Edit v2

This is experiment for introducing more complex suggestions.

V1 restricts suggestions to the comment range. For example,
a suggestion to rename a variable would not apply the change globally,
nor would Gerrit automatically include necessary imports
for the suggestion. V2 will support suggestions on multiple lines.

Feature flag UiFeature__ml_suggested_edit_v2 is introduced.

Release-Notes: skip
Google-Bug-Id: b/322003830
Change-Id: Iab9cbccac2e787a1480c0361bfb81d6a290126e8
diff --git a/polygerrit-ui/app/elements/shared/gr-comment/gr-comment.ts b/polygerrit-ui/app/elements/shared/gr-comment/gr-comment.ts
index b8b4b9a..fbaafd9 100644
--- a/polygerrit-ui/app/elements/shared/gr-comment/gr-comment.ts
+++ b/polygerrit-ui/app/elements/shared/gr-comment/gr-comment.ts
@@ -63,7 +63,11 @@
 import {Subject} from 'rxjs';
 import {debounceTime} from 'rxjs/operators';
 import {changeModelToken} from '../../../models/change/change-model';
-import {ChangeInfo, isBase64FileContent} from '../../../api/rest-api';
+import {
+  ChangeInfo,
+  FixSuggestionInfoInput,
+  isBase64FileContent,
+} from '../../../api/rest-api';
 import {createDiffUrl} from '../../../models/views/change';
 import {userModelToken} from '../../../models/user/user-model';
 import {modalStyles} from '../../../styles/gr-modal-styles';
@@ -219,6 +223,9 @@
   generatedSuggestion?: Suggestion;
 
   @state()
+  generatedFixSuggestion?: FixSuggestionInfoInput;
+
+  @state()
   generatedSuggestionId?: string;
 
   @state()
@@ -366,7 +373,10 @@
       () => this.getConfigModel().docsBaseUrl$,
       docsBaseUrl => (this.docsBaseUrl = docsBaseUrl)
     );
-    if (this.flagsService.isEnabled(KnownExperimentId.ML_SUGGESTED_EDIT)) {
+    if (
+      this.flagsService.isEnabled(KnownExperimentId.ML_SUGGESTED_EDIT) ||
+      this.flagsService.isEnabled(KnownExperimentId.ML_SUGGESTED_EDIT_V2)
+    ) {
       subscribe(
         this,
         () =>
@@ -987,7 +997,8 @@
 
   private showGeneratedSuggestion() {
     return (
-      this.flagsService.isEnabled(KnownExperimentId.ML_SUGGESTED_EDIT) &&
+      (this.flagsService.isEnabled(KnownExperimentId.ML_SUGGESTED_EDIT) ||
+        this.flagsService.isEnabled(KnownExperimentId.ML_SUGGESTED_EDIT_V2)) &&
       this.suggestionsProvider &&
       this.editing &&
       !this.permanentEditingMode &&
@@ -1007,18 +1018,22 @@
   }
 
   private renderGeneratedSuggestionPreview() {
-    if (
-      !this.showGeneratedSuggestion() ||
-      !this.generateSuggestion ||
-      !this.generatedSuggestion
-    )
+    if (!this.showGeneratedSuggestion() || !this.generateSuggestion)
       return nothing;
 
-    return html`<gr-suggestion-diff-preview
-      .showAddSuggestionButton=${true}
-      .suggestion=${this.generatedSuggestion?.replacement}
-      .uuid=${this.generatedSuggestionId}
-    ></gr-suggestion-diff-preview>`;
+    if (this.generatedFixSuggestion) {
+      return html`<gr-suggestion-diff-preview
+        .fixReplacementInfos=${this.generatedFixSuggestion.replacements}
+      ></gr-suggestion-diff-preview>`;
+    } else if (this.generatedSuggestion) {
+      return html`<gr-suggestion-diff-preview
+        .showAddSuggestionButton=${true}
+        .suggestion=${this.generatedSuggestion?.replacement}
+        .uuid=${this.generatedSuggestionId}
+      ></gr-suggestion-diff-preview>`;
+    } else {
+      return nothing;
+    }
   }
 
   private renderGenerateSuggestEditButton() {
@@ -1089,7 +1104,17 @@
     }${USER_SUGGESTION_START_PATTERN}${code}${'\n```'}`;
   }
 
-  private async generateSuggestEdit() {
+  private generateSuggestEdit() {
+    if (this.flagsService.isEnabled(KnownExperimentId.ML_SUGGESTED_EDIT_V2)) {
+      this.generateSuggestEdit_v2();
+    } else if (
+      this.flagsService.isEnabled(KnownExperimentId.ML_SUGGESTED_EDIT)
+    ) {
+      this.generateSuggestEdit_v1();
+    }
+  }
+
+  private async generateSuggestEdit_v1() {
     const suggestionsProvider = this.suggestionsProvider;
     const changeInfo = this.getChangeModel().getChange();
     if (
@@ -1130,6 +1155,7 @@
     // other suggestions.
     this.reporting.reportInteraction(Interaction.GENERATE_SUGGESTION_RESPONSE, {
       uuid: this.generatedSuggestionId,
+      type: 'suggest-code',
       response: suggestionResponse.responseCode,
       numSuggestions: suggestionResponse.suggestions.length,
       hasNewRange: suggestionResponse.suggestions?.[0]?.newRange !== undefined,
@@ -1139,6 +1165,56 @@
     this.generatedSuggestion = suggestion;
   }
 
+  private async generateSuggestEdit_v2() {
+    const suggestionsProvider = this.suggestionsProvider;
+    const changeInfo = this.getChangeModel().getChange();
+    if (
+      !suggestionsProvider?.suggestFix ||
+      !this.showGeneratedSuggestion() ||
+      !this.generateSuggestion ||
+      !changeInfo ||
+      !this.comment ||
+      !this.comment.patch_set ||
+      !this.comment.path ||
+      this.messageText.length === 0
+    )
+      return;
+    this.generatedSuggestionId = uuid();
+    this.reporting.reportInteraction(Interaction.GENERATE_SUGGESTION_REQUEST, {
+      uuid: this.generatedSuggestionId,
+      type: 'suggest-fix',
+      commentId: this.comment.id,
+    });
+    this.suggestionLoading = true;
+    let suggestionResponse;
+    try {
+      suggestionResponse = await suggestionsProvider.suggestFix({
+        prompt: this.messageText,
+        changeInfo: changeInfo as ChangeInfo,
+        patchsetNumber: this.comment?.patch_set,
+        filePath: this.comment.path,
+        range: this.comment.range,
+        lineNumber: this.comment.line,
+      });
+    } finally {
+      this.suggestionLoading = false;
+    }
+
+    if (!suggestionResponse) return;
+    // TODO(milutin): The suggestionResponse can contain multiple suggestion
+    // options. We pick the first one for now. In future we shouldn't ignore
+    // other suggestions.
+    this.reporting.reportInteraction(Interaction.GENERATE_SUGGESTION_RESPONSE, {
+      uuid: this.generatedSuggestionId,
+      type: 'suggest-fix',
+      response: suggestionResponse.responseCode,
+      numSuggestions: suggestionResponse.fix_suggestions.length,
+    });
+    const suggestion = suggestionResponse.fix_suggestions?.[0];
+    if (!suggestion) return;
+    this.generatedFixSuggestion = suggestion;
+  }
+
   private renderRobotActions() {
     if (!this.account || !isRobot(this.comment)) return;
     const endpoint = html`
diff --git a/polygerrit-ui/app/elements/shared/gr-suggestion-diff-preview/gr-suggestion-diff-preview.ts b/polygerrit-ui/app/elements/shared/gr-suggestion-diff-preview/gr-suggestion-diff-preview.ts
index f591388..ee4545a 100644
--- a/polygerrit-ui/app/elements/shared/gr-suggestion-diff-preview/gr-suggestion-diff-preview.ts
+++ b/polygerrit-ui/app/elements/shared/gr-suggestion-diff-preview/gr-suggestion-diff-preview.ts
@@ -19,7 +19,7 @@
 import {GrSyntaxLayerWorker} from '../../../embed/diff/gr-syntax-layer/gr-syntax-layer-worker';
 import {resolve} from '../../../models/dependency';
 import {highlightServiceToken} from '../../../services/highlight/highlight-service';
-import {NumericChangeId} from '../../../api/rest-api';
+import {FixReplacementInfo, NumericChangeId} from '../../../api/rest-api';
 import {changeModelToken} from '../../../models/change/change-model';
 import {subscribe} from '../../lit/subscription-controller';
 import {FilePreview} from '../../diff/gr-apply-fix-dialog/gr-apply-fix-dialog';
@@ -43,11 +43,20 @@
   code: string;
 }
 
+/**
+ * Diff preview for
+ * 1. suggestion vs commented Text
+ * or 2. fixReplacementInfos
+ * that are attached to a comment.
+ */
 @customElement('gr-suggestion-diff-preview')
 export class GrSuggestionDiffPreview extends LitElement {
   @property({type: String})
   suggestion?: string;
 
+  @property({type: Object})
+  fixReplacementInfos?: FixReplacementInfo[];
+
   @property({type: Boolean})
   showAddSuggestionButton = false;
 
@@ -64,7 +73,7 @@
   layers: DiffLayer[] = [];
 
   @state()
-  previewLoadedFor?: string;
+  previewLoadedFor?: string | FixReplacementInfo[];
 
   @state() repo?: RepoName;
 
@@ -167,10 +176,16 @@
         this.fetchFixPreview();
       }
     }
+
+    if (changed.has('changeNum') || changed.has('comment')) {
+      if (this.previewLoadedFor !== this.fixReplacementInfos) {
+        this.fetchFixReplacementInfosPreview();
+      }
+    }
   }
 
   override render() {
-    if (!this.suggestion) return nothing;
+    if (!this.suggestion && !this.fixReplacementInfos) return nothing;
     const code = this.suggestion;
     return html`
       ${when(
@@ -244,6 +259,46 @@
     return res;
   }
 
+  private async fetchFixReplacementInfosPreview() {
+    if (
+      this.suggestion ||
+      !this.changeNum ||
+      !this.comment?.patch_set ||
+      !this.fixReplacementInfos
+    )
+      return;
+
+    // TODO (milutin): This is a temporary fix for the broken path issue.
+    // Our experimental plugin currently returns only the file extension.
+    const replacements = this.fixReplacementInfos.map(fixInfo => {
+      return {
+        ...fixInfo,
+        path: this.comment?.path ?? fixInfo.path,
+      };
+    });
+
+    this.reporting.time(Timing.PREVIEW_FIX_LOAD);
+    const res = await this.restApiService.getFixPreview(
+      this.changeNum,
+      this.comment?.patch_set,
+      replacements
+    );
+
+    if (!res) return;
+    const currentPreviews = Object.keys(res).map(key => {
+      return {filepath: key, preview: res[key]};
+    });
+    this.reporting.timeEnd(Timing.PREVIEW_FIX_LOAD, {
+      uuid: this.uuid,
+    });
+    if (currentPreviews.length > 0) {
+      this.preview = currentPreviews[0];
+      this.previewLoadedFor = this.fixReplacementInfos;
+    }
+
+    return res;
+  }
+
   /**
    * Applies a fix previewed in `suggestion-diff-preview`,
    * navigating to the new change URL with the EDIT patchset.
diff --git a/polygerrit-ui/app/services/flags/flags.ts b/polygerrit-ui/app/services/flags/flags.ts
index b7d36f4..b59d7a8b3 100644
--- a/polygerrit-ui/app/services/flags/flags.ts
+++ b/polygerrit-ui/app/services/flags/flags.ts
@@ -20,5 +20,6 @@
   PUSH_NOTIFICATIONS_DEVELOPER = 'UiFeature__push_notifications_developer',
   PUSH_NOTIFICATIONS = 'UiFeature__push_notifications',
   ML_SUGGESTED_EDIT = 'UiFeature__ml_suggested_edit',
+  ML_SUGGESTED_EDIT_V2 = 'UiFeature__ml_suggested_edit_v2',
   REVISION_PARENTS_DATA = 'UiFeature__revision_parents_data',
 }