diff --git a/api/controllers/console/snippets/snippet_workflow_draft_variable.py b/api/controllers/console/snippets/snippet_workflow_draft_variable.py index d0f98ffa95b..a28ba07b5dd 100644 --- a/api/controllers/console/snippets/snippet_workflow_draft_variable.py +++ b/api/controllers/console/snippets/snippet_workflow_draft_variable.py @@ -34,11 +34,8 @@ from controllers.console.app.workflow_draft_variable import ( ) from controllers.console.snippets.snippet_workflow import get_snippet from controllers.console.wraps import ( - RBACPermission, - RBACResourceScope, account_initialization_required, edit_permission_required, - rbac_permission_required, setup_required, with_current_user, ) @@ -128,9 +125,6 @@ class SnippetWorkflowVariableCollectionApi(Resource): @console_ns.doc(description="Delete all draft workflow variables for the current user (snippet scope)") @console_ns.response(204, "Workflow variables deleted successfully") @_snippet_draft_var_prerequisite - @rbac_permission_required( - RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False - ) def delete(self, current_user: Account, snippet: CustomizedSnippet) -> Response: draft_var_srv = WorkflowDraftVariableService(session=db.session()) draft_var_srv.delete_user_workflow_variables(snippet.id, user_id=current_user.id) @@ -157,9 +151,6 @@ class SnippetNodeVariableCollectionApi(Resource): @console_ns.doc(description="Delete all variables for a specific node (snippet draft workflow)") @console_ns.response(204, "Node variables deleted successfully") @_snippet_draft_var_prerequisite - @rbac_permission_required( - RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False - ) def delete(self, current_user: Account, snippet: CustomizedSnippet, node_id: str) -> Response: validate_node_id(node_id) srv = WorkflowDraftVariableService(db.session()) @@ -194,9 +185,6 @@ class SnippetVariableApi(Resource): @console_ns.response(404, "Variable not found") @_snippet_draft_var_prerequisite @marshal_with(workflow_draft_variable_model) - @rbac_permission_required( - RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False - ) def patch(self, current_user: Account, snippet: CustomizedSnippet, variable_id: str) -> WorkflowDraftVariable: draft_var_srv = WorkflowDraftVariableService(session=db.session()) args_model = WorkflowDraftVariableUpdatePayload.model_validate(console_ns.payload or {}) @@ -244,9 +232,6 @@ class SnippetVariableApi(Resource): @console_ns.response(204, "Variable deleted successfully") @console_ns.response(404, "Variable not found") @_snippet_draft_var_prerequisite - @rbac_permission_required( - RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False - ) def delete(self, current_user: Account, snippet: CustomizedSnippet, variable_id: str) -> Response: draft_var_srv = WorkflowDraftVariableService(session=db.session()) variable = ensure_variable_access( @@ -269,9 +254,6 @@ class SnippetVariableResetApi(Resource): @console_ns.response(204, "Variable reset (no content)") @console_ns.response(404, "Variable not found") @_snippet_draft_var_prerequisite - @rbac_permission_required( - RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False - ) def put(self, current_user: Account, snippet: CustomizedSnippet, variable_id: str) -> Response | Any: draft_var_srv = WorkflowDraftVariableService(session=db.session()) snippet_service = _snippet_service() diff --git a/web/app/components/snippets/hooks/__tests__/use-create-snippet-from-selection.spec.tsx b/web/app/components/snippets/hooks/__tests__/use-create-snippet-from-selection.spec.tsx new file mode 100644 index 00000000000..ddf48659807 --- /dev/null +++ b/web/app/components/snippets/hooks/__tests__/use-create-snippet-from-selection.spec.tsx @@ -0,0 +1,275 @@ +import type { ReactElement } from 'react' +import type { Edge, Node } from '@/app/components/workflow/types' +import type { SnippetCanvasData, SnippetInputField } from '@/models/snippet' +import { act, renderHook } from '@testing-library/react' +import { BlockEnum } from '@/app/components/workflow/types' +import { PipelineInputVarType } from '@/models/pipeline' +import { useCreateSnippetFromSelection } from '../use-create-snippet-from-selection' + +const SNIPPET_INPUT_FIELD_NODE_ID = 'start' +const mockHandleOpenCreateSnippetDialog = vi.fn() +const mockHandleCloseCreateSnippetDialog = vi.fn() +const mockHandleCreateSnippet = vi.fn() + +vi.mock('../use-create-snippet', () => ({ + useCreateSnippet: () => ({ + createSnippetMutation: { + isPending: false, + }, + handleCloseCreateSnippetDialog: mockHandleCloseCreateSnippetDialog, + handleCreateSnippet: mockHandleCreateSnippet, + handleOpenCreateSnippetDialog: mockHandleOpenCreateSnippetDialog, + isCreateSnippetDialogOpen: true, + isCreatingSnippet: false, + }), +})) + +type DialogProps = { + selectedGraph?: SnippetCanvasData + inputFields?: SnippetInputField[] +} + +const createNode = ( + id: string, + data: Record, +): Node => ({ + id, + type: 'custom', + position: { x: 0, y: 0 }, + width: 200, + height: 100, + data, +} as unknown as Node) + +describe('useCreateSnippetFromSelection', () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it('should convert environment, conversation, and system variables into snippet input fields', () => { + const selectedNodes = [ + createNode('llm', { + type: BlockEnum.LLM, + prompt: [ + '{{#env.API_KEY#}}', + '{{#conversation.user_name#}}', + '{{#sys.user_id#}}', + '{{#rag.query#}}', + '{{#source.result#}}', + ].join(' '), + model_selector: ['env', 'MODEL_NAME'], + }), + ] + const edges: Edge[] = [] + const onClose = vi.fn() + + const { result } = renderHook(() => useCreateSnippetFromSelection({ + edges, + selectedNodes, + onClose, + })) + + act(() => { + result.current.handleOpenCreateSnippet() + }) + + const dialogProps = (result.current.createSnippetDialog as ReactElement).props + + expect(dialogProps.inputFields).toEqual([ + { + label: 'API_KEY', + variable: 'API_KEY', + type: PipelineInputVarType.textInput, + required: true, + }, + { + label: 'user_name', + variable: 'user_name', + type: PipelineInputVarType.textInput, + required: true, + }, + { + label: 'user_id', + variable: 'user_id', + type: PipelineInputVarType.textInput, + required: true, + }, + { + label: 'result', + variable: 'result', + type: PipelineInputVarType.textInput, + required: true, + }, + { + label: 'MODEL_NAME', + variable: 'MODEL_NAME', + type: PipelineInputVarType.textInput, + required: true, + }, + ]) + const nodeData = dialogProps.selectedGraph?.nodes[0]?.data as Record | undefined + + expect(nodeData?.prompt).toBe([ + `{{#${SNIPPET_INPUT_FIELD_NODE_ID}.API_KEY#}}`, + `{{#${SNIPPET_INPUT_FIELD_NODE_ID}.user_name#}}`, + `{{#${SNIPPET_INPUT_FIELD_NODE_ID}.user_id#}}`, + '{{#rag.query#}}', + `{{#${SNIPPET_INPUT_FIELD_NODE_ID}.result#}}`, + ].join(' ')) + expect(nodeData?.model_selector).toEqual([ + SNIPPET_INPUT_FIELD_NODE_ID, + 'MODEL_NAME', + ]) + expect(onClose).toHaveBeenCalled() + }) + + it('should convert system variables used by if-else and variable aggregator nodes', () => { + const selectedNodes = [ + createNode('llm', { + type: BlockEnum.LLM, + title: 'LLM', + }), + createNode('if-else', { + type: BlockEnum.IfElse, + cases: [{ + case_id: 'case-1', + conditions: [{ + id: 'condition-1', + variable_selector: ['sys', 'query'], + comparison_operator: 'contains', + value: 'hello', + }], + }], + }), + createNode('variable-aggregator', { + type: BlockEnum.VariableAggregator, + variables: [ + ['sys', 'files'], + ['llm', 'text'], + ], + advanced_settings: { + group_enabled: true, + groups: [{ + groupId: 'group-1', + group_name: 'Group1', + variables: [ + ['sys', 'workflow_id'], + ], + }], + }, + }), + ] + const onClose = vi.fn() + + const { result } = renderHook(() => useCreateSnippetFromSelection({ + edges: [], + selectedNodes, + onClose, + })) + + act(() => { + result.current.handleOpenCreateSnippet() + }) + + const dialogProps = (result.current.createSnippetDialog as ReactElement).props + + expect(dialogProps.inputFields).toEqual([ + { + label: 'query', + variable: 'query', + type: PipelineInputVarType.textInput, + required: true, + }, + { + label: 'files', + variable: 'files', + type: PipelineInputVarType.multiFiles, + required: true, + }, + { + label: 'workflow_id', + variable: 'workflow_id', + type: PipelineInputVarType.textInput, + required: true, + }, + ]) + + const ifElseNode = dialogProps.selectedGraph?.nodes.find(node => node.id === 'if-else') + const aggregatorNode = dialogProps.selectedGraph?.nodes.find(node => node.id === 'variable-aggregator') + const ifElseData = ifElseNode?.data as Record + const aggregatorData = aggregatorNode?.data as { + variables?: string[][] + advanced_settings?: { groups?: Array<{ variables?: string[][] }> } + } + + expect(ifElseData.cases).toEqual([{ + case_id: 'case-1', + conditions: [{ + id: 'condition-1', + variable_selector: [SNIPPET_INPUT_FIELD_NODE_ID, 'query'], + comparison_operator: 'contains', + value: 'hello', + }], + }]) + expect(aggregatorData.variables).toEqual([ + [SNIPPET_INPUT_FIELD_NODE_ID, 'files'], + ['llm', 'text'], + ]) + expect(aggregatorData.advanced_settings?.groups?.[0]?.variables).toEqual([ + [SNIPPET_INPUT_FIELD_NODE_ID, 'workflow_id'], + ]) + }) + + it('should keep #context# prompt placeholders when creating a snippet from workflow selection', () => { + const selectedNodes = [ + createNode('llm', { + type: BlockEnum.LLM, + context: { + enabled: true, + variable_selector: ['code', 'result'], + }, + prompt: '{{#context#}} {{#code.summary#}}', + }), + ] + const onClose = vi.fn() + + const { result } = renderHook(() => useCreateSnippetFromSelection({ + edges: [], + selectedNodes, + onClose, + })) + + act(() => { + result.current.handleOpenCreateSnippet() + }) + + const dialogProps = (result.current.createSnippetDialog as ReactElement).props + const nodeData = dialogProps.selectedGraph?.nodes[0]?.data as { + context?: { + enabled: boolean + variable_selector: string[] + } + prompt?: string + } + + expect(dialogProps.inputFields).toEqual([ + { + label: 'result', + variable: 'result', + type: PipelineInputVarType.textInput, + required: true, + }, + { + label: 'summary', + variable: 'summary', + type: PipelineInputVarType.textInput, + required: true, + }, + ]) + expect(nodeData.context).toEqual({ + enabled: true, + variable_selector: [SNIPPET_INPUT_FIELD_NODE_ID, 'result'], + }) + expect(nodeData.prompt).toBe(`{{#context#}} {{#${SNIPPET_INPUT_FIELD_NODE_ID}.summary#}}`) + }) +}) diff --git a/web/app/components/snippets/hooks/use-create-snippet-from-selection.tsx b/web/app/components/snippets/hooks/use-create-snippet-from-selection.tsx new file mode 100644 index 00000000000..31ed8a5f379 --- /dev/null +++ b/web/app/components/snippets/hooks/use-create-snippet-from-selection.tsx @@ -0,0 +1,329 @@ +import type { Edge, Node, ValueSelector } from '@/app/components/workflow/types' +import type { SnippetCanvasData, SnippetInputField } from '@/models/snippet' +import { useCallback, useState } from 'react' +import { getNodesBounds } from 'reactflow' +import CreateSnippetDialog from '@/app/components/snippets/create-snippet-dialog' +import { PipelineInputVarType } from '@/models/pipeline' +import { useCreateSnippet } from './use-create-snippet' + +const DEFAULT_SNIPPET_VIEWPORT = { x: 0, y: 0, zoom: 1 } +const SNIPPET_INPUT_FIELD_NODE_ID = 'start' +const SNIPPET_VIEWPORT_WIDTH = 1200 +const SNIPPET_VIEWPORT_HEIGHT = 800 +const SNIPPET_VIEWPORT_PADDING = 160 +const VARIABLE_REFERENCE_REGEX = /\{\{#([^#{}]+)#\}\}/g +const RESERVED_VARIABLE_PREFIXES = new Set(['rag']) + +const isRecord = (value: unknown): value is Record => { + return !!value && typeof value === 'object' && !Array.isArray(value) +} + +const isValueSelector = (value: unknown): value is ValueSelector => { + return Array.isArray(value) + && value.length > 0 + && value.every(item => typeof item === 'string') +} + +const isSelectorKey = (key?: string) => { + return key === 'selector' || !!key?.endsWith('_selector') +} + +const isValueSelectorListKey = (key?: string) => { + return key === 'variables' +} + +const isValueSelectorList = (value: unknown[]) => { + return value.length > 0 && value.every(isValueSelector) +} + +const isContextPlaceholderSelector = (selector: ValueSelector) => { + return (selector.length === 1 && selector[0] === 'context') + || selector.at(-1) === '#context#' +} + +const getCenteredViewport = (nodes: Node[]) => { + if (!nodes.length) + return DEFAULT_SNIPPET_VIEWPORT + + const bounds = getNodesBounds(nodes) + if (!bounds.width || !bounds.height) + return DEFAULT_SNIPPET_VIEWPORT + + const zoom = Math.min( + (SNIPPET_VIEWPORT_WIDTH - SNIPPET_VIEWPORT_PADDING * 2) / bounds.width, + (SNIPPET_VIEWPORT_HEIGHT - SNIPPET_VIEWPORT_PADDING * 2) / bounds.height, + 1, + ) + const centerX = bounds.x + bounds.width / 2 + const centerY = bounds.y + bounds.height / 2 + + return { + x: SNIPPET_VIEWPORT_WIDTH / 2 - centerX * zoom, + y: SNIPPET_VIEWPORT_HEIGHT / 2 - centerY * zoom, + zoom, + } +} + +const collectSelectorsFromText = (value: string, selectors: ValueSelector[]) => { + for (const match of value.matchAll(VARIABLE_REFERENCE_REGEX)) { + const variablePath = match[1] + if (!variablePath) + continue + + const selector = variablePath.split('.').filter(Boolean) + if (selector.length > 0 && !isContextPlaceholderSelector(selector)) + selectors.push(selector) + } +} + +const collectVariableSelectors = ( + value: unknown, + selectors: ValueSelector[], + key?: string, +) => { + if (typeof value === 'string') { + collectSelectorsFromText(value, selectors) + return + } + + if (Array.isArray(value)) { + if (isSelectorKey(key) && isValueSelector(value)) + selectors.push(value) + + if (isValueSelectorListKey(key) && isValueSelectorList(value)) { + value.forEach(selector => selectors.push(selector)) + return + } + + value.forEach(item => collectVariableSelectors(item, selectors)) + return + } + + if (!isRecord(value)) + return + + Object.entries(value).forEach(([currentKey, currentValue]) => { + collectVariableSelectors(currentValue, selectors, currentKey) + }) +} + +const isExternalVariableSelector = ( + selector: ValueSelector, + selectedNodeIds: Set, +) => { + const nodeId = selector[0] + if (!nodeId) + return false + + if (nodeId.startsWith('$')) + return false + + if (isContextPlaceholderSelector(selector)) + return false + + if (selectedNodeIds.has(nodeId)) + return false + + return !RESERVED_VARIABLE_PREFIXES.has(nodeId) +} + +const sanitizeInputFieldVariable = (variable: string) => { + const sanitized = variable.replace(/\W/g, '_') + if (!sanitized) + return 'input' + + if (/^\d/.test(sanitized)) + return `input_${sanitized}` + + return sanitized +} + +const getUniqueInputFieldVariable = ( + selector: ValueSelector, + usedVariables: Set, +) => { + const baseVariable = sanitizeInputFieldVariable(selector.at(-1) ?? 'input') + let variable = baseVariable + let index = 2 + + while (usedVariables.has(variable)) { + variable = `${baseVariable}_${index}` + index += 1 + } + + usedVariables.add(variable) + return variable +} + +const getInputFieldType = (selector: ValueSelector) => { + const variable = selector.at(-1) + if (variable === 'files') + return PipelineInputVarType.multiFiles + + return PipelineInputVarType.textInput +} + +const getExternalVariableInputFields = ( + nodes: Node[], + selectedNodeIds: Set, +) => { + const selectors: ValueSelector[] = [] + nodes.forEach(node => collectVariableSelectors(node.data, selectors)) + + const usedVariables = new Set() + const fieldBySelector = new Map() + + selectors.forEach((selector) => { + if (!isExternalVariableSelector(selector, selectedNodeIds)) + return + + const selectorKey = selector.join('.') + if (fieldBySelector.has(selectorKey)) + return + + const variable = getUniqueInputFieldVariable(selector, usedVariables) + fieldBySelector.set(selectorKey, { + label: variable, + variable, + type: getInputFieldType(selector), + required: true, + }) + }) + + return { + inputFields: [...fieldBySelector.values()], + selectorMap: new Map( + [...fieldBySelector.entries()].map(([selectorKey, field]) => [ + selectorKey, + [SNIPPET_INPUT_FIELD_NODE_ID, field.variable] satisfies ValueSelector, + ]), + ), + } +} + +const rewriteVariableReferences = ( + value: unknown, + selectorMap: Map, + key?: string, +): unknown => { + if (typeof value === 'string') { + return value.replace(VARIABLE_REFERENCE_REGEX, (match, variablePath: string) => { + const nextSelector = selectorMap.get(variablePath) + if (!nextSelector) + return match + + return `{{#${nextSelector.join('.')}#}}` + }) + } + + if (Array.isArray(value)) { + if (isSelectorKey(key) && isValueSelector(value)) { + const nextSelector = selectorMap.get(value.join('.')) + if (nextSelector) + return nextSelector + } + + if (isValueSelectorListKey(key) && isValueSelectorList(value)) { + return value.map((selector) => { + const nextSelector = selectorMap.get(selector.join('.')) + return nextSelector || selector + }) + } + + return value.map(item => rewriteVariableReferences(item, selectorMap)) + } + + if (!isRecord(value)) + return value + + return Object.fromEntries( + Object.entries(value).map(([currentKey, currentValue]) => [ + currentKey, + rewriteVariableReferences(currentValue, selectorMap, currentKey), + ]), + ) +} + +const getSelectedSnippetGraph = (selectedNodes: Node[], edges: Edge[]) => { + const selectedNodeIds = new Set(selectedNodes.map(node => node.id)) + const { + inputFields, + selectorMap, + } = getExternalVariableInputFields(selectedNodes, selectedNodeIds) + const nodes = selectedNodes.map(node => ({ + ...node, + data: rewriteVariableReferences(node.data, selectorMap) as Node['data'], + selected: false, + })) + + return { + graph: { + nodes, + edges: edges + .filter(edge => selectedNodeIds.has(edge.source) && selectedNodeIds.has(edge.target)) + .map(edge => ({ + ...edge, + selected: false, + })), + viewport: getCenteredViewport(nodes), + } satisfies SnippetCanvasData, + inputFields, + } +} + +type UseCreateSnippetFromSelectionParams = { + edges: Edge[] + selectedNodes: Node[] + onClose: () => void +} + +export const useCreateSnippetFromSelection = ({ + edges, + selectedNodes, + onClose, +}: UseCreateSnippetFromSelectionParams) => { + const [selectedSnippetGraph, setSelectedSnippetGraph] = useState() + const [selectedSnippetInputFields, setSelectedSnippetInputFields] = useState([]) + const { + createSnippetMutation, + handleCloseCreateSnippetDialog, + handleCreateSnippet, + handleOpenCreateSnippetDialog, + isCreateSnippetDialogOpen, + isCreatingSnippet, + } = useCreateSnippet() + + const handleOpenCreateSnippet = useCallback(() => { + const { + graph, + inputFields, + } = getSelectedSnippetGraph(selectedNodes, edges) + setSelectedSnippetGraph(graph) + setSelectedSnippetInputFields(inputFields) + handleOpenCreateSnippetDialog() + onClose() + }, [edges, handleOpenCreateSnippetDialog, onClose, selectedNodes]) + + const handleCloseCreateSnippet = useCallback(() => { + setSelectedSnippetGraph(undefined) + setSelectedSnippetInputFields([]) + handleCloseCreateSnippetDialog() + }, [handleCloseCreateSnippetDialog]) + + const createSnippetDialog = ( + + ) + + return { + createSnippetDialog, + handleOpenCreateSnippet, + isCreateSnippetDialogOpen, + } +}