feat(core): add ai web-search and web-crawl tools ui components (#12854)

This commit is contained in:
Wu Yue 2025-06-19 19:15:46 +08:00 committed by GitHub
parent 2edc8e43e2
commit 3886babcf4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
20 changed files with 747 additions and 54 deletions

View File

@ -539,6 +539,20 @@ export class StreamObjectParser {
}
break;
}
case 'tool-result': {
const index = acc.findIndex(
item =>
item.type === 'tool-call' &&
item.toolCallId === curr.toolCallId &&
item.toolName === curr.toolName
);
if (index !== -1) {
acc[index] = curr;
} else {
acc.push(curr);
}
break;
}
default: {
acc.push(curr);
break;

View File

@ -10,6 +10,7 @@ import {
buildFinishConfig,
buildGeneratingConfig,
} from '../ai-panel';
import { StreamObjectSchema } from '../components/ai-chat-messages';
import { type AIItemGroupConfig } from '../components/ai-item/types';
import { type AIError, AIProvider } from '../provider';
import { reportResponse } from '../utils/action-reporter';
@ -21,8 +22,12 @@ import {
getSelections,
selectAboveBlocks,
} from '../utils/selection-utils';
import { mergeStreamObjects } from '../utils/stream-objects';
import type { AffineAIPanelWidget } from '../widgets/ai-panel/ai-panel';
import type { AINetworkSearchConfig } from '../widgets/ai-panel/type';
import type {
AIActionAnswer,
AINetworkSearchConfig,
} from '../widgets/ai-panel/type';
import { actionToAnswerRenderer } from './answer-renderer';
export function bindTextStream(
@ -32,13 +37,15 @@ export function bindTextStream(
finish,
signal,
}: {
update: (text: string) => void;
update: (answer: AIActionAnswer) => void;
finish: (state: 'success' | 'error' | 'aborted', err?: AIError) => void;
signal?: AbortSignal;
}
) {
(async () => {
let answer = '';
const answer: AIActionAnswer = {
content: '',
};
signal?.addEventListener('abort', () => {
finish('aborted');
reportResponse('aborted:stop');
@ -47,7 +54,19 @@ export function bindTextStream(
if (signal?.aborted) {
return;
}
answer += data;
try {
const parsed = StreamObjectSchema.safeParse(JSON.parse(data));
if (parsed.success) {
answer.streamObjects = mergeStreamObjects([
...(answer.streamObjects ?? []),
parsed.data,
]);
} else {
answer.content += data;
}
} catch {
answer.content += data;
}
update(answer);
}
finish('success');
@ -137,7 +156,7 @@ function actionToGenerateAnswer<T extends keyof BlockSuitePresets.AIActions>(
}: {
input: string;
signal?: AbortSignal;
update: (text: string) => void;
update: (answer: AIActionAnswer) => void;
finish: (state: 'success' | 'error' | 'aborted', err?: AIError) => void;
}) => {
const { selectedBlocks: blocks } = getSelections(host);

View File

@ -37,7 +37,10 @@ import {
getSelections,
} from '../utils/selection-utils';
import type { AffineAIPanelWidget } from '../widgets/ai-panel/ai-panel';
import type { AINetworkSearchConfig } from '../widgets/ai-panel/type';
import type {
AIActionAnswer,
AINetworkSearchConfig,
} from '../widgets/ai-panel/type';
import type { EdgelessCopilotWidget } from '../widgets/edgeless-copilot';
import { actionToAnswerRenderer } from './answer-renderer';
import { EXCLUDING_COPY_ACTIONS } from './consts';
@ -282,7 +285,7 @@ function actionToGeneration<T extends keyof BlockSuitePresets.AIActions>(
}: {
input: string;
signal?: AbortSignal;
update: (text: string) => void;
update: (answer: AIActionAnswer) => void;
finish: (state: 'success' | 'error' | 'aborted', err?: AIError) => void;
}) => {
if (!extract) {

View File

@ -21,8 +21,10 @@ import {
type ChatMessage,
isChatAction,
isChatMessage,
StreamObjectSchema,
} from '../components/ai-chat-messages';
import { type AIError, AIProvider, UnauthorizedError } from '../provider';
import { mergeStreamObjects } from '../utils/stream-objects';
import { type ChatContextValue } from './chat-context';
import { HISTORY_IMAGE_ACTIONS } from './const';
import { AIPreloadConfig } from './preload-config';
@ -387,7 +389,12 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
last.content = '';
last.createdAt = new Date().toISOString();
}
this.updateContext({ messages, status: 'loading', error: null });
this.updateContext({
messages,
status: 'loading',
error: null,
abortController,
});
const { store } = this.host;
const stream = await AIProvider.actions.chat({
@ -404,11 +411,23 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
reasoning: this._isReasoningActive,
webSearch: this._isNetworkActive,
});
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;
try {
const parsed = StreamObjectSchema.safeParse(JSON.parse(text));
if (parsed.success) {
last.streamObjects = mergeStreamObjects([
...(last.streamObjects ?? []),
parsed.data,
]);
} else {
last.content += text;
}
} catch {
last.content += text;
}
this.updateContext({ messages, status: 'transmitting' });
}

View File

@ -3,6 +3,7 @@ import '../content/rich-text';
import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
import { WithDisposable } from '@blocksuite/affine/global/lit';
import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
import { isInsidePageEditor } from '@blocksuite/affine/shared/utils';
import type { EditorHost } from '@blocksuite/affine/std';
import { ShadowlessElement } from '@blocksuite/affine/std';
@ -18,9 +19,11 @@ import {
type ChatMessage,
type ChatStatus,
isChatMessage,
type StreamObject,
} from '../../components/ai-chat-messages';
import { AIChatErrorRenderer } from '../../messages/error';
import { type AIError } from '../../provider';
import { mergeStreamContent } from '../../utils/stream-objects';
export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
static override styles = css`
@ -29,6 +32,20 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
font-size: var(--affine-font-xs);
font-weight: 400;
}
.reasoning-wrapper {
padding: 16px 20px;
margin: 8px 0;
border-radius: 8px;
background-color: rgba(0, 0, 0, 0.05);
}
.tool-wrapper {
padding: 12px;
margin: 8px 0;
border-radius: 8px;
border: 0.5px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
}
`;
@property({ attribute: false })
@ -78,33 +95,127 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
renderContent() {
const { host, item, isLast, status, error } = this;
const state = isLast
? status !== 'loading' && status !== 'transmitting'
? 'finished'
: 'generating'
: 'finished';
const { streamObjects, content } = item;
const shouldRenderError = isLast && status === 'error' && !!error;
return html`
${item.attachments
? html`<chat-content-images
.images=${item.attachments}
></chat-content-images>`
: nothing}
<chat-content-rich-text
.host=${host}
.text=${item.content}
.state=${state}
.extensions=${this.extensions}
.affineFeatureFlagService=${this.affineFeatureFlagService}
></chat-content-rich-text>
${this.renderImages()}
${streamObjects?.length
? this.renderStreamObjects(streamObjects)
: this.renderRichText(content)}
${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
${this.renderEditorActions()}
`;
}
renderEditorActions() {
private renderImages() {
const { item } = this;
if (!item.attachments) return nothing;
return html`<chat-content-images
.images=${item.attachments}
></chat-content-images>`;
}
private renderStreamObjects(answer: StreamObject[]) {
return html`<div>
${answer.map(data => {
switch (data.type) {
case 'text-delta':
return this.renderRichText(data.textDelta);
case 'reasoning':
return html`
<div class="reasoning-wrapper">
${this.renderRichText(data.textDelta)}
</div>
`;
case 'tool-call':
return this.renderToolCall(data);
case 'tool-result':
return this.renderToolResult(data);
default:
return nothing;
}
})}
</div>`;
}
private renderToolCall(streamObject: StreamObject) {
if (streamObject.type !== 'tool-call') {
return nothing;
}
switch (streamObject.toolName) {
case 'web_crawl_exa':
return html`
<web-crawl-tool
.data=${streamObject}
.host=${this.host}
></web-crawl-tool>
`;
case 'web_search_exa':
return html`
<web-search-tool
.data=${streamObject}
.host=${this.host}
></web-search-tool>
`;
default:
return html`
<div class="tool-wrapper">
${streamObject.toolName} tool calling...
</div>
`;
}
}
private renderToolResult(streamObject: StreamObject) {
if (streamObject.type !== 'tool-result') {
return nothing;
}
switch (streamObject.toolName) {
case 'web_crawl_exa':
return html`
<web-crawl-tool
.data=${streamObject}
.host=${this.host}
></web-crawl-tool>
`;
case 'web_search_exa':
return html`
<web-search-tool
.data=${streamObject}
.host=${this.host}
></web-search-tool>
`;
default:
return html`
<div class="tool-wrapper">
${streamObject.toolName} tool result...
</div>
`;
}
}
private renderRichText(text: string) {
const { host, isLast, status } = this;
const state = isLast
? status !== 'loading' && status !== 'transmitting'
? 'finished'
: 'generating'
: 'finished';
return html`<chat-content-rich-text
.host=${host}
.text=${text}
.state=${state}
.extensions=${this.extensions}
.affineFeatureFlagService=${this.affineFeatureFlagService}
></chat-content-rich-text>`;
}
private renderEditorActions() {
const { item, isLast, status } = this;
if (!isChatMessage(item) || item.role !== 'assistant') return nothing;
@ -118,7 +229,10 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
return nothing;
const { host } = this;
const { content, id: messageId } = item;
const { content, streamObjects, id: messageId } = item;
const markdown = streamObjects?.length
? mergeStreamContent(streamObjects)
: content;
const actions = isInsidePageEditor(host)
? PageEditorActions
@ -128,18 +242,18 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
<chat-copy-more
.host=${host}
.actions=${actions}
.content=${content}
.content=${markdown}
.isLast=${isLast}
.getSessionId=${this.getSessionId}
.messageId=${messageId}
.withMargin=${true}
.retry=${() => this.retry()}
></chat-copy-more>
${isLast && !!content
${isLast && !!markdown
? html`<chat-action-list
.actions=${actions}
.host=${host}
.content=${content}
.content=${markdown}
.getSessionId=${this.getSessionId}
.messageId=${messageId ?? undefined}
.withMargin=${true}

View File

@ -0,0 +1,101 @@
import { WithDisposable } from '@blocksuite/affine/global/lit';
import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
import { ShadowlessElement } from '@blocksuite/affine/std';
import { css, html, type TemplateResult } from 'lit';
import { property, state } from 'lit/decorators.js';
export class ToolCallCard extends WithDisposable(ShadowlessElement) {
static override styles = css`
.ai-tool-call-wrapper {
padding: 12px;
margin: 8px 0;
border-radius: 8px;
border: 0.5px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
.ai-tool-header {
display: flex;
justify-content: space-between;
align-items: center;
gap: 8px;
}
.ai-icon {
width: 24px;
height: 24px;
svg {
width: 24px;
height: 24px;
color: ${unsafeCSSVarV2('icon/activated')};
}
}
.ai-tool-name {
font-size: 14px;
font-weight: 500;
line-height: 24px;
margin-left: 0px;
margin-right: auto;
color: ${unsafeCSSVarV2('icon/activated')};
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
.loading-dots {
display: inline;
margin-left: 2px;
color: ${unsafeCSSVarV2('icon/activated')};
}
}
`;
@property({ attribute: false })
accessor name!: string;
@property({ attribute: false })
accessor icon!: TemplateResult<1>;
@state()
private accessor dotsText = '.';
private animationTimer?: number;
override connectedCallback() {
super.connectedCallback();
this.startDotsAnimation();
}
override disconnectedCallback() {
super.disconnectedCallback();
this.stopDotsAnimation();
}
private startDotsAnimation() {
let dotCount = 1;
this.animationTimer = window.setInterval(() => {
dotCount = (dotCount % 3) + 1;
this.dotsText = '.'.repeat(dotCount);
}, 750);
}
private stopDotsAnimation() {
if (this.animationTimer) {
clearInterval(this.animationTimer);
this.animationTimer = undefined;
}
}
protected override render() {
return html`
<div class="ai-tool-call-wrapper">
<div class="ai-tool-header">
<div class="ai-icon">${this.icon}</div>
<div class="ai-tool-name">
${this.name}<span class="loading-dots">${this.dotsText}</span>
</div>
</div>
</div>
`;
}
}

View File

@ -0,0 +1,160 @@
import { WithDisposable } from '@blocksuite/affine/global/lit';
import { ImageProxyService } from '@blocksuite/affine/shared/adapters';
import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
import { type EditorHost, ShadowlessElement } from '@blocksuite/affine/std';
import { ToggleDownIcon } from '@blocksuite/icons/lit';
import { css, html, nothing, type TemplateResult } from 'lit';
import { property } from 'lit/decorators.js';
interface ToolResult {
title?: string;
icon?: string | TemplateResult<1>;
content?: string;
}
export class ToolResultCard extends WithDisposable(ShadowlessElement) {
static override styles = css`
.ai-tool-wrapper {
padding: 12px;
margin: 8px 0;
border-radius: 8px;
border: 0.5px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
.ai-tool-header {
display: flex;
justify-content: space-between;
align-items: center;
gap: 8px;
}
.ai-icon {
width: 24px;
height: 24px;
svg {
width: 24px;
height: 24px;
color: ${unsafeCSSVarV2('icon/primary')};
}
}
.ai-tool-name {
font-size: 14px;
font-weight: 500;
line-height: 24px;
margin-left: 0px;
margin-right: auto;
color: ${unsafeCSSVarV2('icon/primary')};
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
.ai-tool-results {
display: flex;
flex-direction: column;
gap: 4px;
margin: 4px 2px 4px 12px;
padding-left: 20px;
border-left: 1px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
}
.result-header {
display: flex;
justify-content: space-between;
align-items: center;
}
.result-title {
font-size: 12px;
font-weight: 400;
line-height: 20px;
color: ${unsafeCSSVarV2('icon/primary')};
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
flex: 1;
}
.result-icon {
width: 24px;
height: 24px;
img {
width: 24px;
height: 24px;
border-radius: 100%;
border: 1px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
}
}
.result-content {
font-size: 12px;
line-height: 20px;
color: ${unsafeCSSVarV2('text/secondary')};
margin-top: 8px;
display: -webkit-box;
-webkit-line-clamp: 4;
-webkit-box-orient: vertical;
overflow: hidden;
text-overflow: ellipsis;
}
}
`;
@property({ attribute: false })
accessor host!: EditorHost;
@property({ attribute: false })
accessor name!: string;
@property({ attribute: false })
accessor icon!: TemplateResult<1> | string;
@property({ attribute: false })
accessor results!: ToolResult[];
protected override render() {
const imageProxyService = this.host.store.get(ImageProxyService);
return html`
<div class="ai-tool-wrapper">
<div class="ai-tool-header" data-type="result">
<div class="ai-icon">${this.icon}</div>
<div class="ai-tool-name">${this.name}</div>
<div class="ai-icon">${ToggleDownIcon()}</div>
</div>
<div class="ai-tool-results">
${this.results.map(
result => html`
<div>
<div class="result-header">
<div class="result-title">${result.title || 'Untitled'}</div>
${result.icon
? html`
<div class="result-icon">
${typeof result.icon === 'string'
? html`<img
src=${imageProxyService.buildUrl(result.icon)}
alt="icon"
@error=${(e: Event) => {
const target = e.target as HTMLImageElement;
target.style.display = 'none';
}}
/>`
: result.icon}
</div>
`
: nothing}
</div>
${result.content
? html` <div class="result-content">${result.content}</div> `
: nothing}
</div>
`
)}
</div>
</div>
`;
}
}

View File

@ -0,0 +1,79 @@
import { WithDisposable } from '@blocksuite/affine/global/lit';
import { type EditorHost, ShadowlessElement } from '@blocksuite/affine/std';
import { WebIcon } from '@blocksuite/icons/lit';
import { html, nothing } from 'lit';
import { property } from 'lit/decorators.js';
interface WebCrawlToolCall {
type: 'tool-call';
toolCallId: string;
toolName: string;
args: { url: string };
}
interface WebCrawlToolResult {
type: 'tool-result';
toolCallId: string;
toolName: string;
args: { url: string };
result: Array<{
title: string;
url: string;
content: string;
favicon: string;
publishedDate: string;
author: string;
}>;
}
export class WebCrawlTool extends WithDisposable(ShadowlessElement) {
@property({ attribute: false })
accessor data!: WebCrawlToolCall | WebCrawlToolResult;
@property({ attribute: false })
accessor host!: EditorHost;
renderToolCall() {
return html`
<tool-call-card
.name=${'Reading the website'}
.icon=${WebIcon()}
></tool-call-card>
`;
}
renderToolResult() {
if (this.data.type !== 'tool-result') {
return nothing;
}
const { favicon, title, content } = this.data.result[0];
return html`
<tool-result-card
.host=${this.host}
.name=${'Reading the website'}
.icon=${WebIcon()}
.results=${[
{
title: title,
icon: favicon,
content: content,
},
]}
></tool-result-card>
`;
}
protected override render() {
const { data } = this;
if (data.type === 'tool-call') {
return this.renderToolCall();
}
if (data.type === 'tool-result') {
return this.renderToolResult();
}
return nothing;
}
}

View File

@ -0,0 +1,79 @@
import { WithDisposable } from '@blocksuite/affine/global/lit';
import { type EditorHost, ShadowlessElement } from '@blocksuite/affine/std';
import { WebIcon } from '@blocksuite/icons/lit';
import { html, nothing } from 'lit';
import { property } from 'lit/decorators.js';
interface WebSearchToolCall {
type: 'tool-call';
toolCallId: string;
toolName: string;
args: { url: string };
}
interface WebSearchToolResult {
type: 'tool-result';
toolCallId: string;
toolName: string;
args: { url: string };
result: Array<{
title: string;
url: string;
content: string;
favicon: string;
publishedDate: string;
author: string;
}>;
}
export class WebSearchTool extends WithDisposable(ShadowlessElement) {
@property({ attribute: false })
accessor data!: WebSearchToolCall | WebSearchToolResult;
@property({ attribute: false })
accessor host!: EditorHost;
renderToolCall() {
return html`
<tool-call-card
.name=${'Search from web'}
.icon=${WebIcon()}
></tool-call-card>
`;
}
renderToolResult() {
if (this.data.type !== 'tool-result') {
return nothing;
}
const results = this.data.result.map(item => {
const { favicon, title, content } = item;
return {
title: title,
icon: favicon,
content: content,
};
});
return html`
<tool-result-card
.host=${this.host}
.name=${'The search is complete, and these webpages have been searched'}
.icon=${WebIcon()}
.results=${results}
></tool-result-card>
`;
}
protected override render() {
const { data } = this;
if (data.type === 'tool-call') {
return this.renderToolCall();
}
if (data.type === 'tool-result') {
return this.renderToolResult();
}
return nothing;
}
}

View File

@ -16,9 +16,10 @@ import { ChatAbortIcon } from '../../_common/icons';
import { type AIError, AIProvider } from '../../provider';
import { reportResponse } from '../../utils/action-reporter';
import { readBlobAsURL } from '../../utils/image';
import { mergeStreamObjects } from '../../utils/stream-objects';
import type { ChatChip, DocDisplayConfig } from '../ai-chat-chips/type';
import { isDocChip } from '../ai-chat-chips/utils';
import type { ChatMessage } from '../ai-chat-messages';
import { type ChatMessage, StreamObjectSchema } from '../ai-chat-messages';
import { MAX_IMAGE_COUNT } from './const';
import type {
AIChatInputContext,
@ -610,7 +611,19 @@ export class AIChatInput extends SignalWatcher(
for await (const text of stream) {
const messages = [...this.chatContextValue.messages];
const last = messages[messages.length - 1] as ChatMessage;
last.content += text;
try {
const parsed = StreamObjectSchema.safeParse(JSON.parse(text));
if (parsed.success) {
last.streamObjects = mergeStreamObjects([
...(last.streamObjects ?? []),
parsed.data,
]);
} else {
last.content += text;
}
} catch {
last.content += text;
}
this.updateContext({ messages, status: 'transmitting' });
}

View File

@ -1,6 +1,6 @@
import { z } from 'zod';
const StreamObjectSchema = z.discriminatedUnion('type', [
export const StreamObjectSchema = z.discriminatedUnion('type', [
z.object({
type: z.literal('text-delta'),
textDelta: z.string(),

View File

@ -112,7 +112,7 @@ export const createAIScrollableTextRenderer: (
maxHeight,
autoScroll
) => {
return (answer, state) => {
return (answer: string, state: AffineAIPanelState | undefined) => {
return html`<ai-scrollable-text-renderer
.host=${host}
.answer=${answer}

View File

@ -39,10 +39,7 @@ import React from 'react';
import { filter } from 'rxjs/operators';
import { markDownToDoc } from '../../utils';
import type {
AffineAIPanelState,
AffineAIPanelWidgetConfig,
} from '../widgets/ai-panel/type';
import type { AffineAIPanelState } from '../widgets/ai-panel/type';
export const getCustomPageEditorBlockSpecs: () => ExtensionType[] = () => {
const manager = getViewManager().config.init().value;
@ -430,11 +427,11 @@ export class TextRenderer extends WithDisposable(ShadowlessElement) {
accessor state: AffineAIPanelState | undefined = undefined;
}
export const createTextRenderer: (
export const createTextRenderer = (
host: EditorHost,
options: TextRendererOptions
) => AffineAIPanelWidgetConfig['answerRenderer'] = (host, options) => {
return (answer, state) => {
) => {
return (answer: string, state?: AffineAIPanelState) => {
return html`<text-renderer
contenteditable="false"
.host=${host}

View File

@ -28,6 +28,10 @@ import { ChatContentRichText } from './chat-panel/content/rich-text';
import { ChatMessageAction } from './chat-panel/message/action';
import { ChatMessageAssistant } from './chat-panel/message/assistant';
import { ChatMessageUser } from './chat-panel/message/user';
import { ToolCallCard } from './chat-panel/tools/tool-call-card';
import { ToolResultCard } from './chat-panel/tools/tool-result-card';
import { WebCrawlTool } from './chat-panel/tools/web-crawl';
import { WebSearchTool } from './chat-panel/tools/web-search';
import { ChatPanelAddPopover } from './components/ai-chat-chips/add-popover';
import { ChatPanelCandidatesPopover } from './components/ai-chat-chips/candidates-popover';
import { ChatPanelChips } from './components/ai-chat-chips/chat-panel-chips';
@ -156,6 +160,11 @@ export function registerAIEffects() {
customElements.define('chat-message-assistant', ChatMessageAssistant);
customElements.define('chat-message-user', ChatMessageUser);
customElements.define('tool-call-card', ToolCallCard);
customElements.define('tool-result-card', ToolResultCard);
customElements.define('web-crawl-tool', WebCrawlTool);
customElements.define('web-search-tool', WebSearchTool);
customElements.define(AFFINE_AI_PANEL_WIDGET, AffineAIPanelWidget);
customElements.define(AFFINE_EDGELESS_COPILOT_WIDGET, EdgelessCopilotWidget);

View File

@ -34,10 +34,14 @@ import type {
AIReasoningConfig,
} from '../components/ai-chat-input';
import type { ChatMessage } from '../components/ai-chat-messages';
import { ChatMessagesSchema } from '../components/ai-chat-messages';
import {
ChatMessagesSchema,
StreamObjectSchema,
} from '../components/ai-chat-messages';
import type { TextRendererOptions } from '../components/text-renderer';
import { AIChatErrorRenderer } from '../messages/error';
import { type AIError, AIProvider } from '../provider';
import { mergeStreamObjects } from '../utils/stream-objects';
import { PeekViewStyles } from './styles';
import type { ChatContext } from './types';
import { calcChildBound } from './utils';
@ -345,7 +349,12 @@ export class AIChatBlockPeekView extends LitElement {
last.id = '';
last.createdAt = new Date().toISOString();
}
this.updateContext({ messages, status: 'loading', error: null });
this.updateContext({
messages,
status: 'loading',
error: null,
abortController,
});
const { store } = this.host;
const stream = await AIProvider.actions.chat({
@ -362,11 +371,22 @@ export class AIChatBlockPeekView extends LitElement {
webSearch: this._isNetworkActive,
});
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;
try {
const parsed = StreamObjectSchema.safeParse(JSON.parse(text));
if (parsed.success) {
last.streamObjects = mergeStreamObjects([
...(last.streamObjects ?? []),
parsed.data,
]);
} else {
last.content += text;
}
} catch {
last.content += text;
}
this.updateContext({ messages, status: 'transmitting' });
}

View File

@ -142,7 +142,7 @@ export function textToText({
webSearch,
modelId,
},
workflow ? 'workflow' : undefined
workflow ? 'workflow' : 'stream-object'
);
AIProvider.LAST_ACTION_SESSIONID = sessionId;

View File

@ -0,0 +1,46 @@
import type { StreamObject } from '../components/ai-chat-messages';
export function mergeStreamObjects(chunks: StreamObject[] = []) {
return chunks.reduce((acc, curr) => {
const prev = acc.at(-1);
switch (curr.type) {
case 'reasoning':
case 'text-delta': {
if (prev && prev.type === curr.type) {
prev.textDelta += curr.textDelta;
} else {
acc.push(curr);
}
break;
}
case 'tool-result': {
const index = acc.findIndex(
item =>
item.type === 'tool-call' &&
item.toolCallId === curr.toolCallId &&
item.toolName === curr.toolName
);
if (index !== -1) {
acc[index] = curr;
} else {
acc.push(curr);
}
break;
}
default: {
acc.push(curr);
break;
}
}
return acc;
}, [] as StreamObject[]);
}
export function mergeStreamContent(chunks: StreamObject[]): string {
return chunks.reduce((acc, curr) => {
if (curr.type === 'text-delta') {
acc += curr.textDelta;
}
return acc;
}, '');
}

View File

@ -37,7 +37,12 @@ import { literal, unsafeStatic } from 'lit/static-html.js';
import { type AIError } from '../../provider';
import type { AIPanelGenerating } from './components/index.js';
import type { AffineAIPanelState, AffineAIPanelWidgetConfig } from './type.js';
import type {
AffineAIPanelState,
AffineAIPanelWidgetConfig,
AIActionAnswer,
} from './type.js';
import { mergeAIActionAnswer } from './utils';
export const AFFINE_AI_PANEL_WIDGET = 'affine-ai-panel-widget';
export class AffineAIPanelWidget extends WidgetComponent {
@ -103,7 +108,7 @@ export class AffineAIPanelWidget extends WidgetComponent {
private _abortController = new AbortController();
private _answer: string | null = null;
private _answer: AIActionAnswer | null = null;
private readonly _clearDiscardModal = () => {
if (this._discardModalAbort) {
@ -226,7 +231,7 @@ export class AffineAIPanelWidget extends WidgetComponent {
// reset answer
this._answer = null;
const update = (answer: string) => {
const update = (answer: AIActionAnswer) => {
this._answer = answer;
this.requestUpdate();
};
@ -345,7 +350,7 @@ export class AffineAIPanelWidget extends WidgetComponent {
};
get answer() {
return this._answer;
return this._answer ? mergeAIActionAnswer(this._answer) : null;
}
get inputText() {

View File

@ -1,6 +1,7 @@
import type { Signal } from '@preact/signals-core';
import type { nothing, TemplateResult } from 'lit';
import type { StreamObject } from '../../components/ai-chat-messages';
import type { AIItemGroupConfig } from '../../components/ai-item/types';
import type { AIError } from '../../provider';
@ -34,6 +35,11 @@ export interface AINetworkSearchConfig {
setEnabled: (state: boolean) => void;
}
export type AIActionAnswer = {
content: string;
streamObjects?: StreamObject[];
};
export interface AffineAIPanelWidgetConfig {
answerRenderer: (
answer: string,
@ -41,7 +47,7 @@ export interface AffineAIPanelWidgetConfig {
) => TemplateResult<1> | typeof nothing;
generateAnswer?: (props: {
input: string;
update: (answer: string) => void;
update: (answer: AIActionAnswer) => void;
finish: (type: 'success' | 'error' | 'aborted', err?: AIError) => void;
// Used to allow users to stop actively when generating
signal: AbortSignal;

View File

@ -2,6 +2,8 @@ import { isInsidePageEditor } from '@blocksuite/affine/shared/utils';
import type { EditorHost } from '@blocksuite/affine/std';
import type { AIItemGroupConfig } from '../../components/ai-item/types';
import { mergeStreamContent } from '../../utils/stream-objects';
import type { AIActionAnswer } from './type';
export function filterAIItemGroup(
host: EditorHost,
@ -19,3 +21,10 @@ export function filterAIItemGroup(
}))
.filter(group => group.items.length > 0);
}
export function mergeAIActionAnswer(answer: AIActionAnswer): string {
if (answer.streamObjects?.length) {
return mergeStreamContent(answer.streamObjects);
}
return answer.content;
}