diff --git a/web/app/components/snippets/components/__tests__/snippet-main.spec.tsx b/web/app/components/snippets/components/__tests__/snippet-main.spec.tsx index 911491ce69..f1a4a6726a 100644 --- a/web/app/components/snippets/components/__tests__/snippet-main.spec.tsx +++ b/web/app/components/snippets/components/__tests__/snippet-main.spec.tsx @@ -3,6 +3,7 @@ import type { WorkflowProps } from '@/app/components/workflow' import type { SnippetDetailPayload, SnippetInputField } from '@/models/snippet' import { fireEvent, screen, waitFor } from '@testing-library/react' import { renderWorkflowComponent } from '@/app/components/workflow/__tests__/workflow-test-env' +import { BlockEnum } from '@/app/components/workflow/types' import { PipelineInputVarType } from '@/models/pipeline' import SnippetMain from '../snippet-main' @@ -20,6 +21,7 @@ const mockHandleStartWorkflowRun = vi.fn() const mockHandleStopRun = vi.fn() const mockHandleWorkflowStartRunInWorkflow = vi.fn() const mockHandleCheckBeforePublish = vi.fn() +const mockUseAvailableNodesMetaData = vi.hoisted(() => vi.fn()) const mockInspectVarsCrud = { hasNodeInspectVars: vi.fn(), hasSetInspectVar: vi.fn(), @@ -75,6 +77,10 @@ vi.mock('@/app/components/workflow/hooks/use-checklist', () => ({ }), })) +vi.mock('@/app/components/workflow-app/hooks', () => ({ + useAvailableNodesMetaData: () => mockUseAvailableNodesMetaData(), +})) + vi.mock('@/app/components/snippets/hooks/use-inspect-vars-crud', () => ({ useInspectVarsCrud: () => mockInspectVarsCrud, })) @@ -210,6 +216,14 @@ const renderSnippetMain = () => { ) } +const createNodeMetadata = (type: BlockEnum) => ({ + metaData: { + type, + }, + defaultValue: {}, + checkValid: vi.fn(), +}) + describe('SnippetMain', () => { beforeEach(() => { vi.clearAllMocks() @@ -222,6 +236,24 @@ describe('SnippetMain', () => { }, refetch: vi.fn(), }) + const llmNodeMetadata = createNodeMetadata(BlockEnum.LLM) + const humanInputNodeMetadata = createNodeMetadata(BlockEnum.HumanInput) + const endNodeMetadata = createNodeMetadata(BlockEnum.End) + const knowledgeRetrievalNodeMetadata = createNodeMetadata(BlockEnum.KnowledgeRetrieval) + mockUseAvailableNodesMetaData.mockReturnValue({ + nodes: [ + llmNodeMetadata, + humanInputNodeMetadata, + endNodeMetadata, + knowledgeRetrievalNodeMetadata, + ], + nodesMap: { + [BlockEnum.LLM]: llmNodeMetadata, + [BlockEnum.HumanInput]: humanInputNodeMetadata, + [BlockEnum.End]: endNodeMetadata, + [BlockEnum.KnowledgeRetrieval]: knowledgeRetrievalNodeMetadata, + }, + }) mockHandleCheckBeforePublish.mockResolvedValue(true) capturedHooksStore = undefined snippetDetailStoreState = { @@ -320,6 +352,23 @@ describe('SnippetMain', () => { }) }) + describe('Block Selector', () => { + it('should filter unsupported snippet block types from available node metadata', () => { + renderSnippetMain() + + const availableNodesMetaData = capturedHooksStore?.availableNodesMetaData as { + nodes: Array<{ metaData: { type: BlockEnum } }> + nodesMap: Partial> + } + + expect(availableNodesMetaData.nodes.map(node => node.metaData.type)).toEqual([BlockEnum.LLM]) + expect(availableNodesMetaData.nodesMap[BlockEnum.LLM]).toBeDefined() + expect(availableNodesMetaData.nodesMap[BlockEnum.HumanInput]).toBeUndefined() + expect(availableNodesMetaData.nodesMap[BlockEnum.End]).toBeUndefined() + expect(availableNodesMetaData.nodesMap[BlockEnum.KnowledgeRetrieval]).toBeUndefined() + }) + }) + describe('Run Hooks', () => { it('should pass snippet run handlers to WorkflowWithInnerContext', () => { renderSnippetMain() diff --git a/web/app/components/snippets/components/snippet-main.tsx b/web/app/components/snippets/components/snippet-main.tsx index af176e5921..9dd52a8589 100644 --- a/web/app/components/snippets/components/snippet-main.tsx +++ b/web/app/components/snippets/components/snippet-main.tsx @@ -38,6 +38,12 @@ type SnippetMainContentProps = { onCancel: () => void | Promise } +const unsupportedSnippetBlockTypes = new Set([ + BlockEnum.HumanInput, + BlockEnum.End, + BlockEnum.KnowledgeRetrieval, +]) + const SnippetMainContent = ({ snippetId, fields, @@ -110,7 +116,7 @@ const SnippetMain = ({ } = publishedWorkflowQuery const availableNodesMetaData = useMemo(() => { const nodes = workflowAvailableNodesMetaData.nodes.filter(node => - node.metaData.type !== BlockEnum.HumanInput && node.metaData.type !== BlockEnum.End) + !unsupportedSnippetBlockTypes.has(node.metaData.type)) if (!workflowAvailableNodesMetaData.nodesMap) return { nodes } @@ -118,6 +124,7 @@ const SnippetMain = ({ const { [BlockEnum.HumanInput]: _humanInput, [BlockEnum.End]: _end, + [BlockEnum.KnowledgeRetrieval]: _knowledgeRetrieval, ...nodesMap } = workflowAvailableNodesMetaData.nodesMap diff --git a/web/app/components/workflow/hooks/use-nodes-available-var-list.ts b/web/app/components/workflow/hooks/use-nodes-available-var-list.ts index d3b1d57d0d..7c198aee29 100644 --- a/web/app/components/workflow/hooks/use-nodes-available-var-list.ts +++ b/web/app/components/workflow/hooks/use-nodes-available-var-list.ts @@ -9,7 +9,7 @@ import { } from '@/app/components/workflow/hooks' import { appendSnippetInputFieldVars, - isSnippetCanvas, + filterSnippetSystemVars, } from '@/app/components/workflow/nodes/_base/hooks/snippet-input-field-vars' import { BlockEnum } from '@/app/components/workflow/types' @@ -36,18 +36,6 @@ const getNodeInfo = (nodeId: string, nodes: Node[]) => { } } -const filterSystemVarsForSnippet = (availableVars: NodeOutPutVar[]) => { - if (!isSnippetCanvas()) - return availableVars - - return availableVars - .map(nodeVar => ({ - ...nodeVar, - vars: nodeVar.vars.filter(variable => !variable.variable.startsWith('sys.')), - })) - .filter(nodeVar => nodeVar.vars.length > 0) -} - // TODO: loop type? const useNodesAvailableVarList = (nodes: Node[], { onlyLeafNodeVar, @@ -82,7 +70,7 @@ const useNodesAvailableVarList = (nodes: Node[], { parentNode: iterationNode, } = getNodeInfo(nodeId, nodes) - const availableVars = filterSystemVarsForSnippet([ + const availableVars = filterSnippetSystemVars([ ...snippetInputFieldAvailability.availableVars, ...getNodeAvailableVars({ parentNode: iterationNode, @@ -136,7 +124,7 @@ export const useGetNodesAvailableVarList = () => { parentNode: iterationNode, } = getNodeInfo(nodeId, nodes) - const availableVars = filterSystemVarsForSnippet([ + const availableVars = filterSnippetSystemVars([ ...snippetInputFieldAvailability.availableVars, ...getNodeAvailableVars({ parentNode: iterationNode, diff --git a/web/app/components/workflow/nodes/_base/hooks/__tests__/use-available-var-list.spec.ts b/web/app/components/workflow/nodes/_base/hooks/__tests__/use-available-var-list.spec.ts new file mode 100644 index 0000000000..8784ff381d --- /dev/null +++ b/web/app/components/workflow/nodes/_base/hooks/__tests__/use-available-var-list.spec.ts @@ -0,0 +1,108 @@ +import type { Node, NodeOutPutVar, Var } from '@/app/components/workflow/types' +import { renderHook } from '@testing-library/react' +import { BlockEnum, VarType } from '@/app/components/workflow/types' +import useAvailableVarList from '../use-available-var-list' + +const mockGetTreeLeafNodes = vi.hoisted(() => vi.fn()) +const mockGetBeforeNodesInSameBranchIncludeParent = vi.hoisted(() => vi.fn()) +const mockGetNodeById = vi.hoisted(() => vi.fn()) +const mockGetNodeAvailableVars = vi.hoisted(() => vi.fn()) + +vi.mock('@/app/components/snippets/store', () => ({ + useSnippetDetailStore: (selector: (state: { fields: unknown[] }) => unknown) => selector({ fields: [] }), +})) + +vi.mock('@/app/components/workflow/hooks', () => ({ + useIsChatMode: () => true, + useWorkflow: () => ({ + getTreeLeafNodes: mockGetTreeLeafNodes, + getBeforeNodesInSameBranchIncludeParent: mockGetBeforeNodesInSameBranchIncludeParent, + getNodeById: mockGetNodeById, + }), + useWorkflowVariables: () => ({ + getNodeAvailableVars: mockGetNodeAvailableVars, + }), +})) + +vi.mock('@/app/components/workflow/store', () => ({ + useStore: (selector: (state: { ragPipelineVariables: unknown[] }) => unknown) => selector({ ragPipelineVariables: [] }), +})) + +vi.mock('../use-node-info', () => ({ + default: () => ({ + parentNode: null, + }), +})) + +const createNode = (overrides: Partial = {}): Node => ({ + id: 'node-1', + type: 'custom', + position: { x: 0, y: 0 }, + data: { + type: BlockEnum.LLM, + title: 'Node', + desc: '', + }, + ...overrides, +} as Node) + +const outputVarsWithSystemVars: NodeOutPutVar[] = [ + { + nodeId: 'vars-node', + title: 'Vars', + vars: [ + { + variable: 'answer', + type: VarType.string, + }, + { + variable: 'sys.files', + type: VarType.arrayFile, + }, + ] satisfies Var[], + }, + { + nodeId: 'global', + title: 'SYSTEM', + vars: [{ + variable: 'sys.user_id', + type: VarType.string, + }] satisfies Var[], + }, +] + +describe('useAvailableVarList', () => { + beforeEach(() => { + vi.clearAllMocks() + globalThis.history.pushState({}, '', '/') + mockGetBeforeNodesInSameBranchIncludeParent.mockReturnValue([createNode({ id: 'before-node' })]) + mockGetTreeLeafNodes.mockReturnValue([createNode({ id: 'leaf-node' })]) + mockGetNodeById.mockReturnValue(createNode({ id: 'node-1' })) + mockGetNodeAvailableVars.mockReturnValue(outputVarsWithSystemVars) + }) + + it('filters system variables on snippet canvases', () => { + globalThis.history.pushState({}, '', '/snippets/snippet-1/orchestrate') + + const { result } = renderHook(() => useAvailableVarList('node-1', { + filterVar: () => true, + })) + + expect(result.current.availableVars).toEqual([{ + nodeId: 'vars-node', + title: 'Vars', + vars: [{ + variable: 'answer', + type: VarType.string, + }], + }]) + }) + + it('keeps system variables outside snippet canvases', () => { + const { result } = renderHook(() => useAvailableVarList('node-1', { + filterVar: () => true, + })) + + expect(result.current.availableVars).toEqual(outputVarsWithSystemVars) + }) +}) diff --git a/web/app/components/workflow/nodes/_base/hooks/snippet-input-field-vars.ts b/web/app/components/workflow/nodes/_base/hooks/snippet-input-field-vars.ts index a0c84b7495..96a6473257 100644 --- a/web/app/components/workflow/nodes/_base/hooks/snippet-input-field-vars.ts +++ b/web/app/components/workflow/nodes/_base/hooks/snippet-input-field-vars.ts @@ -14,6 +14,18 @@ export const isSnippetCanvas = () => { return /^\/snippets\/[^/]+\/orchestrate/.test(globalThis.location.pathname) } +export const filterSnippetSystemVars = (availableVars: NodeOutPutVar[]) => { + if (!isSnippetCanvas()) + return availableVars + + return availableVars + .map(nodeVar => ({ + ...nodeVar, + vars: nodeVar.vars.filter(variable => !variable.variable.startsWith('sys.')), + })) + .filter(nodeVar => nodeVar.vars.length > 0) +} + const toWorkflowInputType = (type: SnippetInputField['type']) => type as unknown as InputVarType export const buildSnippetInputFieldNode = ( diff --git a/web/app/components/workflow/nodes/_base/hooks/use-available-var-list.ts b/web/app/components/workflow/nodes/_base/hooks/use-available-var-list.ts index e94f94916f..4e97a3e86a 100644 --- a/web/app/components/workflow/nodes/_base/hooks/use-available-var-list.ts +++ b/web/app/components/workflow/nodes/_base/hooks/use-available-var-list.ts @@ -9,7 +9,7 @@ import { import { useStore as useWorkflowStore } from '@/app/components/workflow/store' import { BlockEnum } from '@/app/components/workflow/types' import { inputVarTypeToVarType } from '../../data-source/utils' -import { appendSnippetInputFieldVars } from './snippet-input-field-vars' +import { appendSnippetInputFieldVars, filterSnippetSystemVars } from './snippet-input-field-vars' import useNodeInfo from './use-node-info' type Params = { @@ -73,7 +73,7 @@ const useAvailableVarList = (nodeId: string, { }) } } - const availableVars = [ + const availableVars = filterSnippetSystemVars([ ...snippetInputFieldAvailability.availableVars, ...getNodeAvailableVars({ parentNode: iterationNode, @@ -84,7 +84,7 @@ const useAvailableVarList = (nodeId: string, { hideChatVar, }), ...dataSourceRagVars, - ] + ]) return { availableVars,