legends810/testingnew
0
1import type { Message } from 'ai';2import { useCallback, useState } from 'react';3import { StreamingMessageParser } from '~/lib/runtime/message-parser';4import { workbenchStore } from '~/lib/stores/workbench';5import { createScopedLogger } from '~/utils/logger';6 7const logger = createScopedLogger('useMessageParser');8 9const messageParser = new StreamingMessageParser({10 callbacks: {11 onArtifactOpen: (data) => {12 logger.trace('onArtifactOpen', data);13 14 workbenchStore.showWorkbench.set(true);15 workbenchStore.addArtifact(data);16 },17 onArtifactClose: (data) => {18 logger.trace('onArtifactClose');19 20 workbenchStore.updateArtifact(data, { closed: true });21 },22 onActionOpen: (data) => {23 logger.trace('onActionOpen', data.action);24 25 // we only add shell actions when when the close tag got parsed because only then we have the content26 if (data.action.type === 'file') {27 workbenchStore.addAction(data);28 }29 },30 onActionClose: (data) => {31 logger.trace('onActionClose', data.action);32 33 if (data.action.type !== 'file') {34 workbenchStore.addAction(data);35 }36 37 workbenchStore.runAction(data);38 },39 onActionStream: (data) => {40 logger.trace('onActionStream', data.action);41 workbenchStore.runAction(data, true);42 },43 },44});45const extractTextContent = (message: Message) =>46 Array.isArray(message.content)47 ? (message.content.find((item) => item.type === 'text')?.text as string) || ''48 : message.content;49 50export function useMessageParser() {51 const [parsedMessages, setParsedMessages] = useState<{ [key: number]: string }>({});52 53 const parseMessages = useCallback((messages: Message[], isLoading: boolean) => {54 let reset = false;55 56 if (import.meta.env.DEV && !isLoading) {57 reset = true;58 messageParser.reset();59 }60 61 for (const [index, message] of messages.entries()) {62 if (message.role === 'assistant' || message.role === 'user') {63 const newParsedContent = messageParser.parse(message.id, extractTextContent(message));64 setParsedMessages((prevParsed) => ({65 ...prevParsed,66 [index]: !reset ? (prevParsed[index] || '') + newParsedContent : newParsedContent,67 }));68 }69 }70 }, []);71 72 return { parsedMessages, parseMessages };73}74 