src/api/predictions/predictionsApi.ts

import { query } from "@/shared/api/query";

import axios from "axios";
import { parseDocumetnsInLlmAnswer } from "@/shared/utils/parseDocumnetsInLlmAnswer";
import {
  ChunkRequestType,
  ChunksResponseType,
  ChunksArrayItem,
  ChunkType,
  LlmRequestType,
  LlmResponseType,
  FeedbackRequestType,
  FeedbackResponseType,
  LogRequestType,
  LogResponseType,
  AbbreviationsDataType,
  AbbreviationsResponseDataType,
} from "./types";

export const getChunks = async (
  data: ChunkRequestType,
  signal: AbortSignal
): Promise<{ chunks: ChunksResponseType; chunksArray: ChunksArrayItem[] }> => {
  const response = await query<ChunksResponseType>({
    url: "/chunks",
    method: "post",
    data,
    signal,
  });

  if ("error" in response) {
    if (axios.isCancel(response.error)) {
      console.log("Запроса был отменен");
    }
    throw new Error(`Ошибка: ${response.error.status}`);
  }
  const chunksArray: ChunksArrayItem[] = [];

  // Проверки на null добавлены для корретной сквозной нумерации результатов в случае,
  // если чанков одного типа не пришло с бэка
  if (response.data.doc_chunks !== null) {
    response.data.doc_chunks?.forEach((element) => {
      chunksArray.push({
        type: ChunkType.doc_chunks,
        number: chunksArray.length,
        chunk: { doc_chunks: element },
      });
    });
  }

  if (response.data.people_search !== null) {
    response.data.people_search?.forEach((element) => {
      chunksArray.push({
        type: ChunkType.people_search,
        number: chunksArray.length,
        chunk: { people_search: element },
      });
    });
  }

  if (response.data.groups_search !== null) {
    chunksArray.push({
      type: ChunkType.groups_search,
      number: chunksArray.length,
      chunk: { groups_search: response.data.groups_search },
    });
  }

  if (response.data.rocks_nn_search !== null) {
    chunksArray.push({
      type: ChunkType.rocks_nn_search,
      number: chunksArray.length,
      chunk: { rocks_nn_search: response.data.rocks_nn_search },
    });
  }

  if (response.data.segmentation_search !== null) {
    chunksArray.push({
      type: ChunkType.segmentation_search,
      number: chunksArray.length,
      chunk: { segmentation_search: response.data.segmentation_search },
    });
  }

  return { chunks: response.data, chunksArray };
};

export const getLlmResponse = async (
  data: LlmRequestType,
  signal: AbortSignal
): Promise<LlmResponseType> => {
  const response = await query<LlmResponseType>({
    url: "/answer_llm",
    method: "post",
    data: { query: data.query, chunks: data.chunks },
    signal,
  });

  if ("error" in response) {
    if (axios.isCancel(response.error)) {
      console.log("Запроса был отменен");
    }
    throw new Error(`Ошибка: ${response.error.status}`);
  }
  const { matches, resultString } = parseDocumetnsInLlmAnswer(
    response.data.answer_llm,
    data.chunksArray
  );

  return {
    answer_llm: resultString,
    log_id: response.data.log_id,
    documentsNumbers: matches,
  };
};

export const postFeedback = async (
  data: FeedbackRequestType
): Promise<FeedbackResponseType> => {
  const response = await query<FeedbackResponseType>({
    url: "/feedback",
    method: "post",
    data,
  });

  if ("error" in response) {
    throw new Error(`Ошибка: ${response.error.status}`);
  }

  return response.data;
};

export const getLogs = async (
  data: LogRequestType | undefined
): Promise<LogResponseType[]> => {
  const response = await query<LogResponseType[]>({
    url: "/logs",
    method: "get",
    params: data,
  });
  if ("error" in response) {
    throw new Error(`Ошибка: ${response.error.status}`);
  }

  return response.data;
};

export const getAbbreviations = async (): Promise<AbbreviationsDataType> => {
  const response = await query<AbbreviationsResponseDataType>({
    url: "/collection/default",
    method: "get",
  });
  if ("error" in response) {
    throw new Error(`Ошибка: ${response.error.status}`);
  }

  return response.data.acronyms;
};