mirror of
https://github.com/langgenius/dify.git
synced 2026-05-06 02:18:08 +08:00
refactor(datasets): extract hooks and components with comprehensive tests (#31707)
Co-authored-by: CodingOnStar <hanxujiang@dify.com> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Coding On Star
GitHub
@ -0,0 +1,763 @@
|
||||
import type { DataSet } from '@/models/datasets'
|
||||
import type { RetrievalConfig } from '@/types/app'
|
||||
import { act, renderHook, waitFor } from '@testing-library/react'
|
||||
import { ChunkingMode, DatasetPermission, DataSourceType, WeightedScoreEnum } from '@/models/datasets'
|
||||
import { RETRIEVE_METHOD } from '@/types/app'
|
||||
import { IndexingType } from '../../../create/step-two'
|
||||
import { useFormState } from './use-form-state'
|
||||
|
||||
// Mock contexts
|
||||
const mockMutateDatasets = vi.fn()
|
||||
const mockInvalidDatasetList = vi.fn()
|
||||
|
||||
vi.mock('@/context/app-context', () => ({
|
||||
useSelector: () => false, // isCurrentWorkspaceDatasetOperator
|
||||
}))
|
||||
|
||||
const createDefaultMockDataset = (): DataSet => ({
|
||||
id: 'dataset-1',
|
||||
name: 'Test Dataset',
|
||||
description: 'Test description',
|
||||
permission: DatasetPermission.onlyMe,
|
||||
icon_info: {
|
||||
icon_type: 'emoji',
|
||||
icon: '📚',
|
||||
icon_background: '#FFFFFF',
|
||||
icon_url: '',
|
||||
},
|
||||
indexing_technique: IndexingType.QUALIFIED,
|
||||
indexing_status: 'completed',
|
||||
data_source_type: DataSourceType.FILE,
|
||||
doc_form: ChunkingMode.text,
|
||||
embedding_model: 'text-embedding-ada-002',
|
||||
embedding_model_provider: 'openai',
|
||||
embedding_available: true,
|
||||
app_count: 0,
|
||||
document_count: 5,
|
||||
total_document_count: 5,
|
||||
word_count: 1000,
|
||||
provider: 'vendor',
|
||||
tags: [],
|
||||
partial_member_list: [],
|
||||
external_knowledge_info: {
|
||||
external_knowledge_id: 'ext-1',
|
||||
external_knowledge_api_id: 'api-1',
|
||||
external_knowledge_api_name: 'External API',
|
||||
external_knowledge_api_endpoint: 'https://api.example.com',
|
||||
},
|
||||
external_retrieval_model: {
|
||||
top_k: 3,
|
||||
score_threshold: 0.7,
|
||||
score_threshold_enabled: true,
|
||||
},
|
||||
retrieval_model_dict: {
|
||||
search_method: RETRIEVE_METHOD.semantic,
|
||||
reranking_enable: false,
|
||||
reranking_model: {
|
||||
reranking_provider_name: '',
|
||||
reranking_model_name: '',
|
||||
},
|
||||
top_k: 3,
|
||||
score_threshold_enabled: false,
|
||||
score_threshold: 0.5,
|
||||
} as RetrievalConfig,
|
||||
retrieval_model: {
|
||||
search_method: RETRIEVE_METHOD.semantic,
|
||||
reranking_enable: false,
|
||||
reranking_model: {
|
||||
reranking_provider_name: '',
|
||||
reranking_model_name: '',
|
||||
},
|
||||
top_k: 3,
|
||||
score_threshold_enabled: false,
|
||||
score_threshold: 0.5,
|
||||
} as RetrievalConfig,
|
||||
built_in_field_enabled: false,
|
||||
keyword_number: 10,
|
||||
created_by: 'user-1',
|
||||
updated_by: 'user-1',
|
||||
updated_at: Date.now(),
|
||||
runtime_mode: 'general',
|
||||
enable_api: true,
|
||||
is_multimodal: false,
|
||||
})
|
||||
|
||||
let mockDataset: DataSet = createDefaultMockDataset()
|
||||
|
||||
vi.mock('@/context/dataset-detail', () => ({
|
||||
useDatasetDetailContextWithSelector: (selector: (state: { dataset: DataSet | null, mutateDatasetRes: () => void }) => unknown) => {
|
||||
const state = {
|
||||
dataset: mockDataset,
|
||||
mutateDatasetRes: mockMutateDatasets,
|
||||
}
|
||||
return selector(state)
|
||||
},
|
||||
}))
|
||||
|
||||
// Mock services
|
||||
vi.mock('@/service/datasets', () => ({
|
||||
updateDatasetSetting: vi.fn().mockResolvedValue({}),
|
||||
}))
|
||||
|
||||
vi.mock('@/service/knowledge/use-dataset', () => ({
|
||||
useInvalidDatasetList: () => mockInvalidDatasetList,
|
||||
}))
|
||||
|
||||
vi.mock('@/service/use-common', () => ({
|
||||
useMembers: () => ({
|
||||
data: {
|
||||
accounts: [
|
||||
{ id: 'user-1', name: 'User 1', email: 'user1@example.com', role: 'owner', avatar: '', avatar_url: '', last_login_at: '', created_at: '', status: 'active' },
|
||||
{ id: 'user-2', name: 'User 2', email: 'user2@example.com', role: 'admin', avatar: '', avatar_url: '', last_login_at: '', created_at: '', status: 'active' },
|
||||
],
|
||||
},
|
||||
}),
|
||||
}))
|
||||
|
||||
vi.mock('@/app/components/header/account-setting/model-provider-page/hooks', () => ({
|
||||
useModelList: () => ({ data: [] }),
|
||||
}))
|
||||
|
||||
vi.mock('@/app/components/datasets/common/check-rerank-model', () => ({
|
||||
isReRankModelSelected: () => true,
|
||||
}))
|
||||
|
||||
vi.mock('@/app/components/base/toast', () => ({
|
||||
default: {
|
||||
notify: vi.fn(),
|
||||
},
|
||||
}))
|
||||
|
||||
describe('useFormState', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
mockDataset = createDefaultMockDataset()
|
||||
})
|
||||
|
||||
describe('Initial State', () => {
|
||||
it('should initialize with dataset values', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
expect(result.current.name).toBe('Test Dataset')
|
||||
expect(result.current.description).toBe('Test description')
|
||||
expect(result.current.permission).toBe(DatasetPermission.onlyMe)
|
||||
expect(result.current.indexMethod).toBe(IndexingType.QUALIFIED)
|
||||
expect(result.current.keywordNumber).toBe(10)
|
||||
})
|
||||
|
||||
it('should initialize icon info from dataset', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
expect(result.current.iconInfo).toEqual({
|
||||
icon_type: 'emoji',
|
||||
icon: '📚',
|
||||
icon_background: '#FFFFFF',
|
||||
icon_url: '',
|
||||
})
|
||||
})
|
||||
|
||||
it('should initialize external retrieval settings', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
expect(result.current.topK).toBe(3)
|
||||
expect(result.current.scoreThreshold).toBe(0.7)
|
||||
expect(result.current.scoreThresholdEnabled).toBe(true)
|
||||
})
|
||||
|
||||
it('should derive member list from API data', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
expect(result.current.memberList).toHaveLength(2)
|
||||
expect(result.current.memberList[0].name).toBe('User 1')
|
||||
})
|
||||
|
||||
it('should return currentDataset from context', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
expect(result.current.currentDataset).toBeDefined()
|
||||
expect(result.current.currentDataset?.id).toBe('dataset-1')
|
||||
})
|
||||
})
|
||||
|
||||
describe('State Setters', () => {
|
||||
it('should update name when setName is called', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setName('New Name')
|
||||
})
|
||||
|
||||
expect(result.current.name).toBe('New Name')
|
||||
})
|
||||
|
||||
it('should update description when setDescription is called', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setDescription('New Description')
|
||||
})
|
||||
|
||||
expect(result.current.description).toBe('New Description')
|
||||
})
|
||||
|
||||
it('should update permission when setPermission is called', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setPermission(DatasetPermission.allTeamMembers)
|
||||
})
|
||||
|
||||
expect(result.current.permission).toBe(DatasetPermission.allTeamMembers)
|
||||
})
|
||||
|
||||
it('should update indexMethod when setIndexMethod is called', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setIndexMethod(IndexingType.ECONOMICAL)
|
||||
})
|
||||
|
||||
expect(result.current.indexMethod).toBe(IndexingType.ECONOMICAL)
|
||||
})
|
||||
|
||||
it('should update keywordNumber when setKeywordNumber is called', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setKeywordNumber(20)
|
||||
})
|
||||
|
||||
expect(result.current.keywordNumber).toBe(20)
|
||||
})
|
||||
|
||||
it('should update selectedMemberIDs when setSelectedMemberIDs is called', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setSelectedMemberIDs(['user-1', 'user-2'])
|
||||
})
|
||||
|
||||
expect(result.current.selectedMemberIDs).toEqual(['user-1', 'user-2'])
|
||||
})
|
||||
})
|
||||
|
||||
describe('Icon Handlers', () => {
|
||||
it('should open app icon picker and save previous icon', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleOpenAppIconPicker()
|
||||
})
|
||||
|
||||
expect(result.current.showAppIconPicker).toBe(true)
|
||||
})
|
||||
|
||||
it('should select emoji icon and close picker', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleOpenAppIconPicker()
|
||||
})
|
||||
|
||||
act(() => {
|
||||
result.current.handleSelectAppIcon({
|
||||
type: 'emoji',
|
||||
icon: '🎉',
|
||||
background: '#FF0000',
|
||||
})
|
||||
})
|
||||
|
||||
expect(result.current.showAppIconPicker).toBe(false)
|
||||
expect(result.current.iconInfo).toEqual({
|
||||
icon_type: 'emoji',
|
||||
icon: '🎉',
|
||||
icon_background: '#FF0000',
|
||||
icon_url: undefined,
|
||||
})
|
||||
})
|
||||
|
||||
it('should select image icon and close picker', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleOpenAppIconPicker()
|
||||
})
|
||||
|
||||
act(() => {
|
||||
result.current.handleSelectAppIcon({
|
||||
type: 'image',
|
||||
fileId: 'file-123',
|
||||
url: 'https://example.com/icon.png',
|
||||
})
|
||||
})
|
||||
|
||||
expect(result.current.showAppIconPicker).toBe(false)
|
||||
expect(result.current.iconInfo).toEqual({
|
||||
icon_type: 'image',
|
||||
icon: 'file-123',
|
||||
icon_background: undefined,
|
||||
icon_url: 'https://example.com/icon.png',
|
||||
})
|
||||
})
|
||||
|
||||
it('should restore previous icon when picker is closed', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleOpenAppIconPicker()
|
||||
})
|
||||
|
||||
act(() => {
|
||||
result.current.handleSelectAppIcon({
|
||||
type: 'emoji',
|
||||
icon: '🎉',
|
||||
background: '#FF0000',
|
||||
})
|
||||
})
|
||||
|
||||
act(() => {
|
||||
result.current.handleOpenAppIconPicker()
|
||||
})
|
||||
|
||||
act(() => {
|
||||
result.current.handleCloseAppIconPicker()
|
||||
})
|
||||
|
||||
expect(result.current.showAppIconPicker).toBe(false)
|
||||
// After close, icon should be restored to the icon before opening
|
||||
expect(result.current.iconInfo).toEqual({
|
||||
icon_type: 'emoji',
|
||||
icon: '🎉',
|
||||
icon_background: '#FF0000',
|
||||
icon_url: undefined,
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('External Retrieval Settings Handler', () => {
|
||||
it('should update topK when provided', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleSettingsChange({ top_k: 5 })
|
||||
})
|
||||
|
||||
expect(result.current.topK).toBe(5)
|
||||
})
|
||||
|
||||
it('should update scoreThreshold when provided', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleSettingsChange({ score_threshold: 0.8 })
|
||||
})
|
||||
|
||||
expect(result.current.scoreThreshold).toBe(0.8)
|
||||
})
|
||||
|
||||
it('should update scoreThresholdEnabled when provided', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleSettingsChange({ score_threshold_enabled: false })
|
||||
})
|
||||
|
||||
expect(result.current.scoreThresholdEnabled).toBe(false)
|
||||
})
|
||||
|
||||
it('should update multiple settings at once', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleSettingsChange({
|
||||
top_k: 10,
|
||||
score_threshold: 0.9,
|
||||
score_threshold_enabled: true,
|
||||
})
|
||||
})
|
||||
|
||||
expect(result.current.topK).toBe(10)
|
||||
expect(result.current.scoreThreshold).toBe(0.9)
|
||||
expect(result.current.scoreThresholdEnabled).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe('Summary Index Setting Handler', () => {
|
||||
it('should update summary index setting', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleSummaryIndexSettingChange({
|
||||
enable: true,
|
||||
})
|
||||
})
|
||||
|
||||
expect(result.current.summaryIndexSetting).toMatchObject({
|
||||
enable: true,
|
||||
})
|
||||
})
|
||||
|
||||
it('should merge with existing settings', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.handleSummaryIndexSettingChange({
|
||||
enable: true,
|
||||
})
|
||||
})
|
||||
|
||||
act(() => {
|
||||
result.current.handleSummaryIndexSettingChange({
|
||||
model_provider_name: 'openai',
|
||||
model_name: 'gpt-4',
|
||||
})
|
||||
})
|
||||
|
||||
expect(result.current.summaryIndexSetting).toMatchObject({
|
||||
enable: true,
|
||||
model_provider_name: 'openai',
|
||||
model_name: 'gpt-4',
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('handleSave', () => {
|
||||
it('should show error toast when name is empty', async () => {
|
||||
const Toast = await import('@/app/components/base/toast')
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setName('')
|
||||
})
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
expect(Toast.default.notify).toHaveBeenCalledWith({
|
||||
type: 'error',
|
||||
message: expect.any(String),
|
||||
})
|
||||
})
|
||||
|
||||
it('should show error toast when name is whitespace only', async () => {
|
||||
const Toast = await import('@/app/components/base/toast')
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setName(' ')
|
||||
})
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
expect(Toast.default.notify).toHaveBeenCalledWith({
|
||||
type: 'error',
|
||||
message: expect.any(String),
|
||||
})
|
||||
})
|
||||
|
||||
it('should call updateDatasetSetting with correct params', async () => {
|
||||
const { updateDatasetSetting } = await import('@/service/datasets')
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
expect(updateDatasetSetting).toHaveBeenCalledWith({
|
||||
datasetId: 'dataset-1',
|
||||
body: expect.objectContaining({
|
||||
name: 'Test Dataset',
|
||||
description: 'Test description',
|
||||
permission: DatasetPermission.onlyMe,
|
||||
}),
|
||||
})
|
||||
})
|
||||
|
||||
it('should show success toast on successful save', async () => {
|
||||
const Toast = await import('@/app/components/base/toast')
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
await waitFor(() => {
|
||||
expect(Toast.default.notify).toHaveBeenCalledWith({
|
||||
type: 'success',
|
||||
message: expect.any(String),
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
it('should call mutateDatasets after successful save', async () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockMutateDatasets).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
it('should call invalidDatasetList after successful save', async () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockInvalidDatasetList).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
it('should set loading to true during save', async () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
expect(result.current.loading).toBe(false)
|
||||
|
||||
const savePromise = act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
// Loading should be true during the save operation
|
||||
await savePromise
|
||||
|
||||
expect(result.current.loading).toBe(false) // After completion
|
||||
})
|
||||
|
||||
it('should not save when already loading', async () => {
|
||||
const { updateDatasetSetting } = await import('@/service/datasets')
|
||||
vi.mocked(updateDatasetSetting).mockImplementation(() => new Promise(resolve => setTimeout(resolve, 100)))
|
||||
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
// Start first save
|
||||
act(() => {
|
||||
result.current.handleSave()
|
||||
})
|
||||
|
||||
// Try to start second save immediately
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
// Should only have been called once
|
||||
expect(updateDatasetSetting).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('should show error toast on save failure', async () => {
|
||||
const { updateDatasetSetting } = await import('@/service/datasets')
|
||||
const Toast = await import('@/app/components/base/toast')
|
||||
vi.mocked(updateDatasetSetting).mockRejectedValueOnce(new Error('Network error'))
|
||||
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
expect(Toast.default.notify).toHaveBeenCalledWith({
|
||||
type: 'error',
|
||||
message: expect.any(String),
|
||||
})
|
||||
})
|
||||
|
||||
it('should include partial_member_list when permission is partialMembers', async () => {
|
||||
const { updateDatasetSetting } = await import('@/service/datasets')
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setPermission(DatasetPermission.partialMembers)
|
||||
result.current.setSelectedMemberIDs(['user-1', 'user-2'])
|
||||
})
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
expect(updateDatasetSetting).toHaveBeenCalledWith({
|
||||
datasetId: 'dataset-1',
|
||||
body: expect.objectContaining({
|
||||
partial_member_list: expect.arrayContaining([
|
||||
expect.objectContaining({ user_id: 'user-1' }),
|
||||
expect.objectContaining({ user_id: 'user-2' }),
|
||||
]),
|
||||
}),
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('Embedding Model', () => {
|
||||
it('should initialize embedding model from dataset', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
expect(result.current.embeddingModel).toEqual({
|
||||
provider: 'openai',
|
||||
model: 'text-embedding-ada-002',
|
||||
})
|
||||
})
|
||||
|
||||
it('should update embedding model when setEmbeddingModel is called', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
act(() => {
|
||||
result.current.setEmbeddingModel({
|
||||
provider: 'cohere',
|
||||
model: 'embed-english-v3.0',
|
||||
})
|
||||
})
|
||||
|
||||
expect(result.current.embeddingModel).toEqual({
|
||||
provider: 'cohere',
|
||||
model: 'embed-english-v3.0',
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('Retrieval Config', () => {
|
||||
it('should initialize retrieval config from dataset', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
expect(result.current.retrievalConfig).toBeDefined()
|
||||
expect(result.current.retrievalConfig.search_method).toBe(RETRIEVE_METHOD.semantic)
|
||||
})
|
||||
|
||||
it('should update retrieval config when setRetrievalConfig is called', () => {
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
const newConfig: RetrievalConfig = {
|
||||
...result.current.retrievalConfig,
|
||||
reranking_enable: true,
|
||||
}
|
||||
|
||||
act(() => {
|
||||
result.current.setRetrievalConfig(newConfig)
|
||||
})
|
||||
|
||||
expect(result.current.retrievalConfig.reranking_enable).toBe(true)
|
||||
})
|
||||
|
||||
it('should include weights in save request when weights are set', async () => {
|
||||
const { updateDatasetSetting } = await import('@/service/datasets')
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
// Set retrieval config with weights
|
||||
const configWithWeights: RetrievalConfig = {
|
||||
...result.current.retrievalConfig,
|
||||
search_method: RETRIEVE_METHOD.hybrid,
|
||||
weights: {
|
||||
weight_type: WeightedScoreEnum.Customized,
|
||||
vector_setting: {
|
||||
vector_weight: 0.7,
|
||||
embedding_provider_name: '',
|
||||
embedding_model_name: '',
|
||||
},
|
||||
keyword_setting: {
|
||||
keyword_weight: 0.3,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
act(() => {
|
||||
result.current.setRetrievalConfig(configWithWeights)
|
||||
})
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
// Verify that weights were included and embedding model info was added
|
||||
expect(updateDatasetSetting).toHaveBeenCalledWith({
|
||||
datasetId: 'dataset-1',
|
||||
body: expect.objectContaining({
|
||||
retrieval_model: expect.objectContaining({
|
||||
weights: expect.objectContaining({
|
||||
vector_setting: expect.objectContaining({
|
||||
embedding_provider_name: 'openai',
|
||||
embedding_model_name: 'text-embedding-ada-002',
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe('External Provider', () => {
|
||||
beforeEach(() => {
|
||||
// Update mock dataset to be external provider
|
||||
mockDataset = {
|
||||
...mockDataset,
|
||||
provider: 'external',
|
||||
external_knowledge_info: {
|
||||
external_knowledge_id: 'ext-123',
|
||||
external_knowledge_api_id: 'api-456',
|
||||
external_knowledge_api_name: 'External API',
|
||||
external_knowledge_api_endpoint: 'https://api.example.com',
|
||||
},
|
||||
external_retrieval_model: {
|
||||
top_k: 5,
|
||||
score_threshold: 0.8,
|
||||
score_threshold_enabled: true,
|
||||
},
|
||||
}
|
||||
})
|
||||
|
||||
it('should include external knowledge info in save request for external provider', async () => {
|
||||
const { updateDatasetSetting } = await import('@/service/datasets')
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
expect(updateDatasetSetting).toHaveBeenCalledWith({
|
||||
datasetId: 'dataset-1',
|
||||
body: expect.objectContaining({
|
||||
external_knowledge_id: 'ext-123',
|
||||
external_knowledge_api_id: 'api-456',
|
||||
external_retrieval_model: expect.objectContaining({
|
||||
top_k: expect.any(Number),
|
||||
score_threshold: expect.any(Number),
|
||||
score_threshold_enabled: expect.any(Boolean),
|
||||
}),
|
||||
}),
|
||||
})
|
||||
})
|
||||
|
||||
it('should use correct external retrieval settings', async () => {
|
||||
const { updateDatasetSetting } = await import('@/service/datasets')
|
||||
const { result } = renderHook(() => useFormState())
|
||||
|
||||
// Update external retrieval settings
|
||||
act(() => {
|
||||
result.current.handleSettingsChange({
|
||||
top_k: 10,
|
||||
score_threshold: 0.9,
|
||||
score_threshold_enabled: false,
|
||||
})
|
||||
})
|
||||
|
||||
await act(async () => {
|
||||
await result.current.handleSave()
|
||||
})
|
||||
|
||||
expect(updateDatasetSetting).toHaveBeenCalledWith({
|
||||
datasetId: 'dataset-1',
|
||||
body: expect.objectContaining({
|
||||
external_retrieval_model: {
|
||||
top_k: 10,
|
||||
score_threshold: 0.9,
|
||||
score_threshold_enabled: false,
|
||||
},
|
||||
}),
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
@ -0,0 +1,264 @@
|
||||
'use client'
|
||||
import type { AppIconSelection } from '@/app/components/base/app-icon-picker'
|
||||
import type { DefaultModel } from '@/app/components/header/account-setting/model-provider-page/declarations'
|
||||
import type { Member } from '@/models/common'
|
||||
import type { IconInfo, SummaryIndexSetting as SummaryIndexSettingType } from '@/models/datasets'
|
||||
import type { RetrievalConfig } from '@/types/app'
|
||||
import { useCallback, useMemo, useRef, useState } from 'react'
|
||||
import { useTranslation } from 'react-i18next'
|
||||
import Toast from '@/app/components/base/toast'
|
||||
import { isReRankModelSelected } from '@/app/components/datasets/common/check-rerank-model'
|
||||
import { ModelTypeEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
|
||||
import { useModelList } from '@/app/components/header/account-setting/model-provider-page/hooks'
|
||||
import { useSelector as useAppContextWithSelector } from '@/context/app-context'
|
||||
import { useDatasetDetailContextWithSelector } from '@/context/dataset-detail'
|
||||
import { DatasetPermission } from '@/models/datasets'
|
||||
import { updateDatasetSetting } from '@/service/datasets'
|
||||
import { useInvalidDatasetList } from '@/service/knowledge/use-dataset'
|
||||
import { useMembers } from '@/service/use-common'
|
||||
import { checkShowMultiModalTip } from '../../utils'
|
||||
|
||||
const DEFAULT_APP_ICON: IconInfo = {
|
||||
icon_type: 'emoji',
|
||||
icon: '📙',
|
||||
icon_background: '#FFF4ED',
|
||||
icon_url: '',
|
||||
}
|
||||
|
||||
export const useFormState = () => {
|
||||
const { t } = useTranslation()
|
||||
const isCurrentWorkspaceDatasetOperator = useAppContextWithSelector(state => state.isCurrentWorkspaceDatasetOperator)
|
||||
const currentDataset = useDatasetDetailContextWithSelector(state => state.dataset)
|
||||
const mutateDatasets = useDatasetDetailContextWithSelector(state => state.mutateDatasetRes)
|
||||
|
||||
// Basic form state
|
||||
const [loading, setLoading] = useState(false)
|
||||
const [name, setName] = useState(currentDataset?.name ?? '')
|
||||
const [description, setDescription] = useState(currentDataset?.description ?? '')
|
||||
|
||||
// Icon state
|
||||
const [iconInfo, setIconInfo] = useState(currentDataset?.icon_info || DEFAULT_APP_ICON)
|
||||
const [showAppIconPicker, setShowAppIconPicker] = useState(false)
|
||||
const previousAppIcon = useRef(DEFAULT_APP_ICON)
|
||||
|
||||
// Permission state
|
||||
const [permission, setPermission] = useState(currentDataset?.permission)
|
||||
const [selectedMemberIDs, setSelectedMemberIDs] = useState<string[]>(currentDataset?.partial_member_list || [])
|
||||
|
||||
// External retrieval state
|
||||
const [topK, setTopK] = useState(currentDataset?.external_retrieval_model.top_k ?? 2)
|
||||
const [scoreThreshold, setScoreThreshold] = useState(currentDataset?.external_retrieval_model.score_threshold ?? 0.5)
|
||||
const [scoreThresholdEnabled, setScoreThresholdEnabled] = useState(currentDataset?.external_retrieval_model.score_threshold_enabled ?? false)
|
||||
|
||||
// Indexing and retrieval state
|
||||
const [indexMethod, setIndexMethod] = useState(currentDataset?.indexing_technique)
|
||||
const [keywordNumber, setKeywordNumber] = useState(currentDataset?.keyword_number ?? 10)
|
||||
const [retrievalConfig, setRetrievalConfig] = useState(currentDataset?.retrieval_model_dict as RetrievalConfig)
|
||||
const [embeddingModel, setEmbeddingModel] = useState<DefaultModel>(
|
||||
currentDataset?.embedding_model
|
||||
? {
|
||||
provider: currentDataset.embedding_model_provider,
|
||||
model: currentDataset.embedding_model,
|
||||
}
|
||||
: {
|
||||
provider: '',
|
||||
model: '',
|
||||
},
|
||||
)
|
||||
|
||||
// Summary index state
|
||||
const [summaryIndexSetting, setSummaryIndexSetting] = useState(currentDataset?.summary_index_setting)
|
||||
|
||||
// Model lists
|
||||
const { data: rerankModelList } = useModelList(ModelTypeEnum.rerank)
|
||||
const { data: embeddingModelList } = useModelList(ModelTypeEnum.textEmbedding)
|
||||
const { data: membersData } = useMembers()
|
||||
const invalidDatasetList = useInvalidDatasetList()
|
||||
|
||||
// Derive member list from API data
|
||||
const memberList = useMemo<Member[]>(() => {
|
||||
return membersData?.accounts ?? []
|
||||
}, [membersData])
|
||||
|
||||
// Icon handlers
|
||||
const handleOpenAppIconPicker = useCallback(() => {
|
||||
setShowAppIconPicker(true)
|
||||
previousAppIcon.current = iconInfo
|
||||
}, [iconInfo])
|
||||
|
||||
const handleSelectAppIcon = useCallback((icon: AppIconSelection) => {
|
||||
const newIconInfo: IconInfo = {
|
||||
icon_type: icon.type,
|
||||
icon: icon.type === 'emoji' ? icon.icon : icon.fileId,
|
||||
icon_background: icon.type === 'emoji' ? icon.background : undefined,
|
||||
icon_url: icon.type === 'emoji' ? undefined : icon.url,
|
||||
}
|
||||
setIconInfo(newIconInfo)
|
||||
setShowAppIconPicker(false)
|
||||
}, [])
|
||||
|
||||
const handleCloseAppIconPicker = useCallback(() => {
|
||||
setIconInfo(previousAppIcon.current)
|
||||
setShowAppIconPicker(false)
|
||||
}, [])
|
||||
|
||||
// External retrieval settings handler
|
||||
const handleSettingsChange = useCallback((data: { top_k?: number, score_threshold?: number, score_threshold_enabled?: boolean }) => {
|
||||
if (data.top_k !== undefined)
|
||||
setTopK(data.top_k)
|
||||
if (data.score_threshold !== undefined)
|
||||
setScoreThreshold(data.score_threshold)
|
||||
if (data.score_threshold_enabled !== undefined)
|
||||
setScoreThresholdEnabled(data.score_threshold_enabled)
|
||||
}, [])
|
||||
|
||||
// Summary index setting handler
|
||||
const handleSummaryIndexSettingChange = useCallback((payload: SummaryIndexSettingType) => {
|
||||
setSummaryIndexSetting(prev => ({ ...prev, ...payload }))
|
||||
}, [])
|
||||
|
||||
// Save handler
|
||||
const handleSave = async () => {
|
||||
if (loading)
|
||||
return
|
||||
|
||||
if (!name?.trim()) {
|
||||
Toast.notify({ type: 'error', message: t('form.nameError', { ns: 'datasetSettings' }) })
|
||||
return
|
||||
}
|
||||
|
||||
if (!isReRankModelSelected({ rerankModelList, retrievalConfig, indexMethod })) {
|
||||
Toast.notify({ type: 'error', message: t('datasetConfig.rerankModelRequired', { ns: 'appDebug' }) })
|
||||
return
|
||||
}
|
||||
|
||||
if (retrievalConfig.weights) {
|
||||
retrievalConfig.weights.vector_setting.embedding_provider_name = embeddingModel.provider || ''
|
||||
retrievalConfig.weights.vector_setting.embedding_model_name = embeddingModel.model || ''
|
||||
}
|
||||
|
||||
try {
|
||||
setLoading(true)
|
||||
const body: Record<string, unknown> = {
|
||||
name,
|
||||
icon_info: iconInfo,
|
||||
doc_form: currentDataset?.doc_form,
|
||||
description,
|
||||
permission,
|
||||
indexing_technique: indexMethod,
|
||||
retrieval_model: {
|
||||
...retrievalConfig,
|
||||
score_threshold: retrievalConfig.score_threshold_enabled ? retrievalConfig.score_threshold : 0,
|
||||
},
|
||||
embedding_model: embeddingModel.model,
|
||||
embedding_model_provider: embeddingModel.provider,
|
||||
keyword_number: keywordNumber,
|
||||
summary_index_setting: summaryIndexSetting,
|
||||
}
|
||||
|
||||
if (currentDataset!.provider === 'external') {
|
||||
body.external_knowledge_id = currentDataset!.external_knowledge_info.external_knowledge_id
|
||||
body.external_knowledge_api_id = currentDataset!.external_knowledge_info.external_knowledge_api_id
|
||||
body.external_retrieval_model = {
|
||||
top_k: topK,
|
||||
score_threshold: scoreThreshold,
|
||||
score_threshold_enabled: scoreThresholdEnabled,
|
||||
}
|
||||
}
|
||||
|
||||
if (permission === DatasetPermission.partialMembers) {
|
||||
body.partial_member_list = selectedMemberIDs.map((id) => {
|
||||
return {
|
||||
user_id: id,
|
||||
role: memberList.find(member => member.id === id)?.role,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
await updateDatasetSetting({ datasetId: currentDataset!.id, body })
|
||||
Toast.notify({ type: 'success', message: t('actionMsg.modifiedSuccessfully', { ns: 'common' }) })
|
||||
|
||||
if (mutateDatasets) {
|
||||
await mutateDatasets()
|
||||
invalidDatasetList()
|
||||
}
|
||||
}
|
||||
catch {
|
||||
Toast.notify({ type: 'error', message: t('actionMsg.modifiedUnsuccessfully', { ns: 'common' }) })
|
||||
}
|
||||
finally {
|
||||
setLoading(false)
|
||||
}
|
||||
}
|
||||
|
||||
// Computed values
|
||||
const showMultiModalTip = useMemo(() => {
|
||||
return checkShowMultiModalTip({
|
||||
embeddingModel,
|
||||
rerankingEnable: retrievalConfig.reranking_enable,
|
||||
rerankModel: {
|
||||
rerankingProviderName: retrievalConfig.reranking_model.reranking_provider_name,
|
||||
rerankingModelName: retrievalConfig.reranking_model.reranking_model_name,
|
||||
},
|
||||
indexMethod,
|
||||
embeddingModelList,
|
||||
rerankModelList,
|
||||
})
|
||||
}, [embeddingModel, rerankModelList, retrievalConfig.reranking_enable, retrievalConfig.reranking_model, embeddingModelList, indexMethod])
|
||||
|
||||
return {
|
||||
// Context values
|
||||
currentDataset,
|
||||
isCurrentWorkspaceDatasetOperator,
|
||||
|
||||
// Loading state
|
||||
loading,
|
||||
|
||||
// Basic form
|
||||
name,
|
||||
setName,
|
||||
description,
|
||||
setDescription,
|
||||
|
||||
// Icon
|
||||
iconInfo,
|
||||
showAppIconPicker,
|
||||
handleOpenAppIconPicker,
|
||||
handleSelectAppIcon,
|
||||
handleCloseAppIconPicker,
|
||||
|
||||
// Permission
|
||||
permission,
|
||||
setPermission,
|
||||
selectedMemberIDs,
|
||||
setSelectedMemberIDs,
|
||||
memberList,
|
||||
|
||||
// External retrieval
|
||||
topK,
|
||||
scoreThreshold,
|
||||
scoreThresholdEnabled,
|
||||
handleSettingsChange,
|
||||
|
||||
// Indexing and retrieval
|
||||
indexMethod,
|
||||
setIndexMethod,
|
||||
keywordNumber,
|
||||
setKeywordNumber,
|
||||
retrievalConfig,
|
||||
setRetrievalConfig,
|
||||
embeddingModel,
|
||||
setEmbeddingModel,
|
||||
embeddingModelList,
|
||||
|
||||
// Summary index
|
||||
summaryIndexSetting,
|
||||
handleSummaryIndexSettingChange,
|
||||
|
||||
// Computed
|
||||
showMultiModalTip,
|
||||
|
||||
// Actions
|
||||
handleSave,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user