src/shared/utils/parseDocumnetsInLlmAnswer.ts

import { ChunksArrayItem, ChunkType } from "@/api/predictions/types";

export const parseDocumetnsInLlmAnswer = (text: string, chunks: ChunksArrayItem[]) => {
  const documentsListString = "Список докумен";
  const initialString = text;
  const index = initialString.indexOf(documentsListString);
  const documentsSubstring = initialString.substring(index);
  const regex = /\[(\d+)\]/g;
  let match;

  const matches: number[] = [];
  let resultString = initialString.substring(0, index);

  while ((match = regex.exec(documentsSubstring)) !== null) {
    const number = Number(match[1]);
    if (number != null) {
      matches.push(number - 1);
    }
  }
  const documentsList = matches
    .map((matchNumber) => {
      let docName: string = "";
      if (chunks[matchNumber]) {
        switch (chunks[matchNumber].type) {
          case ChunkType.doc_chunks:
            if (chunks[matchNumber].chunk.doc_chunks) {
              docName =
                chunks[matchNumber].chunk.doc_chunks?.title && chunks[matchNumber].chunk.doc_chunks?.title !== "unknown"
                  ? chunks[matchNumber].chunk.doc_chunks?.title.trim()
                  : chunks[matchNumber].chunk.doc_chunks?.filename.replace(".json", "").trim();
            }
            break;
          case ChunkType.people_search:
            if (chunks[matchNumber].chunk.people_search) {
              docName = `Информация о сотруднике ${chunks[matchNumber].chunk.people_search?.person_name.trim()}`;
            }
            break;
          case ChunkType.groups_search:
            if (chunks[matchNumber].chunk.groups_search) {
              docName = `Информация о группе ${chunks[matchNumber].chunk.groups_search?.group_name.trim()}`;
            }
            break;
          case ChunkType.rocks_nn_search:
            if (chunks[matchNumber].chunk.rocks_nn_search) {
              docName = `Информация о РОКС ${chunks[matchNumber].chunk.rocks_nn_search?.division.trim()}`;
            }
            break;
          case ChunkType.segmentation_search:
            if (chunks[matchNumber].chunk.segmentation_search) {
              docName = `Информация модели сегментации ${chunks[
                matchNumber
              ].chunk.segmentation_search?.segmentation_model.trim()}`;
            }
            break;
          default:
            break;
        }
        return `Документ ${matchNumber + 1}. ${docName} \n`;
      }
    })
    .join("");

  if (matches.length > 0) {
    resultString = resultString + "\n" + "Список документов:\n" + documentsList + "\n" + "Конец ответа.";
  }

  return { matches, resultString };
};