From 05c15e4ec3fbf36a5c5203a7a84d61e5557eb3e8 Mon Sep 17 00:00:00 2001
From: George Gevoian <george@gevoian.com>
Date: Mon, 31 Jul 2023 00:21:24 -0700
Subject: [PATCH] (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 |  9 +++----
 app/common/AssistancePrompts.ts        |  1 +
 app/server/lib/Assistance.ts           | 34 ++++++++++++++++++++------
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/app/client/widgets/FormulaAssistant.ts b/app/client/widgets/FormulaAssistant.ts
index f0f5c322..df8defa5 100644
--- a/app/client/widgets/FormulaAssistant.ts
+++ b/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',
     };
diff --git a/app/common/AssistancePrompts.ts b/app/common/AssistancePrompts.ts
index 10a925b4..5eef931c 100644
--- a/app/common/AssistancePrompts.ts
+++ b/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.
diff --git a/app/server/lib/Assistance.ts b/app/server/lib/Assistance.ts
index f9af0806..0ece1d65 100644
--- a/app/server/lib/Assistance.ts
+++ b/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,
   };
 }