(core) Porting back AI formula backend

Summary: This is a backend part for the formula AI.

Test Plan: New tests

Reviewers: paulfitz

Reviewed By: paulfitz

Subscribers: cyprien

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

app/server/lib/ActiveDoc.ts

@@ -14,6 +14,7 @@ import {
 } from 'app/common/ActionBundle';
 import {ActionGroup, MinimalActionGroup} from 'app/common/ActionGroup';
 import {ActionSummary} from "app/common/ActionSummary";
+import {Prompt, Suggestion} from "app/common/AssistancePrompts";
 import {
   AclResources,
   AclTableDescription,
@@ -80,6 +81,7 @@ import {parseUserAction} from 'app/common/ValueParser';
 import {ParseOptions} from 'app/plugin/FileParserAPI';
 import {AccessTokenOptions, AccessTokenResult, GristDocAPI} from 'app/plugin/GristAPI';
 import {compileAclFormula} from 'app/server/lib/ACLFormula';
+import {sendForCompletion} from 'app/server/lib/Assistance';
 import {Authorizer} from 'app/server/lib/Authorizer';
 import {checksumFile} from 'app/server/lib/checksumFile';
 import {Client} from 'app/server/lib/Client';
@@ -1313,6 +1315,24 @@ export class ActiveDoc extends EventEmitter {
     return this._pyCall('autocomplete', txt, tableId, columnId, rowId, user.toJSON());
   }
 
+  public async getAssistance(docSession: DocSession, userPrompt: Prompt): Promise<Suggestion> {
+    // Making a prompt can leak names of tables and columns.
+    if (!await this._granularAccess.canScanData(docSession)) {
+      throw new Error("Permission denied");
+    }
+    await this.waitForInitialization();
+    const { tableId, colId, description } = userPrompt;
+    const prompt = await this._pyCall('get_formula_prompt', tableId, colId, description);
+    this._log.debug(docSession, 'getAssistance prompt', {prompt});
+    const completion = await sendForCompletion(prompt);
+    this._log.debug(docSession, 'getAssistance completion', {completion});
+    const formula = await this._pyCall('convert_formula_completion', completion);
+    const action: DocAction = ["ModifyColumn", tableId, colId, {formula}];
+    return {
+      suggestedActions: [action],
+    };
+  }
+
   public fetchURL(docSession: DocSession, url: string, options?: FetchUrlOptions): Promise<UploadResult> {
     return fetchURL(url, this.makeAccessId(docSession.authorizer.getUserId()), options);
   }

app/server/lib/Assistance.ts

@@ -0,0 +1,110 @@
+/**
+ * Module with functions used for AI formula assistance.
+ */
+
+import {delay} from 'app/common/delay';
+import log from 'app/server/lib/log';
+import fetch, { Response as FetchResponse} from 'node-fetch';
+
+
+export async function sendForCompletion(prompt: string): Promise<string> {
+  let completion: string|null = null;
+  if (process.env.OPENAI_API_KEY) {
+    completion = await sendForCompletionOpenAI(prompt);
+  }
+  if (process.env.HUGGINGFACE_API_KEY) {
+    completion = await sendForCompletionHuggingFace(prompt);
+  }
+  if (completion === null) {
+    throw new Error("Please set OPENAI_API_KEY or HUGGINGFACE_API_KEY (and optionally COMPLETION_MODEL)");
+  }
+  log.debug(`Received completion:`, {completion});
+  completion = completion.split(/\n {4}[^ ]/)[0];
+  return completion;
+}
+
+
+async function sendForCompletionOpenAI(prompt: string) {
+  const apiKey = process.env.OPENAI_API_KEY;
+  if (!apiKey) {
+    throw new Error("OPENAI_API_KEY not set");
+  }
+  const response = await fetch(
+    "https://api.openai.com/v1/completions",
+    {
+      method: "POST",
+      headers: {
+        "Authorization": `Bearer ${apiKey}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        prompt,
+        max_tokens: 150,
+        temperature: 0,
+        // COMPLETION_MODEL of `code-davinci-002` may be better if you have access to it.
+        model: process.env.COMPLETION_MODEL || "text-davinci-002",
+        stop: ["\n\n"],
+      }),
+    },
+  );
+  if (response.status !== 200) {
+    log.error(`OpenAI API returned ${response.status}: ${await response.text()}`);
+    throw new Error(`OpenAI API returned status ${response.status}`);
+  }
+  const result = await response.json();
+  const completion = result.choices[0].text;
+  return completion;
+}
+
+async function sendForCompletionHuggingFace(prompt: string) {
+  const apiKey = process.env.HUGGINGFACE_API_KEY;
+  if (!apiKey) {
+    throw new Error("HUGGINGFACE_API_KEY not set");
+  }
+  // COMPLETION_MODEL values I've tried:
+  //   - codeparrot/codeparrot
+  //   - NinedayWang/PolyCoder-2.7B
+  //   - NovelAI/genji-python-6B
+  let completionUrl = process.env.COMPLETION_URL;
+  if (!completionUrl) {
+    if (process.env.COMPLETION_MODEL) {
+      completionUrl = `https://api-inference.huggingface.co/models/${process.env.COMPLETION_MODEL}`;
+    } else {
+      completionUrl = 'https://api-inference.huggingface.co/models/NovelAI/genji-python-6B';
+    }
+  }
+  let retries: number = 0;
+  let response!: FetchResponse;
+  while (retries++ < 3) {
+    response = await fetch(
+      completionUrl,
+      {
+        method: "POST",
+        headers: {
+          "Authorization": `Bearer ${apiKey}`,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          inputs: prompt,
+          parameters: {
+            return_full_text: false,
+            max_new_tokens: 50,
+          },
+        }),
+      },
+    );
+    if (response.status === 503) {
+      log.error(`Sleeping for 10s - HuggingFace API returned ${response.status}: ${await response.text()}`);
+      await delay(10000);
+      continue;
+    }
+  }
+  if (response.status !== 200) {
+    const text = await response.text();
+    log.error(`HuggingFace API returned ${response.status}: ${text}`);
+    throw new Error(`HuggingFace API returned status ${response.status}: ${text}`);
+  }
+  const result = await response.json();
+  const completion = result[0].generated_text;
+  return completion.split('\n\n')[0];
+}

app/server/lib/DocWorker.ts

@@ -110,6 +110,7 @@ export class DocWorker {
       applyUserActionsById:     activeDocMethod.bind(null, 'editors', 'applyUserActionsById'),
       findColFromValues:        activeDocMethod.bind(null, 'viewers', 'findColFromValues'),
       getFormulaError:          activeDocMethod.bind(null, 'viewers', 'getFormulaError'),
+      getAssistance:            activeDocMethod.bind(null, 'editors', 'getAssistance'),
       importFiles:              activeDocMethod.bind(null, 'editors', 'importFiles'),
       finishImportFiles:        activeDocMethod.bind(null, 'editors', 'finishImportFiles'),
       cancelImportFiles:        activeDocMethod.bind(null, 'editors', 'cancelImportFiles'),