(core) Add telemetry for AI Assistant

Summary: Also fixes a few bugs with some telemetry events not being recorded.

Test Plan: Manual.

Reviewers: paulfitz

Reviewed By: paulfitz

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

app/server/lib/Assistance.ts

@@ -5,6 +5,7 @@
 import {AssistanceMessage, AssistanceRequest, AssistanceResponse} from 'app/common/AssistancePrompts';
 import {delay} from 'app/common/delay';
 import {DocAction} from 'app/common/DocActions';
+import {ActiveDoc} from 'app/server/lib/ActiveDoc';
 import {OptDocSession} from 'app/server/lib/DocSession';
 import log from 'app/server/lib/log';
 import fetch from 'node-fetch';
@@ -17,7 +18,7 @@ export const DEPS = { fetch, delayTime: 1000 };
  * An assistant can help a user do things with their document,
  * by interfacing with an external LLM endpoint.
  */
-export interface Assistant {
+interface Assistant {
   apply(session: OptDocSession, doc: AssistanceDoc, request: AssistanceRequest): Promise<AssistanceResponse>;
 }
 
@@ -25,7 +26,7 @@ export interface Assistant {
  * Document-related methods for use in the implementation of assistants.
  * Somewhat ad-hoc currently.
  */
-export interface AssistanceDoc {
+interface AssistanceDoc extends ActiveDoc {
   /**
    * Generate a particular prompt coded in the data engine for some reason.
    * It makes python code for some tables, and starts a function body with
@@ -117,10 +118,11 @@ export class OpenAIAssistant implements Assistant {
   public async apply(
     optSession: OptDocSession, doc: AssistanceDoc, request: AssistanceRequest): Promise<AssistanceResponse> {
     const messages = request.state?.messages || [];
+    const newMessages = [];
     const chatMode = this._chatMode;
     if (chatMode) {
       if (messages.length === 0) {
-        messages.push({
+        newMessages.push({
           role: 'system',
           content: 'You are a helpful assistant for a user of software called Grist. ' +
             'Below are one or more Python classes. ' +
@@ -141,7 +143,7 @@ export class OpenAIAssistant implements Assistant {
             await makeSchemaPromptV1(optSession, doc, request) +
             '\n```',
         });
-        messages.push({
+        newMessages.push({
           role: 'user', content: request.text,
         });
       } else {
@@ -150,21 +152,49 @@ export class OpenAIAssistant implements Assistant {
             messages.pop();
           }
         }
-        messages.push({
+        newMessages.push({
           role: 'user', content: request.text,
         });
       }
     } else {
       messages.length = 0;
-      messages.push({
+      newMessages.push({
         role: 'user', content: await makeSchemaPromptV1(optSession, doc, request),
       });
     }
+    messages.push(...newMessages);
+
+    const newMessagesStartIndex = messages.length - newMessages.length;
+    for (const [index, {role, content}] of newMessages.entries()) {
+      doc.logTelemetryEvent(optSession, 'assistantSend', {
+        full: {
+          conversationId: request.conversationId,
+          context: request.context,
+          prompt: {
+            index: newMessagesStartIndex + index,
+            role,
+            content,
+          },
+        },
+      });
+    }
+
     const completion: string = await this._getCompletion(messages);
     const response = await completionToResponse(doc, request, completion, completion);
     if (chatMode) {
       response.state = {messages};
     }
+    doc.logTelemetryEvent(optSession, 'assistantReceive', {
+      full: {
+        conversationId: request.conversationId,
+        context: request.context,
+        message: {
+          index: messages.length - 1,
+          content: completion,
+        },
+        suggestedFormula: (response.suggestedActions[0]?.[3] as any)?.formula,
+      },
+    });
     return response;
   }
 
@@ -307,7 +337,7 @@ export class HuggingFaceAssistant implements Assistant {
 /**
  * Test assistant that mimics ChatGPT and just returns the input.
  */
-export class EchoAssistant implements Assistant {
+class EchoAssistant implements Assistant {
   public async apply(sess: OptDocSession, doc: AssistanceDoc, request: AssistanceRequest): Promise<AssistanceResponse> {
     if (request.text === "ERROR") {
       throw new Error(`ERROR`);

app/server/lib/Telemetry.ts

@@ -12,6 +12,7 @@ import {
   TelemetryLevels,
   TelemetryMetadata,
   TelemetryMetadataByLevel,
+  TelemetryRetentionPeriod,
 } from 'app/common/Telemetry';
 import {TelemetryPrefsWithSources} from 'app/common/InstallAPI';
 import {Activation} from 'app/gen-server/entity/Activation';
@@ -44,19 +45,15 @@ const MAX_PENDING_FORWARD_EVENT_REQUESTS = 25;
  */
 export class Telemetry implements ITelemetry {
   private _activation: Activation | undefined;
-  private readonly _deploymentType = this._gristServer.getDeploymentType();
-
   private _telemetryPrefs: TelemetryPrefsWithSources | undefined;
-
+  private readonly _deploymentType = this._gristServer.getDeploymentType();
   private readonly _shouldForwardTelemetryEvents = this._deploymentType !== 'saas';
   private readonly _forwardTelemetryEventsUrl = process.env.GRIST_TELEMETRY_URL ||
     'https://telemetry.getgrist.com/api/telemetry';
   private _numPendingForwardEventRequests = 0;
-
-
   private readonly _logger = new LogMethods('Telemetry ', () => ({}));
-  private readonly _telemetryLogger = new LogMethods('Telemetry ', () => ({
-    eventType: 'telemetry',
+  private readonly _telemetryLogger = new LogMethods<string>('Telemetry ', (eventType) => ({
+    eventType,
   }));
 
   private _checkTelemetryEvent: TelemetryEventChecker | undefined;
@@ -153,21 +150,17 @@ export class Telemetry implements ITelemetry {
      */
     app.post('/api/telemetry', expressWrap(async (req, resp) => {
       const mreq = req as RequestWithLogin;
-      const event = stringParam(req.body.event, 'event', TelemetryEvents.values);
+      const event = stringParam(req.body.event, 'event', TelemetryEvents.values) as TelemetryEvent;
       if ('eventSource' in req.body.metadata) {
-        this._telemetryLogger.rawLog('info', null, event, {
-          eventName: event,
+        this._telemetryLogger.rawLog('info', getEventType(event), event, {
           ...(removeNullishKeys(req.body.metadata)),
+          eventName: event,
         });
       } else {
         try {
           this._assertTelemetryIsReady();
-          await this.logEvent(event as TelemetryEvent, merge(
+          await this.logEvent(event, merge(
             {
-              limited: {
-                eventSource: `grist-${this._deploymentType}`,
-                ...(this._deploymentType !== 'saas' ? {installationId: this._activation!.id} : {}),
-              },
               full: {
                 userId: mreq.userId,
                 altSessionId: mreq.altSessionId,
@@ -219,10 +212,11 @@ export class Telemetry implements ITelemetry {
     event: TelemetryEvent,
     metadata?: TelemetryMetadata
   ) {
-    this._telemetryLogger.rawLog('info', null, event, {
+    this._telemetryLogger.rawLog('info', getEventType(event), event, {
+      ...metadata,
       eventName: event,
       eventSource: `grist-${this._deploymentType}`,
-      ...metadata,
+      installationId: this._activation!.id,
     });
   }
 
@@ -238,7 +232,15 @@ export class Telemetry implements ITelemetry {
 
     try {
       this._numPendingForwardEventRequests += 1;
-      await this._doForwardEvent(JSON.stringify({event, metadata}));
+      await this._doForwardEvent(JSON.stringify({
+        event,
+        metadata: {
+          ...metadata,
+          eventName: event,
+          eventSource: `grist-${this._deploymentType}`,
+          installationId: this._activation!.id,
+        }
+      }));
     } catch (e) {
       this._logger.error(undefined, `failed to forward telemetry event ${event}`, e);
     } finally {
@@ -342,3 +344,15 @@ export function hashDigestKeys(metadata: TelemetryMetadata): TelemetryMetadata {
   });
   return filteredMetadata;
 }
+
+type TelemetryEventType = 'telemetry' | 'telemetry-short-retention';
+
+const EventTypeByRetentionPeriod: Record<TelemetryRetentionPeriod, TelemetryEventType> = {
+  indefinitely: 'telemetry',
+  short: 'telemetry-short-retention',
+};
+
+function getEventType(event: TelemetryEvent) {
+  const {retentionPeriod} = TelemetryContracts[event];
+  return EventTypeByRetentionPeriod[retentionPeriod];
+}