lib/messageUtils.ts

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

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

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

const ANSWERS_PREFIX = 'answers';
const INPUT_PREFIX = 'input';

export const generateAnswersImageMarkdown = (index: number, url: string) => {
  return `![${ANSWERS_PREFIX}-${index}](${url})`;
};

export const generateInputImageMarkdown = (url: string, index = 0) => {
  const prefix = 'input';
  return `![${INPUT_PREFIX}-${index}](${url})`;
};

export const cleanInputMessage = (content: string) => {
  return content.replace(/!\[input-.*?\)/g, '');
};

export const cleanAnswerMessage = (content: string) => {
  return content.replace(/!\[answers.*?\.png\)/g, '');
};

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>', '') +
      '\n\n' +
      images
        .map((_, index) => generateAnswersImageMarkdown(index, '/loading.gif'))
        .join('') +
      rest.join(''),
    images: images,
  };
};