refactor: refine ai tracker (#6778)

upstream https://github.com/toeverything/blocksuite/pull/6966 Added a new solution that inspect on actions event stream and adapt them into the mixpanel format. ![image.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/T2klNLEk0wxLh4NRDzhk/2631f0dc-5626-45d5-bcaf-60987aec3c7e.png)
2026-02-12 12:28:42 +00:00 · 2024-05-06 13:34:43 +00:00
parent 1ac16a48bf
commit a5e4730a5f
4 changed files with 309 additions and 93 deletions
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/event-source.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/event-source.ts
@@ -13,6 +13,7 @@ type AffineTextStream = AsyncIterable<AffineTextEvent>;

 type toTextStreamOptions = {
  timeout?: number;
+  signal?: AbortSignal;
 };

 // todo: may need to extend the error type
@@ -28,7 +29,7 @@ const safeParseError = (data: string): { status: number } => {

 export function toTextStream(
  eventSource: EventSource,
-  { timeout }: toTextStreamOptions = {}
+  { timeout, signal }: toTextStreamOptions = {}
 ): AffineTextStream {
  return {
    [Symbol.asyncIterator]: async function* () {
@@ -73,14 +74,19 @@ export function toTextStream(
      });

      try {
-        while (eventSource.readyState !== EventSource.CLOSED) {
+        while (
+          eventSource.readyState !== EventSource.CLOSED &&
+          !signal?.aborted
+        ) {
          if (messageQueue.length === 0) {
            // Wait for the next message or timeout
            await (timeout
              ? Promise.race([
                  messagePromise,
                  delay(timeout).then(() => {
-                    throw new Error('Timeout');
+                    if (!signal?.aborted) {
+                      throw new Error('Timeout');
+                    }
                  }),
                ])
              : messagePromise);
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.tsx
@@ -1,6 +1,5 @@
 import { notify } from '@affine/component';
 import { authAtom, openSettingModalAtom } from '@affine/core/atoms';
-import { mixpanel } from '@affine/core/utils';
 import { getBaseUrl } from '@affine/graphql';
 import { Trans } from '@affine/i18n';
 import { UnauthorizedError } from '@blocksuite/blocks';
@@ -15,60 +14,7 @@ import {
  textToText,
  toImage,
 } from './request';
-
-type AIAction = keyof BlockSuitePresets.AIActions;
-
-const TRACKED_ACTIONS: Record<AIAction, boolean> = {
-  chat: true,
-  summary: true,
-  translate: true,
-  changeTone: true,
-  improveWriting: true,
-  improveGrammar: true,
-  fixSpelling: true,
-  createHeadings: true,
-  makeLonger: true,
-  makeShorter: true,
-  checkCodeErrors: true,
-  explainCode: true,
-  writeArticle: true,
-  writeTwitterPost: true,
-  writePoem: true,
-  writeOutline: true,
-  writeBlogPost: true,
-  brainstorm: true,
-  findActions: true,
-  brainstormMindmap: true,
-  explain: true,
-  explainImage: true,
-  makeItReal: true,
-  createSlides: true,
-  createImage: true,
-  expandMindmap: true,
-  continueWriting: true,
-};
-
-const provideAction = <T extends AIAction>(
-  id: T,
-  action: (
-    ...options: Parameters<BlockSuitePresets.AIActions[T]>
-  ) => ReturnType<BlockSuitePresets.AIActions[T]>
-) => {
-  if (TRACKED_ACTIONS[id]) {
-    const wrappedFn: typeof action = (opts, ...rest) => {
-      mixpanel.track('AI', {
-        resolve: id,
-        docId: opts.docId,
-        workspaceId: opts.workspaceId,
-      });
-      // @ts-expect-error - todo: add a middleware in blocksuite instead?
-      return action(opts, ...rest);
-    };
-    AIProvider.provide(id, wrappedFn);
-  } else {
-    AIProvider.provide(id, action);
-  }
-};
+import { setupTracker } from './tracker';

 export function setupAIProvider() {
  // a single workspace should have only a single chat session
@@ -104,7 +50,7 @@ export function setupAIProvider() {
  }

  //#region actions
-  provideAction('chat', options => {
+  AIProvider.provide('chat', options => {
    const sessionId = getChatSessionId(options.workspaceId, options.docId);
    return textToText({
      ...options,
@@ -113,7 +59,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('summary', options => {
+  AIProvider.provide('summary', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -121,7 +67,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('translate', options => {
+  AIProvider.provide('translate', options => {
    return textToText({
      ...options,
      promptName: 'Translate to',
@@ -132,7 +78,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('changeTone', options => {
+  AIProvider.provide('changeTone', options => {
    return textToText({
      ...options,
      params: {
@@ -143,7 +89,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('improveWriting', options => {
+  AIProvider.provide('improveWriting', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -151,7 +97,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('improveGrammar', options => {
+  AIProvider.provide('improveGrammar', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -159,7 +105,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('fixSpelling', options => {
+  AIProvider.provide('fixSpelling', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -167,7 +113,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('createHeadings', options => {
+  AIProvider.provide('createHeadings', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -175,7 +121,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('makeLonger', options => {
+  AIProvider.provide('makeLonger', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -183,7 +129,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('makeShorter', options => {
+  AIProvider.provide('makeShorter', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -191,7 +137,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('checkCodeErrors', options => {
+  AIProvider.provide('checkCodeErrors', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -199,7 +145,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('explainCode', options => {
+  AIProvider.provide('explainCode', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -207,7 +153,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('writeArticle', options => {
+  AIProvider.provide('writeArticle', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -215,7 +161,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('writeTwitterPost', options => {
+  AIProvider.provide('writeTwitterPost', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -223,7 +169,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('writePoem', options => {
+  AIProvider.provide('writePoem', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -231,7 +177,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('writeOutline', options => {
+  AIProvider.provide('writeOutline', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -239,7 +185,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('writeBlogPost', options => {
+  AIProvider.provide('writeBlogPost', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -247,7 +193,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('brainstorm', options => {
+  AIProvider.provide('brainstorm', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -255,7 +201,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('findActions', options => {
+  AIProvider.provide('findActions', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -263,7 +209,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('brainstormMindmap', options => {
+  AIProvider.provide('brainstormMindmap', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -271,7 +217,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('expandMindmap', options => {
+  AIProvider.provide('expandMindmap', options => {
    return textToText({
      ...options,
      params: {
@@ -283,7 +229,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('explain', options => {
+  AIProvider.provide('explain', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -291,7 +237,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('explainImage', options => {
+  AIProvider.provide('explainImage', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -299,7 +245,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('makeItReal', options => {
+  AIProvider.provide('makeItReal', options => {
    return textToText({
      ...options,
      promptName: 'Make it real',
@@ -309,7 +255,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('createSlides', options => {
+  AIProvider.provide('createSlides', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -317,7 +263,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('createImage', options => {
+  AIProvider.provide('createImage', options => {
    // test to image
    let promptName: PromptKey = 'debug:action:dalle3';
    // image to image
@@ -330,7 +276,7 @@ export function setupAIProvider() {
    });
  });

-  provideAction('continueWriting', options => {
+  AIProvider.provide('continueWriting', options => {
    return textToText({
      ...options,
      content: options.input,
@@ -384,9 +330,6 @@ export function setupAIProvider() {
  });

  AIProvider.slots.requestUpgradePlan.on(() => {
-    mixpanel.track('AI', {
-      action: 'requestUpgradePlan',
-    });
    getCurrentStore().set(openSettingModalAtom, {
      activeTab: 'billing',
      open: true,
@@ -394,9 +337,6 @@ export function setupAIProvider() {
  });

  AIProvider.slots.requestLogin.on(() => {
-    mixpanel.track('AI', {
-      action: 'requestLogin',
-    });
    getCurrentStore().set(authAtom, s => ({
      ...s,
      openModal: true,
@@ -410,4 +350,6 @@ export function setupAIProvider() {
      ),
    });
  });
+
+  setupTracker();
 }
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
@@ -137,7 +137,10 @@ export function textToText({
            eventSource.close();
          };
        }
-        for await (const event of toTextStream(eventSource, { timeout })) {
+        for await (const event of toTextStream(eventSource, {
+          timeout,
+          signal,
+        })) {
          if (event.type === 'message') {
            yield event.data;
          }
@@ -180,6 +183,7 @@ export function toImage({
  attachments,
  params,
  seed,
+  signal,
  timeout = TIMEOUT,
 }: ToImageOptions) {
  return {
@@ -194,7 +198,10 @@ export function toImage({
      });

      const eventSource = client.imagesStream(messageId, sessionId, seed);
-      for await (const event of toTextStream(eventSource, { timeout })) {
+      for await (const event of toTextStream(eventSource, {
+        timeout,
+        signal,
+      })) {
        if (event.type === 'attachment') {
          yield event.data;
        }
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/tracker.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/tracker.ts
@@ -0,0 +1,261 @@
+import { mixpanel } from '@affine/core/utils';
+import { DebugLogger } from '@affine/debug';
+import type { EditorHost } from '@blocksuite/block-std';
+import type { ElementModel } from '@blocksuite/blocks';
+import { AIProvider } from '@blocksuite/presets';
+import type { BlockModel } from '@blocksuite/store';
+import { lowerCase, omit } from 'lodash-es';
+
+type AIActionEventName =
+  | 'AI action invoked'
+  | 'AI action aborted'
+  | 'AI result discarded'
+  | 'AI result accepted';
+
+type AIActionEventProperties = {
+  page: 'doc-editor' | 'whiteboard-editor';
+  segment:
+    | 'AI action panel'
+    | 'right side bar'
+    | 'inline chat panel'
+    | 'AI result panel';
+  module:
+    | 'exit confirmation'
+    | 'AI action panel'
+    | 'AI chat panel'
+    | 'inline chat panel'
+    | 'AI result panel';
+  control:
+    | 'stop button'
+    | 'format toolbar'
+    | 'AI chat send button'
+    | 'paywall'
+    | 'policy wall'
+    | 'server error'
+    | 'insert'
+    | 'replace'
+    | 'discard'
+    | 'retry'
+    | 'add note'
+    | 'add page'
+    | 'continue in chat';
+  type:
+    | 'doc' // synced doc
+    | 'note' // note shape
+    | 'text'
+    | 'image'
+    | 'draw object'
+    | 'chatbox text'
+    | 'other';
+  category: string;
+  other: Record<string, unknown>;
+  docId: string;
+  workspaceId: string;
+};
+
+type BlocksuiteActionEvent = Parameters<
+  Parameters<typeof AIProvider.slots.actions.on>[0]
+>[0];
+
+const logger = new DebugLogger('affine:ai-tracker');
+
+const trackAction = ({
+  eventName,
+  properties,
+}: {
+  eventName: AIActionEventName;
+  properties: AIActionEventProperties;
+}) => {
+  logger.debug('trackAction', eventName, properties);
+  mixpanel.track(eventName, properties);
+};
+
+const inferPageMode = (host: EditorHost) => {
+  return host.querySelector('affine-page-root')
+    ? 'doc-editor'
+    : 'whiteboard-editor';
+};
+
+const defaultActionOptions = [
+  'stream',
+  'input',
+  'content',
+  'stream',
+  'attachments',
+  'signal',
+  'docId',
+  'workspaceId',
+  'host',
+  'models',
+  'control',
+  'where',
+  'seed',
+];
+
+function isElementModel(
+  model: BlockModel | ElementModel
+): model is ElementModel {
+  return !isBlockModel(model);
+}
+
+function isBlockModel(model: BlockModel | ElementModel): model is BlockModel {
+  return 'flavour' in model;
+}
+
+function inferObjectType(event: BlocksuiteActionEvent) {
+  const models: (BlockModel | ElementModel)[] | undefined =
+    event.options.models;
+  if (!models) {
+    if (event.action === 'chat') {
+      return 'chatbox text';
+    } else if (event.options.attachments?.length) {
+      return 'image';
+    } else {
+      return 'text';
+    }
+  } else if (models.every(isElementModel)) {
+    return 'draw object';
+  } else if (models.every(isBlockModel)) {
+    const flavour = models[0].flavour;
+    if (flavour === 'affine:note') {
+      return 'note';
+    } else if (
+      ['affine:paragraph', 'affine:list', 'affine:code'].includes(flavour)
+    ) {
+      return 'text';
+    } else if (flavour === 'affine:image') {
+      return 'image';
+    }
+  }
+  return 'other';
+}
+
+function inferSegment(
+  event: BlocksuiteActionEvent
+): AIActionEventProperties['segment'] {
+  if (event.action === 'chat') {
+    return 'inline chat panel';
+  } else if (event.event.startsWith('result:')) {
+    return 'AI result panel';
+  } else if (event.options.where === 'chat-panel') {
+    return 'right side bar';
+  } else {
+    return 'AI action panel';
+  }
+}
+
+function inferModule(
+  event: BlocksuiteActionEvent
+): AIActionEventProperties['module'] {
+  if (event.action === 'chat') {
+    return 'AI chat panel';
+  } else if (event.event === 'result:discard') {
+    return 'exit confirmation';
+  } else if (event.event.startsWith('result:')) {
+    return 'AI result panel';
+  } else if (event.options.where === 'chat-panel') {
+    return 'inline chat panel';
+  } else {
+    return 'AI action panel';
+  }
+}
+
+function inferEventName(
+  event: BlocksuiteActionEvent
+): AIActionEventName | null {
+  if (['result:discard', 'result:retry'].includes(event.event)) {
+    return 'AI result discarded';
+  } else if (event.event.startsWith('result:')) {
+    return 'AI result accepted';
+  } else if (event.event.startsWith('aborted:')) {
+    return 'AI action aborted';
+  } else if (event.event === 'started') {
+    return 'AI action invoked';
+  }
+  return null;
+}
+
+function inferControl(
+  event: BlocksuiteActionEvent
+): AIActionEventProperties['control'] {
+  if (event.event === 'aborted:stop') {
+    return 'stop button';
+  } else if (event.event === 'aborted:paywall') {
+    return 'paywall';
+  } else if (event.event === 'aborted:server-error') {
+    return 'server error';
+  } else if (event.options.control === 'chat-send') {
+    return 'AI chat send button';
+  } else if (event.event === 'result:add-note') {
+    return 'add note';
+  } else if (event.event === 'result:add-page') {
+    return 'add page';
+  } else if (event.event === 'result:continue-in-chat') {
+    return 'continue in chat';
+  } else if (event.event === 'result:insert') {
+    return 'insert';
+  } else if (event.event === 'result:replace') {
+    return 'replace';
+  } else if (event.event === 'result:discard') {
+    return 'discard';
+  } else if (event.event === 'result:retry') {
+    return 'retry';
+  } else {
+    return 'format toolbar';
+  }
+}
+
+const toTrackedOptions = (
+  event: BlocksuiteActionEvent
+): {
+  eventName: AIActionEventName;
+  properties: AIActionEventProperties;
+} | null => {
+  const eventName = inferEventName(event);
+
+  if (!eventName) return null;
+
+  const pageMode = inferPageMode(event.options.host);
+  const otherProperties = omit(event.options, defaultActionOptions);
+  const type = inferObjectType(event);
+  const segment = inferSegment(event);
+  const module = inferModule(event);
+  const control = inferControl(event);
+  const category = lowerCase(event.action);
+
+  return {
+    eventName,
+    properties: {
+      page: pageMode,
+      segment,
+      category,
+      module,
+      control,
+      type,
+      other: otherProperties,
+      docId: event.options.docId,
+      workspaceId: event.options.workspaceId,
+    },
+  };
+};
+
+export function setupTracker() {
+  AIProvider.slots.requestUpgradePlan.on(() => {
+    mixpanel.track('AI', {
+      action: 'requestUpgradePlan',
+    });
+  });
+
+  AIProvider.slots.requestLogin.on(() => {
+    mixpanel.track('AI', {
+      action: 'requestLogin',
+    });
+  });
+
+  AIProvider.slots.actions.on(event => {
+    const properties = toTrackedOptions(event);
+    if (properties) {
+      trackAction(properties);
+    }
+  });
+}