lib/hooks/useCleanedUpMessages.ts

import { useMemo } from 'react';
import { MessageBase } from '../types';
import { CLEANED_SEPARATOR } from '../constants';

const PAIRS: Record<string, string> = {
  '┍': '┑',
  '┝': '┥',
  '├': '┤',
  '┕': '┙',
};

const MIDDLE_STARTER = '┝';
const MIDDLE_SEPARATOR = '┿';

export const getCleanedUpMessages = ({
  content,
  role,
}: Pick<MessageBase, 'role' | 'content'>) => {
  if (role === 'user') {
    return {
      content,
    };
  }
  if (content.split(CLEANED_SEPARATOR).length === 2) {
    return {
      logs: content.split(CLEANED_SEPARATOR)[0],
      content: content.split(CLEANED_SEPARATOR)[1],
    };
  }
  const [logs = '', answer = ''] = content.split('<ANSWER>');
  const cleanedLogs = [];
  let left = 0;
  let right = 0;
  while (right < logs.length) {
    if (Object.keys(PAIRS).includes(content[right])) {
      cleanedLogs.push(content.substring(left, right));
      left = right++;
      while (
        right < content.length &&
        PAIRS[content[left]] !== content[right]
      ) {
        right++;
      }
      if (content[left] === MIDDLE_STARTER) {
        // add the text alignment so it can be shown as a table
        const separators = logs
          .substring(left, right)
          .split(MIDDLE_SEPARATOR).length;
        if (separators > 0) {
          cleanedLogs.push(
            Array(separators + 1)
              .fill('|')
              .join(' :- '),
          );
        }
      }
      left = ++right;
    } else {
      right++;
    }
  }
  cleanedLogs.push(content.substring(left, right));
  const [answerText, imagesStr = ''] = answer.split('<VIZ>');
  const [imagesArrayStr, ...rest] = imagesStr.split('</VIZ>');
  const images = imagesArrayStr
    .split('</IMG>')
    .map(str => str.replace('<IMG>', ''))
    .slice(0, -1);
  return {
    logs: cleanedLogs.join('').replace(/│/g, '|').split('|\n\n|').join('|\n|'),
    content:
      answerText.replace('</</ANSWER>', '').replace('</ANSWER>', '') +
      images.map((_, index) => `![answers-${index}](/loading.gif)`).join('') +
      rest.join(''),
    images: images,
  };
};

export const useCleanedUpMessages = ({ content, role }: MessageBase) => {
  const cleanedMessage = useMemo(() => {
    return getCleanedUpMessages({ content, role });
  }, [content, role]);
  return cleanedMessage;
};