From a5e4730a5fa457028e2a0ca23a2564fcee31f7d9 Mon Sep 17 00:00:00 2001 From: pengx17 Date: Mon, 6 May 2024 13:34:43 +0000 Subject: [PATCH] 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) --- .../block-suite-editor/ai/event-source.ts | 12 +- .../block-suite-editor/ai/provider.tsx | 118 ++------ .../block-suite-editor/ai/request.ts | 11 +- .../block-suite-editor/ai/tracker.ts | 261 ++++++++++++++++++ 4 files changed, 309 insertions(+), 93 deletions(-) create mode 100644 packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/tracker.ts diff --git 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 index eec95bd0f0..18cf89db40 100644 --- 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; 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); diff --git 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 index ae16126822..22f32bffc1 100644 --- 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 = { - 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 = ( - id: T, - action: ( - ...options: Parameters - ) => ReturnType -) => { - 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(); } diff --git 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 index 014bb428d9..5f08c964bc 100644 --- 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; } diff --git 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 new file mode 100644 index 0000000000..604a589210 --- /dev/null +++ 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; + docId: string; + workspaceId: string; +}; + +type BlocksuiteActionEvent = Parameters< + Parameters[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); + } + }); +}