lib/hooks/useCleanedUpMessages.ts

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

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

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

export const useCleanedUpMessages = ({ content, role }: MessageBase) => {
  return useMemo(() => {
    if (role === 'user') {
      return {
        content,
      };
    }
    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));
    return {
      logs: cleanedLogs
        .join('')
        .replace(/│/g, '|')
        .split('|\n\n|')
        .join('|\n|'),
      content: answer.replace('</</ANSWER>', '').replace('</ANSWER>', ''),
    };
  }, [content, role]);
};