Feat/parent child retrieval (#12086)

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
Signed-off-by: -LAN- <laipz8200@outlook.com>
Co-authored-by: AkaraChen <akarachen@outlook.com>
Co-authored-by: nite-knite <nkCoding@gmail.com>
Co-authored-by: Joel <iamjoel007@gmail.com>
Co-authored-by: Warren Chen <warren.chen830@gmail.com>
Co-authored-by: crazywoola <427733928@qq.com>
Co-authored-by: Yi Xiao <54782454+YIXIAO0@users.noreply.github.com>
Co-authored-by: yihong <zouzou0208@gmail.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>
Co-authored-by: KVOJJJin <jzongcode@gmail.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: JzoNgKVO <27049666+JzoNgKVO@users.noreply.github.com>
Co-authored-by: Charlie.Wei <luowei@cvte.com>
Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com>
Co-authored-by: huayaoyue6 <huayaoyue@163.com>
Co-authored-by: kurokobo <kuro664@gmail.com>
Co-authored-by: Matsuda <yiyth.fcb6@gmail.com>
Co-authored-by: shirochan <s.yusuke0711@gmail.com>
Co-authored-by: Jyong <76649700+JohnJyong@users.noreply.github.com>
Co-authored-by: Huỳnh Gia Bôi <boihuynh147@gmail.com>
Co-authored-by: Julian Huynh <julian.huynh@immersio.io>
Co-authored-by: Hash Brown <hi@xzd.me>
Co-authored-by: 非法操作 <hjlarry@163.com>
Co-authored-by: Kazuki Takamatsu <kazuki.takamatsu@chowagiken.co.jp>
Co-authored-by: Trey Dong <1346650911@qq.com>
Co-authored-by: VoidIsVoid <343750470@qq.com>
Co-authored-by: Gimling <huangjl@ruyi.ai>
Co-authored-by: xiandan-erizo <xiandan.erizo@gmail.com>
Co-authored-by: Muneyuki Noguchi <nogu.dev@gmail.com>
Co-authored-by: zhaobingshuang <1475195565@qq.com>
Co-authored-by: zhaobs <zhaobs@cailian.net>
Co-authored-by: suzuki.sh <s2terminal@users.noreply.github.com>
Co-authored-by: Yingchun Lai <laiyingchun@apache.org>
Co-authored-by: huanshare <huanshare@live.com>
Co-authored-by: huanshare <liuhuan101@longfor.com>
Co-authored-by: orangeclk <orangeclk@users.noreply.github.com>
Co-authored-by: 문정현 <120004247+JungHyunMoon@users.noreply.github.com>
Co-authored-by: barabicu <kztk533@gmail.com>
Co-authored-by: Wei Mingzhi <whistler_wmz@users.sf.net>
Co-authored-by: Paul van Oorschot <20116814+pvoo@users.noreply.github.com>
Co-authored-by: zkyTech <zhangkunyuan@hotmail.com>
Co-authored-by: zhangkunyuan <zhangkunyuan@cmhi.chinamobile.com>
Co-authored-by: Tommy <34446820+Asterovim@users.noreply.github.com>
Co-authored-by: zxhlyh <jasonapring2015@outlook.com>
Co-authored-by: Novice <857526207@qq.com>
Co-authored-by: Novice Lee <novicelee@NovicedeMacBook-Pro.local>
Co-authored-by: Novice Lee <novicelee@NoviPro.local>
Co-authored-by: zxhlyh <16177003+zxhlyh@users.noreply.github.com>
Co-authored-by: liuzhenghua <1090179900@qq.com>
Co-authored-by: Jiang <65766008+AlwaysBluer@users.noreply.github.com>
Co-authored-by: jiangzhijie <jiangzhijie.jzj@alibaba-inc.com>
Co-authored-by: Joe <79627742+ZhouhaoJiang@users.noreply.github.com>
Co-authored-by: Alok Shrivastwa <alok.shrivastwa@gmail.com>
Co-authored-by: Alok Shrivastwa <Alok.Shrivastwa@microland.com>
Co-authored-by: JasonVV <jasonwangiii@outlook.com>
Co-authored-by: Hiroshi Fujita <fujita-h@users.noreply.github.com>
Co-authored-by: Kevin9703 <51311316+Kevin9703@users.noreply.github.com>
Co-authored-by: NFish <douxc512@gmail.com>
Co-authored-by: Junyan Qin <1010553892@qq.com>
Co-authored-by: IWAI, Masaharu <iwaim.sub@gmail.com>
Co-authored-by: IWAI, Masaharu <iwai_masaharu@funkit.co.jp>
Co-authored-by: Bowen Liang <liangbowen@gf.com.cn>
Co-authored-by: luckylhb90 <luckylhb90@gmail.com>
Co-authored-by: hobo.l <hobo.l@binance.com>
Co-authored-by: douxc <7553076+douxc@users.noreply.github.com>

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

@@ -0,0 +1,223 @@
+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,
+  })
+}

