CoolFace
Apppublic

legends810/testingnew

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
useMessageParser.ts74 linesDownload Raw Back to hooks
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