Jayden/AFFiNE-Mirror
0
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-07-15 17:16:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(core): ai session create (#11539)

Browse Source 
Close [BS-3079](https://linear.app/affine-design/issue/BS-3079).

- Separate the create session logic from the `createMessage function`.
- Ensure the session is created before executing any chat or actions.
- Convert the `AIActions` into asynchronous functions.
- Transfer the update prompt name logic to the chat action.
- Introduce a networkSearch field in `AITextActionOptions`.
- Eliminate the redundant `LAST_ROOT_SESSION_ID`.
This commit is contained in:
akumatus
2025-04-09 06:18:56 +00:00
parent 1e9561b46c
commit c9790ed854
14 changed files with 403 additions and 331 deletions
Download Patch File Download Diff File Expand all files Collapse all files
blocksuite/affine/model/src/elements/mindmap/mindmap.ts
+1 -1
View File
@@ -743,7 +743,7 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
const targetPos =
typeof targetXYWH === 'string' ? deserializeXYWH(targetXYWH) : targetXYWH;
const offsetX = targetPos[0] - x;
const offsetY = targetPos[1] - y + targetPos[3];
const offsetY = targetPos[1] - y;
this.surface.doc.transact(() => {
this.childElements.forEach(el => {
packages/frontend/core/src/blocksuite/ai/actions/doc-handler.ts
+1 -1
View File
@@ -108,7 +108,7 @@ function actionToStream<T extends keyof BlockSuitePresets.AIActions>(
workspaceId: host.doc.workspace.id,
} as Parameters<typeof action>[0];
// @ts-expect-error TODO(@Peng): maybe fix this
stream = action(options);
stream = await action(options);
if (!stream) return;
yield* stream;
},
packages/frontend/core/src/blocksuite/ai/actions/edgeless-handler.ts
+2 -2
View File
@@ -207,7 +207,7 @@ function actionToStream<T extends keyof BlockSuitePresets.AIActions>(
}
// @ts-expect-error TODO(@Peng): maybe fix this
stream = action(options);
stream = await action(options);
if (!stream) return;
yield* stream;
},
@@ -237,7 +237,7 @@ function actionToStream<T extends keyof BlockSuitePresets.AIActions>(
} as Parameters<typeof action>[0];
// @ts-expect-error TODO(@Peng): maybe fix this
stream = action(options);
stream = await action(options);
if (!stream) return;
yield* stream;
},
packages/frontend/core/src/blocksuite/ai/actions/page-response.ts
+10 -3
View File
@@ -5,7 +5,11 @@ import {
} from '@blocksuite/affine/blocks/surface';
import { fitContent } from '@blocksuite/affine/gfx/shape';
import { createTemplateJob } from '@blocksuite/affine/gfx/template';
import { Bound, getCommonBound } from '@blocksuite/affine/global/gfx';
import {
Bound,
getCommonBound,
type XYWH,
} from '@blocksuite/affine/global/gfx';
import type {
MindmapElementModel,
ShapeElementModel,
@@ -83,7 +87,10 @@ function responseToBrainstormMindmap(
});
// wait for mindmap xywh update
setTimeout(() => {
const frameBound = expandBound(mindmap.elementBound, PADDING);
const { x, y, w, h } = mindmap.elementBound;
const targetBound: XYWH = [x, y + h / 2 + PADDING - 15, w, h];
mindmap.moveTo(targetBound);
const frameBound = expandBound(new Bound(...targetBound), PADDING);
addSurfaceRefBlock(host, frameBound, place);
}, 0);
});
@@ -107,7 +114,7 @@ function responseToMakeItReal(host: EditorHost, ctx: AIContext, place: Place) {
const bound = getEdgelessContentBound(host);
const x = bound ? bound.x + bound.w + PADDING * 2 : 0;
const y = bound ? bound.y : 0;
const htmlBound = new Bound(x, y, width || 800, height || 600);
const htmlBound = new Bound(x, y + PADDING, width || 800, height || 600);
const html = preprocessHtml(aiPanel.answer);
host.doc.transact(() => {
host.doc.addBlock(
packages/frontend/core/src/blocksuite/ai/actions/types.ts
+42 -35
View File
@@ -13,6 +13,8 @@ import type { EditorHost } from '@blocksuite/affine/std';
import type { GfxModel } from '@blocksuite/affine/std/gfx';
import type { BlockModel } from '@blocksuite/affine/store';
import type { PromptKey } from '../provider/prompt';
export const translateLangs = [
'English',
'Spanish',
@@ -131,6 +133,7 @@ declare global {
interface ChatOptions extends AITextActionOptions {
sessionId?: string;
isRootSession?: boolean;
networkSearch?: boolean;
contexts?: {
docs: AIDocContextOption[];
files: AIFileContextOption[];
@@ -155,107 +158,107 @@ declare global {
interface AIActions {
// chat is a bit special because it's has a internally maintained session
chat<T extends ChatOptions>(options: T): AIActionTextResponse<T>;
chat<T extends ChatOptions>(options: T): Promise<AIActionTextResponse<T>>;
summary<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
improveWriting<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
improveGrammar<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
fixSpelling<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
createHeadings<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
makeLonger<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
makeShorter<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
continueWriting<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
checkCodeErrors<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
explainCode<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
writeArticle<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
writeTwitterPost<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
writePoem<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
writeBlogPost<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
brainstorm<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
writeOutline<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
explainImage<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
findActions<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
// mindmap
brainstormMindmap<T extends BrainstormMindMap>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
expandMindmap<T extends ExpandMindMap>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
// presentation
createSlides<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
// explain this
explain<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
// actions with variants
translate<T extends TranslateOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
changeTone<T extends ChangeToneOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
// make it real, image to text
makeItReal<T extends AIImageActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
createImage<T extends AIImageActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
processImage<T extends ProcessImageOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
filterImage<T extends FilterImageOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
generateCaption<T extends AITextActionOptions>(
options: T
): AIActionTextResponse<T>;
): Promise<AIActionTextResponse<T>>;
}
type AIDocsAndFilesContext = {
@@ -357,12 +360,16 @@ declare global {
>[];
};
interface CreateSessionOptions {
docId: string;
workspaceId: string;
promptName: PromptKey;
sessionId?: string;
retry?: boolean;
}
interface AISessionService {
createSession: (
workspaceId: string,
docId: string,
promptName?: string
) => Promise<string>;
createSession: (options: CreateSessionOptions) => Promise<string>;
getSessions: (
workspaceId: string,
docId?: string,
packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-messages.ts
+11 -13
View File
@@ -348,10 +348,10 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
}
retry = async () => {
const { doc } = this.host;
try {
const sessionId = await this.createSessionId();
if (!sessionId) return;
if (!AIProvider.actions.chat) return;
const abortController = new AbortController();
const messages = [...this.chatContextValue.messages];
@@ -362,7 +362,8 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
}
this.updateContext({ messages, status: 'loading', error: null });
const stream = AIProvider.actions.chat?.({
const { doc } = this.host;
const stream = await AIProvider.actions.chat({
sessionId,
retry: true,
docId: doc.id,
@@ -374,18 +375,15 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
control: 'chat-send',
isRootSession: true,
});
if (stream) {
this.updateContext({ abortController });
for await (const text of stream) {
const messages = [...this.chatContextValue.messages];
const last = messages[messages.length - 1] as ChatMessage;
last.content += text;
this.updateContext({ messages, status: 'transmitting' });
}
this.updateContext({ status: 'success' });
this.updateContext({ abortController });
for await (const text of stream) {
const messages = [...this.chatContextValue.messages];
const last = messages[messages.length - 1] as ChatMessage;
last.content += text;
this.updateContext({ messages, status: 'transmitting' });
}
this.updateContext({ status: 'success' });
} catch (error) {
this.updateContext({ status: 'error', error: error as AIError });
} finally {
packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
+5 -5
View File
@@ -141,7 +141,6 @@ export class ChatPanel extends SignalWatcher(
const history = histories?.find(history => history.sessionId === sessionId);
if (history) {
messages.push(...history.messages);
AIProvider.LAST_ROOT_SESSION_ID = history.sessionId;
}
this.chatContextValue = {
@@ -182,10 +181,11 @@ export class ChatPanel extends SignalWatcher(
if (this._sessionId) {
return this._sessionId;
}
this._sessionId = await AIProvider.session?.createSession(
this.doc.workspace.id,
this.doc.id
);
this._sessionId = await AIProvider.session?.createSession({
docId: this.doc.id,
workspaceId: this.doc.workspace.id,
promptName: 'Chat With AFFiNE AI',
});
return this._sessionId;
};
packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts
+14 -32
View File
@@ -25,7 +25,6 @@ import type {
} from '../ai-chat-chips/type';
import { isDocChip, isFileChip } from '../ai-chat-chips/utils';
import type { ChatMessage } from '../ai-chat-messages';
import { PROMPT_NAME_AFFINE_AI, PROMPT_NAME_NETWORK_SEARCH } from './const';
import type { AIChatInputContext, AINetworkSearchConfig } from './type';
const MaximumImageCount = 32;
@@ -255,7 +254,8 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
private get _isNetworkActive() {
return (
!!this.networkSearchConfig.visible.value &&
!!this.networkSearchConfig.enabled.value
!!this.networkSearchConfig.enabled.value &&
!this._isNetworkDisabled
);
}
@@ -274,22 +274,6 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
);
}
private _getPromptName() {
if (this._isNetworkDisabled) {
return PROMPT_NAME_AFFINE_AI;
}
return this._isNetworkActive
? PROMPT_NAME_NETWORK_SEARCH
: PROMPT_NAME_AFFINE_AI;
}
private async _updatePromptName(promptName: string) {
const sessionId = await this.createSessionId();
if (sessionId && AIProvider.session) {
await AIProvider.session.updateSession(sessionId, promptName);
}
}
override connectedCallback() {
super.connectedCallback();
this._disposables.add(
@@ -313,7 +297,7 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
const { images, status } = this.chatContextValue;
const hasImages = images.length > 0;
const maxHeight = hasImages ? 272 + 2 : 200 + 2;
const uploadDisabled = this._isNetworkActive && !this._isNetworkDisabled;
const uploadDisabled = this._isNetworkActive;
return html` <div
class="chat-panel-input"
data-if-focused=${this.focused}
@@ -380,9 +364,7 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
data-testid="chat-network-search"
aria-disabled=${this._isNetworkDisabled}
data-active=${this._isNetworkActive}
@click=${this._isNetworkDisabled
? undefined
: this._toggleNetworkSearch}
@click=${this._toggleNetworkSearch}
@pointerdown=${stopPropagation}
>
${PublishIcon()}
@@ -473,6 +455,9 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
e.preventDefault();
e.stopPropagation();
if (this._isNetworkDisabled) {
return;
}
const enable = this.networkSearchConfig.enabled.value;
this.networkSearchConfig.setEnabled(!enable);
};
@@ -514,13 +499,12 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
};
send = async (text: string) => {
const { status, markdown, images } = this.chatContextValue;
if (status === 'loading' || status === 'transmitting') return;
if (!text) return;
if (!AIProvider.actions.chat) return;
try {
const promptName = this._getPromptName();
const { status, markdown, images } = this.chatContextValue;
if (status === 'loading' || status === 'transmitting') return;
if (!text) return;
if (!AIProvider.actions.chat) return;
const abortController = new AbortController();
this.updateContext({
images: [],
@@ -538,16 +522,13 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
// optimistic update messages
await this._preUpdateMessages(userInput, attachments);
// must update prompt name after local chat message is updated
// otherwise, the unauthorized error can not be rendered properly
await this._updatePromptName(promptName);
const sessionId = await this.createSessionId();
const contexts = await this._getMatchedContexts(userInput);
if (abortController.signal.aborted) {
return;
}
const stream = AIProvider.actions.chat({
const stream = await AIProvider.actions.chat({
sessionId,
input: userInput,
contexts,
@@ -560,6 +541,7 @@ export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
isRootSession: this.isRootSession,
where: this.trackOptions.where,
control: this.trackOptions.control,
networkSearch: this._isNetworkActive,
});
for await (const text of stream) {
packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/const.ts
-2
View File
@@ -1,2 +0,0 @@
export const PROMPT_NAME_AFFINE_AI = 'Chat With AFFiNE AI';
export const PROMPT_NAME_NETWORK_SEARCH = 'Search With AFFiNE AI';
packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/index.ts
-1
View File
@@ -1,3 +1,2 @@
export * from './ai-chat-input';
export * from './const';
export * from './type';
packages/frontend/core/src/blocksuite/ai/peek-view/chat-block-peek-view.ts
+16 -24
View File
@@ -320,16 +320,12 @@ export class AIChatBlockPeekView extends LitElement {
* Retry the last chat message
*/
retry = async () => {
const { doc } = this.host;
const { _forkBlockId, _forkSessionId } = this;
if (!_forkBlockId || !_forkSessionId) {
return;
}
let content = '';
try {
const abortController = new AbortController();
const { _forkBlockId, _forkSessionId } = this;
if (!_forkBlockId || !_forkSessionId) return;
if (!AIProvider.actions.chat) return;
const abortController = new AbortController();
const messages = [...this.chatContext.messages];
const last = messages[messages.length - 1];
if ('content' in last) {
@@ -339,7 +335,8 @@ export class AIChatBlockPeekView extends LitElement {
}
this.updateContext({ messages, status: 'loading', error: null });
const stream = AIProvider.actions.chat?.({
const { doc } = this.host;
const stream = await AIProvider.actions.chat({
sessionId: _forkSessionId,
retry: true,
docId: doc.id,
@@ -351,26 +348,21 @@ export class AIChatBlockPeekView extends LitElement {
control: 'chat-send',
});
if (stream) {
this.updateContext({ abortController });
for await (const text of stream) {
const messages = [...this.chatContext.messages];
const last = messages[messages.length - 1] as ChatMessage;
last.content += text;
this.updateContext({ messages, status: 'transmitting' });
content += text;
}
this.updateContext({ status: 'success' });
this.updateContext({ abortController });
for await (const text of stream) {
const messages = [...this.chatContext.messages];
const last = messages[messages.length - 1] as ChatMessage;
last.content += text;
this.updateContext({ messages, status: 'transmitting' });
}
this.updateContext({ status: 'success' });
// Update new chat block messages if there are contents returned from AI
await this.updateChatBlockMessages();
} catch (error) {
this.updateContext({ status: 'error', error: error as AIError });
} finally {
this.updateContext({ abortController: null });
if (content) {
// Update new chat block messages if there are contents returned from AI
await this.updateChatBlockMessages();
}
}
};
packages/frontend/core/src/blocksuite/ai/provider/ai-provider.ts
+5 -8
View File
@@ -100,8 +100,6 @@ export class AIProvider {
static LAST_ACTION_SESSIONID = '';
static LAST_ROOT_SESSION_ID = '';
static MAX_LOCAL_HISTORY = 10;
private readonly actions: Partial<BlockSuitePresets.AIActions> = {};
@@ -158,10 +156,10 @@ export class AIProvider {
id: T,
action: (
...options: Parameters<BlockSuitePresets.AIActions[T]>
) => ReturnType<BlockSuitePresets.AIActions[T]>
) => Promise<ReturnType<BlockSuitePresets.AIActions[T]>>
): void {
// @ts-expect-error TODO: maybe fix this
this.actions[id] = (
this.actions[id] = async (
...args: Parameters<BlockSuitePresets.AIActions[T]>
) => {
const options = args[0];
@@ -176,9 +174,8 @@ export class AIProvider {
this.actionHistory.shift();
}
// wrap the action with slot actions
const result: BlockSuitePresets.TextStream | Promise<string> = action(
...args
);
const result: BlockSuitePresets.TextStream | Promise<string> =
await action(...args);
const isTextStream = (
m: BlockSuitePresets.TextStream | Promise<string>
): m is BlockSuitePresets.TextStream =>
@@ -315,7 +312,7 @@ export class AIProvider {
id: T,
action: (
...options: Parameters<BlockSuitePresets.AIActions[T]>
) => ReturnType<BlockSuitePresets.AIActions[T]>
) => Promise<ReturnType<BlockSuitePresets.AIActions[T]>>
): void;
static provide(id: unknown, action: unknown) {
packages/frontend/core/src/blocksuite/ai/provider/request.ts
+36 -99
View File
@@ -3,16 +3,12 @@ import { partition } from 'lodash-es';
import { AIProvider } from './ai-provider';
import type { CopilotClient } from './copilot-client';
import { delay, toTextStream } from './event-source';
import type { PromptKey } from './prompt';
const TIMEOUT = 50000;
export type TextToTextOptions = {
client: CopilotClient;
docId: string;
workspaceId: string;
promptName?: PromptKey;
sessionId?: string | Promise<string>;
sessionId: string;
content?: string;
attachments?: (string | Blob | File)[];
params?: Record<string, any>;
@@ -61,30 +57,22 @@ async function resizeImage(blob: Blob | File): Promise<Blob | null> {
return null;
}
async function createSessionMessage({
interface CreateMessageOptions {
client: CopilotClient;
sessionId: string;
content?: string;
attachments?: (string | Blob | File)[];
params?: Record<string, any>;
}
async function createMessage({
client,
docId,
workspaceId,
promptName = 'Chat With AFFiNE AI',
sessionId,
content,
sessionId: providedSessionId,
attachments,
params,
}: TextToTextOptions): Promise<{
sessionId: string;
messageId: string;
}> {
if (!promptName && !providedSessionId) {
throw new Error('promptName or sessionId is required');
}
}: CreateMessageOptions): Promise<string> {
const hasAttachments = attachments && attachments.length > 0;
const sessionId = await (providedSessionId ??
client.createSession({
workspaceId,
docId,
promptName,
}));
const options: Parameters<CopilotClient['createMessage']>[0] = {
sessionId,
content,
@@ -110,67 +98,44 @@ async function createSessionMessage({
).filter(Boolean) as File[];
}
const messageId = await client.createMessage(options);
return {
messageId,
sessionId,
};
return await client.createMessage(options);
}
export function textToText({
client,
docId,
workspaceId,
promptName,
sessionId,
content,
attachments,
params,
sessionId,
stream,
signal,
timeout = TIMEOUT,
retry = false,
workflow = false,
isRootSession = false,
postfix,
}: TextToTextOptions) {
let _sessionId: string;
let _messageId: string | undefined;
let messageId: string | undefined;
if (stream) {
return {
[Symbol.asyncIterator]: async function* () {
if (retry) {
const retrySessionId =
(await sessionId) ?? AIProvider.LAST_ACTION_SESSIONID;
_sessionId = retrySessionId;
_messageId = undefined;
} else {
const message = await createSessionMessage({
if (!retry) {
messageId = await createMessage({
client,
docId,
workspaceId,
promptName,
sessionId,
content,
attachments,
params,
sessionId,
});
_sessionId = message.sessionId;
_messageId = message.messageId;
}
const eventSource = client.chatTextStream(
{
sessionId: _sessionId,
messageId: _messageId,
sessionId,
messageId,
},
workflow ? 'workflow' : undefined
);
AIProvider.LAST_ACTION_SESSIONID = _sessionId;
if (isRootSession) {
AIProvider.LAST_ROOT_SESSION_ID = _sessionId;
}
AIProvider.LAST_ACTION_SESSIONID = sessionId;
if (signal) {
if (signal.aborted) {
@@ -212,34 +177,20 @@ export function textToText({
})
: null,
(async function () {
if (retry) {
const retrySessionId =
(await sessionId) ?? AIProvider.LAST_ACTION_SESSIONID;
_sessionId = retrySessionId;
_messageId = undefined;
} else {
const message = await createSessionMessage({
if (!retry) {
messageId = await createMessage({
client,
docId,
workspaceId,
promptName,
sessionId,
content,
attachments,
params,
sessionId,
});
_sessionId = message.sessionId;
_messageId = message.messageId;
}
AIProvider.LAST_ACTION_SESSIONID = _sessionId;
if (isRootSession) {
AIProvider.LAST_ROOT_SESSION_ID = _sessionId;
}
AIProvider.LAST_ACTION_SESSIONID = sessionId;
return client.chatText({
sessionId: _sessionId,
messageId: _messageId,
sessionId,
messageId,
});
})(),
]);
@@ -248,50 +199,36 @@ export function textToText({
// Only one image is currently being processed
export function toImage({
docId,
workspaceId,
promptName,
content,
sessionId,
attachments,
params,
seed,
sessionId,
signal,
timeout = TIMEOUT,
retry = false,
workflow = false,
client,
}: ToImageOptions) {
let _sessionId: string;
let _messageId: string | undefined;
let messageId: string | undefined;
return {
[Symbol.asyncIterator]: async function* () {
if (retry) {
const retrySessionId =
(await sessionId) ?? AIProvider.LAST_ACTION_SESSIONID;
_sessionId = retrySessionId;
_messageId = undefined;
} else {
const { messageId, sessionId } = await createSessionMessage({
docId,
workspaceId,
promptName,
if (!retry) {
messageId = await createMessage({
client,
sessionId,
content,
attachments,
params,
client,
});
_sessionId = sessionId;
_messageId = messageId;
}
const eventSource = client.imagesStream(
_sessionId,
_messageId,
sessionId,
messageId,
seed,
workflow ? 'workflow' : undefined
);
AIProvider.LAST_ACTION_SESSIONID = _sessionId;
AIProvider.LAST_ACTION_SESSIONID = sessionId;
for await (const event of toTextStream(eventSource, {
timeout,
packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
+260 -105
View File
@@ -14,7 +14,7 @@ import type { PromptKey } from './prompt';
import { textToText, toImage } from './request';
import { setupTracker } from './tracker';
const filterStyleToPromptName = new Map(
const filterStyleToPromptName = new Map<string, PromptKey>(
Object.entries({
'Clay style': 'workflow:image-clay',
'Pixel style': 'workflow:image-pixel',
@@ -23,7 +23,7 @@ const filterStyleToPromptName = new Map(
})
);
const processTypeToPromptName = new Map(
const processTypeToPromptName = new Map<string, PromptKey>(
Object.entries({
Clearer: 'debug:action:fal-upscaler',
'Remove background': 'debug:action:fal-remove-bg',
@@ -35,31 +35,78 @@ export function setupAIProvider(
client: CopilotClient,
globalDialogService: GlobalDialogService
) {
async function createSession({
workspaceId,
docId,
promptName,
sessionId,
retry,
}: {
workspaceId: string;
docId: string;
promptName: PromptKey;
sessionId?: string;
retry?: boolean;
}) {
if (sessionId) return sessionId;
if (retry) return AIProvider.LAST_ACTION_SESSIONID;
return client.createSession({
workspaceId,
docId,
promptName,
});
}
//#region actions
AIProvider.provide('chat', options => {
const { input, contexts, ...rest } = options;
AIProvider.provide('chat', async options => {
const { input, contexts, attachments, networkSearch, retry } = options;
const disableSearch =
!!contexts?.files.length ||
!!contexts?.docs.length ||
!!attachments?.length;
const promptName =
networkSearch && !disableSearch
? 'Search With AFFiNE AI'
: 'Chat With AFFiNE AI';
const sessionId = await createSession({
promptName,
...options,
});
if (!retry) {
await AIProvider.session?.updateSession(sessionId, promptName);
}
return textToText({
...rest,
...options,
client,
sessionId,
content: input,
params: contexts,
});
});
AIProvider.provide('summary', options => {
AIProvider.provide('summary', async options => {
const sessionId = await createSession({
promptName: 'Summary',
...options,
});
return textToText({
...options,
client,
sessionId,
content: options.input,
promptName: 'Summary',
});
});
AIProvider.provide('translate', options => {
AIProvider.provide('translate', async options => {
const sessionId = await createSession({
promptName: 'Translate to',
...options,
});
return textToText({
...options,
client,
promptName: 'Translate to',
sessionId,
content: options.input,
params: {
language: options.lang,
@@ -67,200 +114,280 @@ export function setupAIProvider(
});
});
AIProvider.provide('changeTone', options => {
AIProvider.provide('changeTone', async options => {
const sessionId = await createSession({
promptName: 'Change tone to',
...options,
});
return textToText({
...options,
client,
sessionId,
params: {
tone: options.tone.toLowerCase(),
},
content: options.input,
promptName: 'Change tone to',
});
});
AIProvider.provide('improveWriting', options => {
return textToText({
...options,
client,
content: options.input,
AIProvider.provide('improveWriting', async options => {
const sessionId = await createSession({
promptName: 'Improve writing for it',
...options,
});
});
AIProvider.provide('improveGrammar', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('improveGrammar', async options => {
const sessionId = await createSession({
promptName: 'Improve grammar for it',
...options,
});
});
AIProvider.provide('fixSpelling', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('fixSpelling', async options => {
const sessionId = await createSession({
promptName: 'Fix spelling for it',
...options,
});
});
AIProvider.provide('createHeadings', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('createHeadings', async options => {
const sessionId = await createSession({
promptName: 'Create headings',
...options,
});
});
AIProvider.provide('makeLonger', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('makeLonger', async options => {
const sessionId = await createSession({
promptName: 'Make it longer',
...options,
});
});
AIProvider.provide('makeShorter', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('makeShorter', async options => {
const sessionId = await createSession({
promptName: 'Make it shorter',
...options,
});
});
AIProvider.provide('checkCodeErrors', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('checkCodeErrors', async options => {
const sessionId = await createSession({
promptName: 'Check code error',
...options,
});
});
AIProvider.provide('explainCode', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('explainCode', async options => {
const sessionId = await createSession({
promptName: 'Explain this code',
...options,
});
});
AIProvider.provide('writeArticle', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('writeArticle', async options => {
const sessionId = await createSession({
promptName: 'Write an article about this',
...options,
});
});
AIProvider.provide('writeTwitterPost', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('writeTwitterPost', async options => {
const sessionId = await createSession({
promptName: 'Write a twitter about this',
...options,
});
});
AIProvider.provide('writePoem', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('writePoem', async options => {
const sessionId = await createSession({
promptName: 'Write a poem about this',
...options,
});
});
AIProvider.provide('writeOutline', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('writeOutline', async options => {
const sessionId = await createSession({
promptName: 'Write outline',
...options,
});
});
AIProvider.provide('writeBlogPost', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('writeBlogPost', async options => {
const sessionId = await createSession({
promptName: 'Write a blog post about this',
...options,
});
});
AIProvider.provide('brainstorm', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('brainstorm', async options => {
const sessionId = await createSession({
promptName: 'Brainstorm ideas about this',
...options,
});
});
AIProvider.provide('findActions', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('findActions', async options => {
const sessionId = await createSession({
promptName: 'Find action items from it',
...options,
});
});
AIProvider.provide('brainstormMindmap', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('brainstormMindmap', async options => {
const sessionId = await createSession({
promptName: 'workflow:brainstorm',
...options,
});
return textToText({
...options,
client,
sessionId,
content: options.input,
// 3 minutes
timeout: 180000,
workflow: true,
});
});
AIProvider.provide('expandMindmap', options => {
AIProvider.provide('expandMindmap', async options => {
if (!options.input) {
throw new Error('expandMindmap action requires input');
}
const sessionId = await createSession({
promptName: 'Expand mind map',
...options,
});
return textToText({
...options,
client,
sessionId,
params: {
mindmap: options.mindmap,
node: options.input,
},
content: options.input,
promptName: 'Expand mind map',
});
});
AIProvider.provide('explain', options => {
return textToText({
...options,
client,
content: options.input,
AIProvider.provide('explain', async options => {
const sessionId = await createSession({
promptName: 'Explain this',
...options,
});
});
AIProvider.provide('explainImage', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
promptName: 'Explain this image',
});
});
AIProvider.provide('makeItReal', options => {
AIProvider.provide('explainImage', async options => {
const sessionId = await createSession({
promptName: 'Explain this image',
...options,
});
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('makeItReal', async options => {
let promptName: PromptKey = 'Make it real';
let content = options.input || '';
@@ -275,15 +402,20 @@ Here are our design notes:\n ${content}.`;
Could you make a new website based on these notes and send back just the html file?`;
}
const sessionId = await createSession({
promptName,
...options,
});
return textToText({
...options,
client,
sessionId,
content,
promptName,
});
});
AIProvider.provide('createSlides', options => {
AIProvider.provide('createSlides', async options => {
const SlideSchema = z.object({
page: z.number(),
type: z.enum(['name', 'title', 'content']),
@@ -320,11 +452,15 @@ Could you make a new website based on these notes and send back just the html fi
})
.join('\n');
};
const sessionId = await createSession({
promptName: 'workflow:presentation',
...options,
});
return textToText({
...options,
client,
sessionId,
content: options.input,
promptName: 'workflow:presentation',
// 3 minutes
timeout: 180000,
workflow: true,
@@ -332,79 +468,98 @@ Could you make a new website based on these notes and send back just the html fi
});
});
AIProvider.provide('createImage', options => {
AIProvider.provide('createImage', async options => {
// test to image
let promptName: PromptKey = 'debug:action:dalle3';
// image to image
if (options.attachments?.length) {
promptName = 'debug:action:fal-sd15';
}
const sessionId = await createSession({
promptName,
...options,
});
return toImage({
...options,
client,
sessionId,
content: options.input,
promptName,
});
});
AIProvider.provide('filterImage', options => {
AIProvider.provide('filterImage', async options => {
// test to image
const promptName = filterStyleToPromptName.get(options.style as string);
const promptName: PromptKey | undefined = filterStyleToPromptName.get(
options.style
);
if (!promptName) {
throw new Error('filterImage requires a promptName');
}
const sessionId = await createSession({
promptName,
...options,
});
return toImage({
...options,
client,
sessionId,
content: options.input,
timeout: 180000,
promptName: promptName as PromptKey,
workflow: !!promptName?.startsWith('workflow:'),
});
});
AIProvider.provide('processImage', options => {
AIProvider.provide('processImage', async options => {
// test to image
const promptName = processTypeToPromptName.get(
options.type as string
) as PromptKey;
const promptName: PromptKey | undefined = processTypeToPromptName.get(
options.type
);
if (!promptName) {
throw new Error('processImage requires a promptName');
}
const sessionId = await createSession({
promptName,
...options,
});
return toImage({
...options,
client,
sessionId,
content: options.input,
timeout: 180000,
promptName,
});
});
AIProvider.provide('generateCaption', options => {
return textToText({
...options,
client,
content: options.input,
AIProvider.provide('generateCaption', async options => {
const sessionId = await createSession({
promptName: 'Generate a caption',
...options,
});
});
AIProvider.provide('continueWriting', options => {
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
AIProvider.provide('continueWriting', async options => {
const sessionId = await createSession({
promptName: 'Continue writing',
...options,
});
return textToText({
...options,
client,
sessionId,
content: options.input,
});
});
//#endregion
AIProvider.provide('session', {
createSession: async (
workspaceId: string,
docId: string,
promptName = 'Chat With AFFiNE AI'
) => {
return client.createSession({
workspaceId,
docId,
promptName,
});
},
createSession,
getSessions: async (
workspaceId: string,
docId?: string,