web/service/knowledge/use-document.ts

@@ -0,0 +1,124 @@
+import {
+  useMutation,
+  useQuery,
+} from '@tanstack/react-query'
+import { del, get, patch } from '../base'
+import { useInvalid } from '../use-base'
+import type { MetadataType } from '../datasets'
+import type { DocumentDetailResponse, SimpleDocumentDetail, UpdateDocumentBatchParams } from '@/models/datasets'
+import { DocumentActionType } from '@/models/datasets'
+import type { CommonResponse } from '@/models/common'
+
+const NAME_SPACE = 'knowledge/document'
+
+const useDocumentListKey = [NAME_SPACE, 'documentList']
+export const useDocumentList = (payload: {
+  datasetId: string
+  query: {
+    keyword: string
+    page: number
+    limit: number
+  }
+}) => {
+  const { query, datasetId } = payload
+  return useQuery<{ data: SimpleDocumentDetail[] }>({
+    queryKey: [...useDocumentListKey, datasetId, query],
+    queryFn: () => get<{ data: SimpleDocumentDetail[] }>(`/datasets/${datasetId}/documents`, {
+      params: query,
+    }),
+  })
+}
+
+const useAutoDisabledDocumentKey = [NAME_SPACE, 'autoDisabledDocument']
+export const useAutoDisabledDocuments = (datasetId: string) => {
+  return useQuery({
+    queryKey: [...useAutoDisabledDocumentKey, datasetId],
+    queryFn: () => get<{ document_ids: string[] }>(`/datasets/${datasetId}/auto-disable-logs`),
+  })
+}
+
+export const useInvalidDisabledDocument = () => {
+  return useInvalid(useAutoDisabledDocumentKey)
+}
+
+const toBatchDocumentsIdParams = (documentIds: string[] | string) => {
+  const ids = Array.isArray(documentIds) ? documentIds : [documentIds]
+  return ids.map(id => `document_id=${id}`).join('&')
+}
+
+export const useDocumentBatchAction = (action: DocumentActionType) => {
+  return useMutation({
+    mutationFn: ({ datasetId, documentIds, documentId }: UpdateDocumentBatchParams) => {
+      return patch<CommonResponse>(`/datasets/${datasetId}/documents/status/${action}/batch?${toBatchDocumentsIdParams(documentId || documentIds!)}`)
+    },
+  })
+}
+
+export const useDocumentEnable = () => {
+  return useDocumentBatchAction(DocumentActionType.enable)
+}
+
+export const useDocumentDisable = () => {
+  return useDocumentBatchAction(DocumentActionType.disable)
+}
+
+export const useDocumentArchive = () => {
+  return useDocumentBatchAction(DocumentActionType.archive)
+}
+
+export const useDocumentUnArchive = () => {
+  return useDocumentBatchAction(DocumentActionType.unArchive)
+}
+
+export const useDocumentDelete = () => {
+  return useMutation({
+    mutationFn: ({ datasetId, documentIds, documentId }: UpdateDocumentBatchParams) => {
+      return del<CommonResponse>(`/datasets/${datasetId}/documents?${toBatchDocumentsIdParams(documentId || documentIds!)}`)
+    },
+  })
+}
+
+export const useSyncDocument = () => {
+  return useMutation({
+    mutationFn: ({ datasetId, documentId }: UpdateDocumentBatchParams) => {
+      return get<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/notion/sync`)
+    },
+  })
+}
+
+export const useSyncWebsite = () => {
+  return useMutation({
+    mutationFn: ({ datasetId, documentId }: UpdateDocumentBatchParams) => {
+      return get<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/website-sync`)
+    },
+  })
+}
+
+const useDocumentDetailKey = [NAME_SPACE, 'documentDetail']
+export const useDocumentDetail = (payload: {
+  datasetId: string
+  documentId: string
+  params: { metadata: MetadataType }
+}) => {
+  const { datasetId, documentId, params } = payload
+  return useQuery<DocumentDetailResponse>({
+    queryKey: [...useDocumentDetailKey, 'withoutMetaData', datasetId, documentId],
+    queryFn: () => get<DocumentDetailResponse>(`/datasets/${datasetId}/documents/${documentId}`, { params }),
+  })
+}
+
+export const useDocumentMetadata = (payload: {
+  datasetId: string
+  documentId: string
+  params: { metadata: MetadataType }
+}) => {
+  const { datasetId, documentId, params } = payload
+  return useQuery<DocumentDetailResponse>({
+    queryKey: [...useDocumentDetailKey, 'withMetaData', datasetId, documentId],
+    queryFn: () => get<DocumentDetailResponse>(`/datasets/${datasetId}/documents/${documentId}`, { params }),
+  })
+}
+
+export const useInvalidDocumentDetailKey = () => {
+  return useInvalid(useDocumentDetailKey)
+}

