tjdify/web/service/knowledge/use-create-dataset.ts

import groupBy from 'lodash-es/groupBy'
import type { MutationOptions } from '@tanstack/react-query'
import { useMutation } from '@tanstack/react-query'
import { createDocument, createFirstDocument, fetchDefaultProcessRule, fetchFileIndexingEstimate } from '../datasets'
import { type IndexingType } from '@/app/components/datasets/create/step-two'
import type { ChunkingMode, CrawlOptions, CrawlResultItem, CreateDocumentReq, CustomFile, DataSourceType, FileIndexingEstimateResponse, IndexingEstimateParams, NotionInfo, ProcessRule, ProcessRuleResponse, createDocumentResponse } from '@/models/datasets'
import type { DataSourceProvider, NotionPage } from '@/models/common'

export const getNotionInfo = (
  notionPages: NotionPage[],
) => {
  const workspacesMap = groupBy(notionPages, 'workspace_id')
  const workspaces = Object.keys(workspacesMap).map((workspaceId) => {
    return {
      workspaceId,
      pages: workspacesMap[workspaceId],
    }
  })
  return workspaces.map((workspace) => {
    return {
      workspace_id: workspace.workspaceId,
      pages: workspace.pages.map((page) => {
        const { page_id, page_name, page_icon, type } = page
        return {
          page_id,
          page_name,
          page_icon,
          type,
        }
      }),
    }
  }) as NotionInfo[]
}

export const getWebsiteInfo = (
  opts: {
    websiteCrawlProvider: DataSourceProvider
    websiteCrawlJobId: string
    websitePages: CrawlResultItem[]
    crawlOptions?: CrawlOptions
  },
) => {
  const { websiteCrawlProvider, websiteCrawlJobId, websitePages, crawlOptions } = opts
  return {
    provider: websiteCrawlProvider,
    job_id: websiteCrawlJobId,
    urls: websitePages.map(page => page.source_url),
    only_main_content: crawlOptions?.only_main_content,
  }
}

type GetFileIndexingEstimateParamsOptionBase = {
  docForm: ChunkingMode
  docLanguage: string
  indexingTechnique: IndexingType
  processRule: ProcessRule
  dataset_id: string
}

type GetFileIndexingEstimateParamsOptionFile = GetFileIndexingEstimateParamsOptionBase & {
  dataSourceType: DataSourceType.FILE
  files: CustomFile[]
}

const getFileIndexingEstimateParamsForFile = ({
  docForm,
  docLanguage,
  dataSourceType,
  files,
  indexingTechnique,
  processRule,
  dataset_id,
}: GetFileIndexingEstimateParamsOptionFile): IndexingEstimateParams => {
  return {
    info_list: {
      data_source_type: dataSourceType,
      file_info_list: {
        file_ids: files.map(file => file.id) as string[],
      },
    },
    indexing_technique: indexingTechnique,
    process_rule: processRule,
    doc_form: docForm,
    doc_language: docLanguage,
    dataset_id,
  }
}

export const useFetchFileIndexingEstimateForFile = (
  options: GetFileIndexingEstimateParamsOptionFile,
  mutationOptions: MutationOptions<FileIndexingEstimateResponse> = {},
) => {
  return useMutation({
    mutationFn: async () => {
      return fetchFileIndexingEstimate(getFileIndexingEstimateParamsForFile(options))
    },
    ...mutationOptions,
  })
}

type GetFileIndexingEstimateParamsOptionNotion = GetFileIndexingEstimateParamsOptionBase & {
  dataSourceType: DataSourceType.NOTION
  notionPages: NotionPage[]
}

const getFileIndexingEstimateParamsForNotion = ({
  docForm,
  docLanguage,
  dataSourceType,
  notionPages,
  indexingTechnique,
  processRule,
  dataset_id,
}: GetFileIndexingEstimateParamsOptionNotion): IndexingEstimateParams => {
  return {
    info_list: {
      data_source_type: dataSourceType,
      notion_info_list: getNotionInfo(notionPages),
    },
    indexing_technique: indexingTechnique,
    process_rule: processRule,
    doc_form: docForm,
    doc_language: docLanguage,
    dataset_id,
  }
}

export const useFetchFileIndexingEstimateForNotion = (
  options: GetFileIndexingEstimateParamsOptionNotion,
  mutationOptions: MutationOptions<FileIndexingEstimateResponse> = {},
) => {
  return useMutation({
    mutationFn: async () => {
      return fetchFileIndexingEstimate(getFileIndexingEstimateParamsForNotion(options))
    },
    ...mutationOptions,
  })
}

type GetFileIndexingEstimateParamsOptionWeb = GetFileIndexingEstimateParamsOptionBase & {
  dataSourceType: DataSourceType.WEB
  websitePages: CrawlResultItem[]
  crawlOptions?: CrawlOptions
  websiteCrawlProvider: DataSourceProvider
  websiteCrawlJobId: string
}

const getFileIndexingEstimateParamsForWeb = ({
  docForm,
  docLanguage,
  dataSourceType,
  websitePages,
  crawlOptions,
  websiteCrawlProvider,
  websiteCrawlJobId,
  indexingTechnique,
  processRule,
  dataset_id,
}: GetFileIndexingEstimateParamsOptionWeb): IndexingEstimateParams => {
  return {
    info_list: {
      data_source_type: dataSourceType,
      website_info_list: getWebsiteInfo({
        websiteCrawlProvider,
        websiteCrawlJobId,
        websitePages,
        crawlOptions,
      }),
    },
    indexing_technique: indexingTechnique,
    process_rule: processRule,
    doc_form: docForm,
    doc_language: docLanguage,
    dataset_id,
  }
}

export const useFetchFileIndexingEstimateForWeb = (
  options: GetFileIndexingEstimateParamsOptionWeb,
  mutationOptions: MutationOptions<FileIndexingEstimateResponse> = {},
) => {
  return useMutation({
    mutationFn: async () => {
      return fetchFileIndexingEstimate(getFileIndexingEstimateParamsForWeb(options))
    },
    ...mutationOptions,
  })
}

export const useCreateFirstDocument = (
  mutationOptions: MutationOptions<createDocumentResponse, Error, CreateDocumentReq> = {},
) => {
  return useMutation({
    mutationFn: async (createDocumentReq: CreateDocumentReq,
    ) => {
      return createFirstDocument({ body: createDocumentReq })
    },
    ...mutationOptions,
  })
}

export const useCreateDocument = (
  datasetId: string,
  mutationOptions: MutationOptions<createDocumentResponse, Error, CreateDocumentReq> = {},
) => {
  return useMutation({
    mutationFn: async (req: CreateDocumentReq) => {
      return createDocument({ datasetId, body: req })
    },
    ...mutationOptions,
  })
}

export const useFetchDefaultProcessRule = (
  mutationOptions: MutationOptions<ProcessRuleResponse, Error, string> = {},
) => {
  return useMutation({
    mutationFn: async (url: string) => {
      return fetchDefaultProcessRule({ url })
    },
    ...mutationOptions,
  })
}