From a2194300ccc33e2d2d3a1dcecafaef457a51edf4 Mon Sep 17 00:00:00 2001 From: Stephen Zhou <38493346+hyoban@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:35:24 +0800 Subject: [PATCH] fix(knowledge-fs): stabilize initial source creation --- .../knowledge_fs_initial_source_tasks.py | 10 +- .../test_knowledge_fs_initial_source_tasks.py | 49 +++- .../__tests__/create-knowledge-page.spec.tsx | 276 ++++++++++++++++++ .../new-rag/create-connected-source-setup.tsx | 144 ++++++--- web/features/new-rag/create-source-setup.tsx | 14 +- 5 files changed, 441 insertions(+), 52 deletions(-) diff --git a/api/tasks/knowledge_fs_initial_source_tasks.py b/api/tasks/knowledge_fs_initial_source_tasks.py index 9b273d652a6..5af888d9597 100644 --- a/api/tasks/knowledge_fs_initial_source_tasks.py +++ b/api/tasks/knowledge_fs_initial_source_tasks.py @@ -32,7 +32,7 @@ from services.knowledge_fs.product_dto import ( KnowledgeFSSourceUpdatePayload, KnowledgeFSSourceWorkflowImportPayload, ) -from services.knowledge_fs.product_remote import KnowledgeFSProductResourceNotFoundError +from services.knowledge_fs.product_remote import KnowledgeFSProductRemoteError, KnowledgeFSProductResourceNotFoundError from services.knowledge_fs.runtime import get_knowledge_fs_runtime _LEGACY_WEBSITE_PLUGIN_IDS = { @@ -554,6 +554,10 @@ def _run_initial_source_task( }, ) raise task.retry(exc=exc) + except KnowledgeFSProductResourceNotFoundError: + raise + except KnowledgeFSProductRemoteError as exc: + raise task.retry(exc=exc) @shared_task(bind=True, queue="knowledge_fs_lifecycle", max_retries=180, default_retry_delay=2) @@ -614,6 +618,10 @@ def import_initial_website_source( }, ) raise self.retry(exc=exc) + except KnowledgeFSProductResourceNotFoundError: + raise + except KnowledgeFSProductRemoteError as exc: + raise self.retry(exc=exc) __all__ = [ diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py index 9e3dca2aaa6..e9b0d2f383f 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py @@ -10,7 +10,7 @@ from services.knowledge_fs.product_dto import ( KnowledgeFSInitialOnlineDriveSourcePayload, KnowledgeFSInitialWebsiteSourcePayload, ) -from services.knowledge_fs.product_remote import KnowledgeFSProductResourceNotFoundError +from services.knowledge_fs.product_remote import KnowledgeFSProductRemoteError, KnowledgeFSProductResourceNotFoundError from tasks.knowledge_fs_initial_source_tasks import ( KnowledgeFSInitialSourceNotReadyError, import_initial_source, @@ -693,3 +693,50 @@ def test_initial_source_task_validates_discriminated_connector_payload() -> None parsed_payload = start_import.call_args.kwargs["payload"] assert isinstance(parsed_payload, KnowledgeFSInitialOnlineDocumentSourcePayload) assert parsed_payload.credential_id == "notion-credential-1" + + +def test_initial_source_task_retries_transient_remote_error() -> None: + serialized_payload = _drive_payload().model_dump(mode="json", by_alias=True) + remote_error = KnowledgeFSProductRemoteError("temporary outage") + retry_error = RuntimeError("retry requested") + with ( + patch( + "tasks.knowledge_fs_initial_source_tasks.start_initial_source_import", + side_effect=remote_error, + ), + patch.object(import_initial_source, "retry", side_effect=retry_error) as retry, + pytest.raises(RuntimeError, match="retry requested"), + ): + import_initial_source.run( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + operation_id="operation-1", + payload=serialized_payload, + workflow_id="workflow-1", + ) + + retry.assert_called_once_with(exc=remote_error) + + +def test_initial_source_task_does_not_retry_authoritative_missing_resource() -> None: + serialized_payload = _drive_payload().model_dump(mode="json", by_alias=True) + missing_resource = KnowledgeFSProductResourceNotFoundError("workflow was not found") + with ( + patch( + "tasks.knowledge_fs_initial_source_tasks.start_initial_source_import", + side_effect=missing_resource, + ), + patch.object(import_initial_source, "retry") as retry, + pytest.raises(KnowledgeFSProductResourceNotFoundError, match="workflow was not found"), + ): + import_initial_source.run( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + operation_id="operation-1", + payload=serialized_payload, + workflow_id="workflow-1", + ) + + retry.assert_not_called() diff --git a/web/features/new-rag/__tests__/create-knowledge-page.spec.tsx b/web/features/new-rag/__tests__/create-knowledge-page.spec.tsx index 8ee890b0ea8..8dbbb432a2e 100644 --- a/web/features/new-rag/__tests__/create-knowledge-page.spec.tsx +++ b/web/features/new-rag/__tests__/create-knowledge-page.spec.tsx @@ -239,6 +239,37 @@ const firecrawlDatasourceAuth = { provider: 'firecrawl', } +const jinaDatasourcePlugin = { + ...firecrawlDatasourcePlugin, + declaration: { + ...firecrawlDatasourcePlugin.declaration, + identity: { + ...firecrawlDatasourcePlugin.declaration.identity, + label: { en_US: 'Jina Reader', zh_Hans: 'Jina Reader' }, + name: 'jinareader', + }, + }, + plugin_id: 'langgenius/jina_datasource', + plugin_unique_identifier: 'langgenius/jina_datasource:1.0.0@local', + provider: 'jinareader', +} + +const jinaDatasourceAuth = { + ...firecrawlDatasourceAuth, + credentials_list: [ + { + ...firecrawlDatasourceAuth.credentials_list[0], + id: 'jina-credential-1', + name: 'Default Jina Reader', + }, + ], + label: { en_US: 'Jina Reader' }, + name: 'jinareader', + plugin_id: 'langgenius/jina_datasource', + plugin_unique_identifier: 'langgenius/jina_datasource:1.0.0@local', + provider: 'jinareader', +} + const notionDatasourcePlugin = { ...firecrawlDatasourcePlugin, declaration: { @@ -1107,6 +1138,54 @@ describe('CreateKnowledgePage', () => { expect(screen.getByText('Getting started')).toBeInTheDocument() }) + it('creates an initial source from a synchronous Jina Reader preview', async () => { + const user = userEvent.setup() + navigationMock.startMode = 'source' + datasourceQueryMock.plugins.data = [firecrawlDatasourcePlugin, jinaDatasourcePlugin] + datasourceQueryMock.auth.data = { result: [firecrawlDatasourceAuth, jinaDatasourceAuth] } + serviceMock.createCrawl.mockResolvedValueOnce({ + data: { + content: '# Dify introduction', + description: 'Introduction', + title: 'Dify introduction', + url: 'https://docs.dify.ai/introduction', + }, + }) + renderPage() + await fillRequiredFields(user) + await user.click(screen.getByRole('radio', { name: 'Jina Reader' })) + await user.type( + screen.getByPlaceholderText('dataset.newKnowledge.rootUrlPlaceholder'), + 'https://docs.dify.ai/introduction', + ) + await user.type( + screen.getByPlaceholderText('dataset.newKnowledge.sourceNamePlaceholder'), + 'Dify introduction', + ) + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.crawlAndPreview' })) + + await user.click(await screen.findByRole('checkbox', { name: 'Dify introduction' })) + expect(serviceMock.getCrawlStatus).not.toHaveBeenCalled() + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.createTitle' })) + + await waitFor(() => expect(serviceMock.create).toHaveBeenCalledOnce()) + expect(serviceMock.create).toHaveBeenCalledWith({ + body: expect.objectContaining({ + initial_source: expect.objectContaining({ + credentialId: 'jina-credential-1', + pluginId: 'langgenius/jina_datasource', + provider: 'jinareader', + selection: [ + expect.objectContaining({ + source_url: 'https://docs.dify.ai/introduction', + title: 'Dify introduction', + }), + ], + }), + }), + }) + }) + it('shows and can stop an ongoing website crawl', async () => { const user = userEvent.setup() navigationMock.startMode = 'source' @@ -1338,6 +1417,203 @@ describe('CreateKnowledgePage', () => { }) }) + it('uses the online-drive transport when creating with Google Docs', async () => { + const user = userEvent.setup() + navigationMock.startMode = 'source' + datasourceQueryMock.plugins.data = [ + firecrawlDatasourcePlugin, + notionDatasourcePlugin, + googleDriveDatasourcePlugin, + ] + datasourceQueryMock.auth.data = { + result: [firecrawlDatasourceAuth, notionDatasourceAuth, googleDriveDatasourceAuth], + } + serviceMock.previewInitialSource.mockResolvedValue({ + files: [ + { + bucket: null, + id: 'doc-1', + mime_type: 'application/vnd.google-apps.document', + name: 'Launch plan', + provider_item_id: '["","doc-1"]', + size: 1024, + type: 'application/vnd.google-apps.document', + }, + ], + kind: 'online_drive', + next_page_parameters: null, + }) + renderPage() + await fillRequiredFields(user) + await user.click(screen.getByRole('radio', { name: 'dataset.newKnowledge.onlineDocuments' })) + await user.click(screen.getByRole('radio', { name: 'Google Docs' })) + await user.type( + screen.getByPlaceholderText('dataset.newKnowledge.sourceNamePlaceholder'), + 'Team docs', + ) + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.preview' })) + + await waitFor(() => + expect(serviceMock.previewInitialSource).toHaveBeenCalledWith({ + body: expect.objectContaining({ kind: 'online_drive' }), + }), + ) + await user.click(await screen.findByRole('checkbox', { name: 'Launch plan' })) + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.createTitle' })) + + await waitFor(() => expect(serviceMock.create).toHaveBeenCalledOnce()) + expect(serviceMock.create).toHaveBeenCalledWith({ + body: expect.objectContaining({ + initial_source: expect.objectContaining({ + kind: 'online_drive', + selection: [expect.objectContaining({ id: 'doc-1' })], + }), + }), + }) + }) + + it('preserves root pagination while expanding a drive folder', async () => { + const user = userEvent.setup() + navigationMock.startMode = 'source' + datasourceQueryMock.plugins.data = [firecrawlDatasourcePlugin, googleDriveDatasourcePlugin] + datasourceQueryMock.auth.data = { + result: [firecrawlDatasourceAuth, googleDriveDatasourceAuth], + } + serviceMock.previewInitialSource.mockImplementation( + ({ body }: { body: { parameters: Record } }) => { + if (body.parameters.prefix === 'folder-1') { + return Promise.resolve({ + files: [ + { + bucket: null, + id: 'child-1', + mime_type: 'application/pdf', + name: 'Folder child.pdf', + provider_item_id: '["","child-1"]', + size: 128, + type: 'application/pdf', + }, + ], + kind: 'online_drive', + next_page_parameters: null, + }) + } + if (body.parameters.next_page_parameters) { + return Promise.resolve({ + files: [ + { + bucket: null, + id: 'root-2', + mime_type: 'application/pdf', + name: 'Second root file.pdf', + provider_item_id: '["","root-2"]', + size: 256, + type: 'application/pdf', + }, + ], + kind: 'online_drive', + next_page_parameters: null, + }) + } + return Promise.resolve({ + files: [ + { + bucket: null, + id: 'folder-1', + mime_type: null, + name: 'Plans', + provider_item_id: '["","folder-1"]', + size: 0, + type: 'folder', + }, + ], + kind: 'online_drive', + next_page_parameters: { cursor: 'root-next' }, + }) + }, + ) + renderPage() + await user.click(screen.getByRole('radio', { name: 'dataset.newKnowledge.onlineDrive' })) + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.preview' })) + + const folderButton = await screen.findByRole('button', { name: 'Plans' }) + expect(folderButton).toHaveAttribute('aria-expanded', 'false') + await user.click(folderButton) + expect(folderButton).toHaveAttribute('aria-expanded', 'true') + expect(await screen.findByRole('checkbox', { name: 'Folder child.pdf' })).toBeInTheDocument() + + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.loadMore' })) + + expect( + await screen.findByRole('checkbox', { name: 'Second root file.pdf' }), + ).toBeInTheDocument() + expect(serviceMock.previewInitialSource).toHaveBeenLastCalledWith({ + body: expect.objectContaining({ + parameters: { next_page_parameters: { cursor: 'root-next' } }, + }), + }) + }) + + it('limits a paginated drive selection to the backend maximum', async () => { + const user = userEvent.setup() + navigationMock.startMode = 'source' + datasourceQueryMock.plugins.data = [firecrawlDatasourcePlugin, googleDriveDatasourcePlugin] + datasourceQueryMock.auth.data = { + result: [firecrawlDatasourceAuth, googleDriveDatasourceAuth], + } + const files = Array.from({ length: 200 }, (_, index) => ({ + bucket: null, + id: `file-${index + 1}`, + mime_type: 'application/pdf', + name: `File ${index + 1}.pdf`, + provider_item_id: `["","file-${index + 1}"]`, + size: 128, + type: 'application/pdf', + })) + serviceMock.previewInitialSource + .mockResolvedValueOnce({ + files, + kind: 'online_drive', + next_page_parameters: { cursor: 'next' }, + }) + .mockResolvedValueOnce({ + files: [ + { + bucket: null, + id: 'file-201', + mime_type: 'application/pdf', + name: 'File 201.pdf', + provider_item_id: '["","file-201"]', + size: 128, + type: 'application/pdf', + }, + ], + kind: 'online_drive', + next_page_parameters: null, + }) + renderPage() + await fillRequiredFields(user) + await user.click(screen.getByRole('radio', { name: 'dataset.newKnowledge.onlineDrive' })) + await user.type( + screen.getByPlaceholderText('dataset.newKnowledge.sourceNamePlaceholder'), + 'Drive archive', + ) + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.preview' })) + await user.click(await screen.findByRole('button', { name: 'dataset.newKnowledge.loadMore' })) + await screen.findByRole('checkbox', { name: 'File 201.pdf' }) + await user.click(screen.getByRole('checkbox', { name: 'dataset.newKnowledge.selectAll' })) + + expect(screen.getByRole('checkbox', { name: 'File 201.pdf' })).toHaveAttribute( + 'aria-disabled', + 'true', + ) + expect(screen.getByText(/^dataset\.newKnowledge\.pagesSelected/)).toHaveTextContent('200') + await user.click(screen.getByRole('button', { name: 'dataset.newKnowledge.createTitle' })) + + await waitFor(() => expect(serviceMock.create).toHaveBeenCalledOnce()) + expect(serviceMock.create.mock.calls[0]?.[0].body.initial_source.selection).toHaveLength(200) + }) + it('requires a selected website preview page before creating with an initial source', async () => { const user = userEvent.setup() navigationMock.startMode = 'source' diff --git a/web/features/new-rag/create-connected-source-setup.tsx b/web/features/new-rag/create-connected-source-setup.tsx index 52c9cbbdd26..90af5a37e62 100644 --- a/web/features/new-rag/create-connected-source-setup.tsx +++ b/web/features/new-rag/create-connected-source-setup.tsx @@ -41,6 +41,10 @@ type NextPageRequest = { prefix?: string } +const MAX_SELECTION = 200 +const ROOT_PAGE_SCOPE = 'root' +const SELECTION_LIMIT_ID = 'create-connected-source-selection-limit' + function isDriveContainer(file: PreviewFile) { return /bucket|directory|folder|workspace/i.test(file.type) } @@ -80,7 +84,11 @@ export function CreateConnectedSourceSetup({ const [loadingMore, setLoadingMore] = useState(false) const [error, setError] = useState(false) const [previewed, setPreviewed] = useState(false) - const [nextPageRequest, setNextPageRequest] = useState() + const [nextPageRequests, setNextPageRequests] = useState>( + () => new Map(), + ) + const driveTransport = providerOption.providerType === 'online_drive' + const selectionAtLimit = selected.size >= MAX_SELECTION const selectableResources = useMemo( () => resources.filter( @@ -113,7 +121,7 @@ export function CreateConnectedSourceSetup({ pluginId: providerOption.plugin.plugin_id, provider: providerOption.plugin.provider, } - if (draft.sourceType === 'onlineDocuments') { + if (!driveTransport) { onInitialSourceChange({ ...binding, kind: 'online_document', @@ -158,6 +166,7 @@ export function CreateConnectedSourceSetup({ }, [ credential.id, draft, + driveTransport, onInitialSourceChange, providerOption.datasource.identity.name, providerOption.plugin.plugin_id, @@ -189,7 +198,7 @@ export function CreateConnectedSourceSetup({ body: { credentialId: credential.id, datasource: providerOption.datasource.identity.name, - kind: draft.sourceType === 'onlineDocuments' ? 'online_document' : 'online_drive', + kind: driveTransport ? 'online_drive' : 'online_document', parameters: { ...(bucket ? { bucket } : {}), ...(prefix ? { prefix } : {}), @@ -199,39 +208,43 @@ export function CreateConnectedSourceSetup({ provider: providerOption.plugin.provider, }, }) - const nextResources: PreviewResource[] = - draft.sourceType === 'onlineDocuments' - ? (response.documents ?? []).map((document) => ({ - depth, - document, - key: `document:${document.provider_item_id}`, - kind: 'document' as const, - parentKey, - })) - : (response.files ?? []).map((file) => ({ - depth, - file, - key: `file:${file.provider_item_id}`, - kind: 'file' as const, - parentKey, - })) + const nextResources: PreviewResource[] = driveTransport + ? (response.files ?? []).map((file) => ({ + depth, + file, + key: `file:${file.provider_item_id}`, + kind: 'file' as const, + parentKey, + })) + : (response.documents ?? []).map((document) => ({ + depth, + document, + key: `document:${document.provider_item_id}`, + kind: 'document' as const, + parentKey, + })) setResources((current) => { const next = new Map((append ? current : []).map((resource) => [resource.key, resource])) for (const resource of nextResources) next.set(resource.key, resource) return [...next.values()] }) if (parentKey) setExpanded((current) => new Set(current).add(parentKey)) - setNextPageRequest( - response.next_page_parameters - ? { - bucket, - depth, - nextPage: response.next_page_parameters, - parentKey, - prefix, - } - : null, - ) + setNextPageRequests((current) => { + const next = append ? new Map(current) : new Map() + const scope = parentKey ?? ROOT_PAGE_SCOPE + if (response.next_page_parameters) { + next.set(scope, { + bucket, + depth, + nextPage: response.next_page_parameters, + parentKey, + prefix, + }) + } else { + next.delete(scope) + } + return next + }) setPreviewed(true) } catch { setError(true) @@ -240,14 +253,14 @@ export function CreateConnectedSourceSetup({ setLoadingMore(false) } }, - [credential.id, draft.sourceType, providerOption], + [credential.id, driveTransport, providerOption], ) const toggle = (key: string) => { setSelected((current) => { const next = new Set(current) if (next.has(key)) next.delete(key) - else next.add(key) + else if (next.size < MAX_SELECTION) next.add(key) return next }) } @@ -259,7 +272,10 @@ export function CreateConnectedSourceSetup({ selectableResources.every((resource) => current.has(resource.key)) for (const resource of selectableResources) { if (allSelected) next.delete(resource.key) - else next.add(resource.key) + else { + if (next.size >= MAX_SELECTION) break + next.add(resource.key) + } } return next }) @@ -281,6 +297,9 @@ export function CreateConnectedSourceSetup({ prefix: resource.file.id || undefined, }) } + const visibleNextPageRequests = [...nextPageRequests.entries()].filter( + ([scope]) => scope === ROOT_PAGE_SCOPE || expanded.has(scope), + ) return (
@@ -332,11 +351,21 @@ export function CreateConnectedSourceSetup({
$['newKnowledge.selectAll'])} + aria-describedby={selectionAtLimit ? SELECTION_LIMIT_ID : undefined} checked={ selectableResources.length > 0 && selectableResources.every((resource) => selected.has(resource.key)) } - disabled={disabled || !selectableResources.length} + disabled={ + disabled || + !selectableResources.length || + (selectionAtLimit && + !selectableResources.every((resource) => selected.has(resource.key))) + } + indeterminate={ + selected.size > 0 && + !selectableResources.every((resource) => selected.has(resource.key)) + } onCheckedChange={toggleAll} /> @@ -344,8 +373,13 @@ export function CreateConnectedSourceSetup({ ? t(($) => $['newKnowledge.selectPagesToSync']) : t(($) => $['newKnowledge.selectFilesAndFolders'])} - + {t(($) => $['newKnowledge.pagesSelected'], { count: selected.size })} + {selectionAtLimit && ( + + {t(($) => $['newKnowledge.maxPages'])}: {MAX_SELECTION} + + )}
    @@ -364,12 +398,13 @@ export function CreateConnectedSourceSetup({ size="small" className="size-5 px-0" aria-label={resourceName(resource)} + aria-expanded={expanded.has(resource.key)} disabled={disabled || loadingMore} onClick={() => expandContainer(resource)} > @@ -377,8 +412,13 @@ export function CreateConnectedSourceSetup({ ) : ( toggle(resource.key)} /> )} @@ -390,17 +430,25 @@ export function CreateConnectedSourceSetup({ ) })}
- {nextPageRequest && ( -
- + {visibleNextPageRequests.length > 0 && ( +
+ {visibleNextPageRequests.map(([scope, request]) => { + const parent = + scope === ROOT_PAGE_SCOPE ? undefined : resources.find(({ key }) => key === scope) + return ( + + ) + })}
)} diff --git a/web/features/new-rag/create-source-setup.tsx b/web/features/new-rag/create-source-setup.tsx index 49fda285fc8..66908709ad4 100644 --- a/web/features/new-rag/create-source-setup.tsx +++ b/web/features/new-rag/create-source-setup.tsx @@ -66,8 +66,12 @@ type LocalCrawlState = 'error' | 'idle' | 'running' | 'stopped' | 'success' type InitialSource = NonNullable function crawlPages(response: Record): CrawlResultItem[] { - if (!Array.isArray(response.data)) return [] - return response.data.flatMap((item) => { + const items = Array.isArray(response.data) + ? response.data + : response.data && typeof response.data === 'object' + ? [response.data] + : [] + return items.flatMap((item) => { if (!item || typeof item !== 'object') return [] const page = item as Record const sourceUrl = @@ -262,6 +266,12 @@ export function CreateSourceSetup({ }, url: draft.rootUrl, })) as Record + const synchronousPages = crawlPages(created) + if (synchronousPages.length) { + setPreviewPages(synchronousPages) + setCrawlState('success') + return + } const jobId = typeof created.job_id === 'string' ? created.job_id : undefined if (!jobId) throw new Error('Website crawl did not return a job id')