web/service/knowledge/use-segment.ts

@@ -0,0 +1,169 @@
+import { useMutation, useQuery } from '@tanstack/react-query'
+import { del, get, patch, post } from '../base'
+import type { CommonResponse } from '@/models/common'
+import type {
+  BatchImportResponse,
+  ChildChunkDetail,
+  ChildSegmentsResponse,
+  ChunkingMode,
+  SegmentDetailModel,
+  SegmentUpdater,
+  SegmentsResponse,
+} from '@/models/datasets'
+
+const NAME_SPACE = 'segment'
+
+export const useSegmentListKey = [NAME_SPACE, 'chunkList']
+
+export const useSegmentList = (
+  payload: {
+    datasetId: string
+    documentId: string
+    params: {
+      page: number
+      limit: number
+      keyword: string
+      enabled: boolean | 'all'
+    }
+  },
+  disable?: boolean,
+) => {
+  const { datasetId, documentId, params } = payload
+  const { page, limit, keyword, enabled } = params
+  return useQuery<SegmentsResponse>({
+    queryKey: [...useSegmentListKey, datasetId, documentId, page, limit, keyword, enabled],
+    queryFn: () => {
+      return get<SegmentsResponse>(`/datasets/${datasetId}/documents/${documentId}/segments`, { params })
+    },
+    enabled: !disable,
+  })
+}
+
+export const useUpdateSegment = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'update'],
+    mutationFn: (payload: { datasetId: string; documentId: string; segmentId: string; body: SegmentUpdater }) => {
+      const { datasetId, documentId, segmentId, body } = payload
+      return patch<{ data: SegmentDetailModel; doc_form: ChunkingMode }>(`/datasets/${datasetId}/documents/${documentId}/segments/${segmentId}`, { body })
+    },
+  })
+}
+
+export const useAddSegment = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'add'],
+    mutationFn: (payload: { datasetId: string; documentId: string; body: SegmentUpdater }) => {
+      const { datasetId, documentId, body } = payload
+      return post<{ data: SegmentDetailModel; doc_form: ChunkingMode }>(`/datasets/${datasetId}/documents/${documentId}/segment`, { body })
+    },
+  })
+}
+
+export const useEnableSegment = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'enable'],
+    mutationFn: (payload: { datasetId: string; documentId: string; segmentIds: string[] }) => {
+      const { datasetId, documentId, segmentIds } = payload
+      const query = segmentIds.map(id => `segment_id=${id}`).join('&')
+      return patch<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/segment/enable?${query}`)
+    },
+  })
+}
+
+export const useDisableSegment = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'disable'],
+    mutationFn: (payload: { datasetId: string; documentId: string; segmentIds: string[] }) => {
+      const { datasetId, documentId, segmentIds } = payload
+      const query = segmentIds.map(id => `segment_id=${id}`).join('&')
+      return patch<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/segment/disable?${query}`)
+    },
+  })
+}
+
+export const useDeleteSegment = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'delete'],
+    mutationFn: (payload: { datasetId: string; documentId: string; segmentIds: string[] }) => {
+      const { datasetId, documentId, segmentIds } = payload
+      const query = segmentIds.map(id => `segment_id=${id}`).join('&')
+      return del<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/segments?${query}`)
+    },
+  })
+}
+
+export const useChildSegmentListKey = [NAME_SPACE, 'childChunkList']
+
+export const useChildSegmentList = (
+  payload: {
+    datasetId: string
+    documentId: string
+    segmentId: string
+    params: {
+      page: number
+      limit: number
+      keyword: string
+    }
+  },
+  disable?: boolean,
+) => {
+  const { datasetId, documentId, segmentId, params } = payload
+  const { page, limit, keyword } = params
+  return useQuery({
+    queryKey: [...useChildSegmentListKey, datasetId, documentId, segmentId, page, limit, keyword],
+    queryFn: () => {
+      return get<ChildSegmentsResponse>(`/datasets/${datasetId}/documents/${documentId}/segments/${segmentId}/child_chunks`, { params })
+    },
+    enabled: !disable,
+  })
+}
+
+export const useDeleteChildSegment = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'childChunk', 'delete'],
+    mutationFn: (payload: { datasetId: string; documentId: string; segmentId: string; childChunkId: string }) => {
+      const { datasetId, documentId, segmentId, childChunkId } = payload
+      return del<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/segments/${segmentId}/child_chunks/${childChunkId}`)
+    },
+  })
+}
+
+export const useAddChildSegment = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'childChunk', 'add'],
+    mutationFn: (payload: { datasetId: string; documentId: string; segmentId: string; body: { content: string } }) => {
+      const { datasetId, documentId, segmentId, body } = payload
+      return post<{ data: ChildChunkDetail }>(`/datasets/${datasetId}/documents/${documentId}/segments/${segmentId}/child_chunks`, { body })
+    },
+  })
+}
+
+export const useUpdateChildSegment = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'childChunk', 'update'],
+    mutationFn: (payload: { datasetId: string; documentId: string; segmentId: string; childChunkId: string; body: { content: string } }) => {
+      const { datasetId, documentId, segmentId, childChunkId, body } = payload
+      return patch<{ data: ChildChunkDetail }>(`/datasets/${datasetId}/documents/${documentId}/segments/${segmentId}/child_chunks/${childChunkId}`, { body })
+    },
+  })
+}
+
+export const useSegmentBatchImport = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'batchImport'],
+    mutationFn: (payload: { url: string; body: FormData }) => {
+      const { url, body } = payload
+      return post<BatchImportResponse>(url, { body }, { bodyStringify: false, deleteContentType: true })
+    },
+  })
+}
+
+export const useCheckSegmentBatchImportProgress = () => {
+  return useMutation({
+    mutationKey: [NAME_SPACE, 'batchImport', 'checkProgress'],
+    mutationFn: (payload: { jobID: string }) => {
+      const { jobID } = payload
+      return get<BatchImportResponse>(`/datasets/batch_import_status/${jobID}`)
+    },
+  })
+}