107 lines
3.4 KiB
TypeScript
107 lines
3.4 KiB
TypeScript
import { useCallback, useRef, useState } from 'react';
|
|
import { StreamingMessageParser } from '~/.client/runtime/message-parser';
|
|
import { webBuilderStore } from '~/.client/stores/web-builder';
|
|
import type { UPageUIMessage } from '~/types/message';
|
|
import { createScopedLogger } from '~/utils/logger';
|
|
|
|
const logger = createScopedLogger('useMessageParser');
|
|
|
|
const chatStore = webBuilderStore.chatStore;
|
|
const messageParser = new StreamingMessageParser({
|
|
callbacks: {
|
|
onArtifactOpen: (data) => {
|
|
logger.trace('onArtifactOpen', data);
|
|
|
|
webBuilderStore.showWorkbench.set(true);
|
|
chatStore.addArtifact(data);
|
|
chatStore.setCurrentMessageId(data.messageId);
|
|
},
|
|
onArtifactClose: (data) => {
|
|
logger.trace('onArtifactClose');
|
|
|
|
chatStore.updateArtifact(data, { closed: true });
|
|
},
|
|
onActionOpen: (data) => {
|
|
logger.trace('onActionOpen', data.action);
|
|
chatStore.addAction(data);
|
|
},
|
|
onActionStream: (data) => {
|
|
logger.trace('onActionStream', data.action);
|
|
chatStore.runAction(data, true);
|
|
},
|
|
onActionClose: (data) => {
|
|
logger.trace('onActionClose', data.action);
|
|
chatStore.addAction(data);
|
|
chatStore.runAction(data);
|
|
},
|
|
},
|
|
});
|
|
const extractTextContent = (message: UPageUIMessage) =>
|
|
message.parts
|
|
.filter((part) => part.type === 'text')
|
|
.map((part) => part.text)
|
|
.join('\n');
|
|
|
|
export function useMessageParser() {
|
|
const [parsedMessages, setParsedMessages] = useState<{ [key: number]: string }>({});
|
|
const messageIdMap = useRef<Map<number, string>>(new Map());
|
|
|
|
const parseMessages = useCallback((messages: UPageUIMessage[], isLoading: boolean) => {
|
|
let reset = false;
|
|
|
|
if (import.meta.env.DEV && !isLoading) {
|
|
reset = true;
|
|
messageParser.reset();
|
|
}
|
|
|
|
for (const [index, message] of messages.entries()) {
|
|
if (message.role === 'assistant' || message.role === 'user') {
|
|
if (!messageIdMap.current.has(index)) {
|
|
messageIdMap.current.set(index, message.id);
|
|
}
|
|
// 当对应位置的 message id 发生变化时，重置解析
|
|
if (messageIdMap.current.get(index) !== message.id) {
|
|
reset = true;
|
|
messageParser.reset();
|
|
messageIdMap.current.set(index, message.id);
|
|
}
|
|
|
|
try {
|
|
const textContent = extractTextContent(message);
|
|
// 检查消息内容是否存在
|
|
if (textContent === undefined || textContent === null) {
|
|
logger.warn(`Message ${message.id} has no text content`);
|
|
continue;
|
|
}
|
|
|
|
// 解析消息内容
|
|
const newParsedContent = messageParser.parse(message.id, textContent);
|
|
if (!newParsedContent) {
|
|
continue;
|
|
}
|
|
|
|
// 更新解析后的消息
|
|
setParsedMessages((prevParsed) => {
|
|
const updatedContent = !reset ? (prevParsed[index] || '') + newParsedContent : newParsedContent;
|
|
return {
|
|
...prevParsed,
|
|
[index]: updatedContent,
|
|
};
|
|
});
|
|
} catch (error) {
|
|
// 捕获并记录解析过程中的错误
|
|
logger.error(`Error parsing message ${message.id}:`, error);
|
|
|
|
// 出错时保留原始消息内容
|
|
setParsedMessages((prevParsed) => ({
|
|
...prevParsed,
|
|
[index]: extractTextContent(message),
|
|
}));
|
|
}
|
|
}
|
|
}
|
|
}, []);
|
|
|
|
return { parsedMessages, parseMessages };
|
|
}
|