(core) Show tweaked formula in AI responses

Summary:
The formula that's used when the Apply button is clicked, and the formula that's
shown in responses from the Formula Assistant should now be the same. Previously, they
would differ slightly.

Test Plan: Server tests.

Reviewers: alexmojaki

Reviewed By: alexmojaki

Subscribers: alexmojaki

Differential Revision: https://phab.getgrist.com/D3977

app/client/widgets/FormulaAssistant.ts

@@ -635,25 +635,24 @@ export class FormulaAssistant extends Disposable {
     // Get the state of the chat from the column.
     const conversationId = this._chat.conversationId.get();
     const prevState = column.chatHistory.peek().get().state;
-    const {reply, suggestedActions, state} = await askAI(gristDoc, {
+    const {reply, suggestedActions, suggestedFormula, state} = await askAI(gristDoc, {
       conversationId,
       column,
       description,
       state: prevState,
     });
-    console.debug('suggestedActions', {suggestedActions, reply, state});
+    console.debug('received formula assistant response: ', {suggestedActions, suggestedFormula, reply, state});
     // If back-end is capable of conversation, keep its state.
     const chatHistoryNew = column.chatHistory.peek();
     const value = chatHistoryNew.get();
     value.state = state;
-    const formula = (suggestedActions[0]?.[3] as any)?.formula as string;
     // If model has a conversational skills (and maintains a history), we might get actually
     // some markdown text back, so we need to parse it.
-    const prettyMessage = state ? (reply || formula || '') : (formula || reply || '');
+    const prettyMessage = state ? (reply || suggestedFormula || '') : (suggestedFormula || reply || '');
     // Add it to the chat.
     return {
       message: prettyMessage,
-      formula,
+      formula: suggestedFormula,
       action: suggestedActions[0],
       sender: 'ai',
     };

app/common/AssistancePrompts.ts

@@ -50,6 +50,7 @@ export interface AssistanceRequest {
  */
 export interface AssistanceResponse {
   suggestedActions: DocAction[];
+  suggestedFormula?: string;
   state?: AssistanceState;
   // If the model can be trusted to issue a self-contained
   // markdown-friendly string, it can be included here.

app/server/lib/Assistance.ts

@@ -198,7 +198,14 @@ export class OpenAIAssistant implements Assistant {
 
     const userIdHash = getUserHash(optSession);
     const completion: string = await this._getCompletion(messages, userIdHash);
-    const response = await completionToResponse(doc, request, completion, completion);
+    const response = await completionToResponse(doc, request, completion);
+    if (response.suggestedFormula) {
+      // Show the tweaked version of the suggested formula to the user (i.e. the one that's
+      // copied when the Apply button is clicked).
+      response.reply = replaceMarkdownCode(completion, response.suggestedFormula);
+    } else {
+      response.reply = completion;
+    }
     response.state = {messages};
     doc.logTelemetryEvent(optSession, 'assistantReceive', {
       full: {
@@ -208,7 +215,7 @@ export class OpenAIAssistant implements Assistant {
           index: messages.length - 1,
           content: completion,
         },
-        suggestedFormula: (response.suggestedActions[0]?.[3] as any)?.formula,
+        suggestedFormula: response.suggestedFormula,
       },
     });
     return response;
@@ -407,6 +414,14 @@ export async function sendForCompletion(
   return await assistant.apply(optSession, doc, request);
 }
 
+/**
+ * Returns a new Markdown string with the contents of its first multi-line code block
+ * replaced with `replaceValue`.
+ */
+export function replaceMarkdownCode(markdown: string, replaceValue: string) {
+  return markdown.replace(/```\w*\n(.*)```/s, '```python\n' + replaceValue + '\n```');
+}
+
 async function makeSchemaPromptV1(session: OptDocSession, doc: AssistanceDoc, request: AssistanceRequest) {
   if (request.context.type !== 'formula') {
     throw new Error('makeSchemaPromptV1 only works for formulas');
@@ -418,23 +433,28 @@ async function makeSchemaPromptV1(session: OptDocSession, doc: AssistanceDoc, re
   });
 }
 
-async function completionToResponse(doc: AssistanceDoc, request: AssistanceRequest,
-                                    completion: string, reply?: string): Promise<AssistanceResponse> {
+async function completionToResponse(
+  doc: AssistanceDoc,
+  request: AssistanceRequest,
+  completion: string,
+  reply?: string
+): Promise<AssistanceResponse> {
   if (request.context.type !== 'formula') {
     throw new Error('completionToResponse only works for formulas');
   }
-  completion = await doc.assistanceFormulaTweak(completion);
+  const suggestedFormula = await doc.assistanceFormulaTweak(completion) || undefined;
   // Suggest an action only if the completion is non-empty (that is,
   // it actually looked like code).
-  const suggestedActions: DocAction[] = completion ? [[
+  const suggestedActions: DocAction[] = suggestedFormula ? [[
     "ModifyColumn",
     request.context.tableId,
     request.context.colId, {
-      formula: completion,
+      formula: suggestedFormula,
     }
   ]] : [];
   return {
     suggestedActions,
+    suggestedFormula,
     reply,
   };
 }