59 lines
1.9 KiB
TypeScript
Raw Normal View History

import { useTranslation } from 'react-i18next'
import { AddDocumentsStep } from './types'
2025-05-22 14:49:40 +08:00
import type { DataSourceOption, Datasource } from '@/app/components/rag-pipeline/components/panel/test-run/types'
import { useMemo } from 'react'
import { BlockEnum, type Node } from '@/app/components/workflow/types'
import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
import type { DatasourceType } from '@/models/pipeline'
export const useAddDocumentsSteps = () => {
const { t } = useTranslation()
const steps = [
{
label: t('datasetPipeline.addDocuments.steps.chooseDatasource'),
value: AddDocumentsStep.dataSource,
},
{
label: t('datasetPipeline.addDocuments.steps.processDocuments'),
value: AddDocumentsStep.processDocuments,
},
{
label: t('datasetPipeline.addDocuments.steps.processingDocuments'),
value: AddDocumentsStep.processingDocuments,
},
]
return steps
}
2025-05-22 14:49:40 +08:00
export const useDatasourceOptions = (pipelineNodes: Node<DataSourceNodeType>[]) => {
const datasourceNodes = pipelineNodes.filter(node => node.data.type === BlockEnum.DataSource)
2025-05-22 14:49:40 +08:00
const datasources: Datasource[] = useMemo(() => {
return datasourceNodes.map((node) => {
return {
nodeId: node.id,
type: node.data.provider_type as DatasourceType,
variables: node.data.variables || [],
description: node.data.desc || '',
docTitle: '', // todo: Add docTitle and docLink if needed, or remove these properties if not used
docLink: '',
2025-05-22 14:49:40 +08:00
}
})
}, [datasourceNodes])
2025-05-22 14:49:40 +08:00
const options = useMemo(() => {
const options: DataSourceOption[] = []
datasourceNodes.forEach((node) => {
const type = node.data.provider_type as DatasourceType
const label = node.data.title
options.push({
label,
value: node.id,
type,
})
2025-05-22 14:49:40 +08:00
})
return options
}, [datasourceNodes])
2025-05-22 14:49:40 +08:00
return { datasources, options }
}