From f43fde57972fae698029aa5b9f629e8fef694a2f Mon Sep 17 00:00:00 2001 From: zhsama Date: Thu, 15 Jan 2026 23:26:19 +0800 Subject: [PATCH] feat: Enhance context variable handling for Agent and LLM nodes --- .../plugins/component-picker-block/index.tsx | 2 +- .../workflow-variable-block/component.tsx | 5 +- .../plugins/workflow-variable-block/node.tsx | 5 +- web/app/components/workflow/constants.ts | 10 ++++ .../_base/hooks/use-available-var-list.ts | 31 +++++++++-- .../nodes/llm/components/config-prompt.tsx | 53 +++++-------------- .../mixed-variable-text-input/index.tsx | 20 +++++-- .../tool/components/sub-graph-modal/index.tsx | 27 ++++++---- 8 files changed, 88 insertions(+), 65 deletions(-) diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx index 4731f9abe7..1511754128 100644 --- a/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx +++ b/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx @@ -220,7 +220,7 @@ const ComponentPicker = ({ ({ ...node, - type: BlockEnum.Agent, + type: BlockEnum.Agent || BlockEnum.LLM, }))} onSelect={handleSelectAgent} onClose={handleClose} diff --git a/web/app/components/base/prompt-editor/plugins/workflow-variable-block/component.tsx b/web/app/components/base/prompt-editor/plugins/workflow-variable-block/component.tsx index 1ff0da78ab..a0388042fc 100644 --- a/web/app/components/base/prompt-editor/plugins/workflow-variable-block/component.tsx +++ b/web/app/components/base/prompt-editor/plugins/workflow-variable-block/component.tsx @@ -67,7 +67,8 @@ const WorkflowVariableBlockComponent = ({ )() const [localWorkflowNodesMap, setLocalWorkflowNodesMap] = useState(workflowNodesMap) const node = localWorkflowNodesMap![variables[isRagVar ? 1 : 0]] - const isAgentContextVariable = node?.type === BlockEnum.Agent && variables[variablesLength - 1] === 'context' + const isContextVariable = (node?.type === BlockEnum.Agent || node?.type === BlockEnum.LLM) + && variables[variablesLength - 1] === 'context' const isException = isExceptionVariable(varName, node?.type) const variableValid = useMemo(() => { @@ -136,7 +137,7 @@ const WorkflowVariableBlockComponent = ({ }) }, [node, reactflow, store]) - if (isAgentContextVariable) + if (isContextVariable) return const Item = ( diff --git a/web/app/components/base/prompt-editor/plugins/workflow-variable-block/node.tsx b/web/app/components/base/prompt-editor/plugins/workflow-variable-block/node.tsx index 75ceb82f2d..64d6baf001 100644 --- a/web/app/components/base/prompt-editor/plugins/workflow-variable-block/node.tsx +++ b/web/app/components/base/prompt-editor/plugins/workflow-variable-block/node.tsx @@ -123,8 +123,9 @@ export class WorkflowVariableBlockNode extends DecoratorNode getTextContent(): string { const variables = this.getVariables() const node = this.getWorkflowNodesMap()?.[variables[0]] - const isAgentContextVariable = node?.type === BlockEnum.Agent && variables[variables.length - 1] === 'context' - const marker = isAgentContextVariable ? '@' : '#' + const isContextVariable = (node?.type === BlockEnum.Agent || node?.type === BlockEnum.LLM) + && variables[variables.length - 1] === 'context' + const marker = isContextVariable ? '@' : '#' return `{{${marker}${variables.join('.')}${marker}}}` } } diff --git a/web/app/components/workflow/constants.ts b/web/app/components/workflow/constants.ts index 4d95db7fcf..39bc15c6ce 100644 --- a/web/app/components/workflow/constants.ts +++ b/web/app/components/workflow/constants.ts @@ -131,6 +131,11 @@ export const SUPPORT_OUTPUT_VARS_NODE = [ ] export const AGENT_OUTPUT_STRUCT: Var[] = [ + { + variable: 'context', + type: VarType.arrayObject, + schemaType: 'List[promptMessage]', + }, { variable: 'usage', type: VarType.object, @@ -142,6 +147,11 @@ export const LLM_OUTPUT_STRUCT: Var[] = [ variable: 'text', type: VarType.string, }, + { + variable: 'context', + type: VarType.arrayObject, + schemaType: 'List[promptMessage]', + }, { variable: 'reasoning_content', type: VarType.string, 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 f226900899..e687813b69 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 @@ -71,14 +71,35 @@ const useAvailableVarList = (nodeId: string, { hideEnv, hideChatVar, }), ...dataSourceRagVars] + const availableNodesWithParent = [ + ...availableNodes, + ...(isDataSourceNode ? [currNode] : []), + ] + const llmNodeIds = new Set( + availableNodesWithParent + .filter(node => node?.data.type === BlockEnum.LLM) + .map(node => node!.id), + ) + const filteredAvailableVars = llmNodeIds.size + ? availableVars + .map((nodeVar) => { + if (!llmNodeIds.has(nodeVar.nodeId)) + return nodeVar + const nextVars = nodeVar.vars.filter(item => item.variable !== 'context') + if (nextVars.length === nodeVar.vars.length) + return nodeVar + return { + ...nodeVar, + vars: nextVars, + } + }) + .filter(nodeVar => nodeVar.vars.length > 0) + : availableVars return { - availableVars, + availableVars: filteredAvailableVars, availableNodes, - availableNodesWithParent: [ - ...availableNodes, - ...(isDataSourceNode ? [currNode] : []), - ], + availableNodesWithParent, } } diff --git a/web/app/components/workflow/nodes/llm/components/config-prompt.tsx b/web/app/components/workflow/nodes/llm/components/config-prompt.tsx index 147e69b6e1..d88ec95f34 100644 --- a/web/app/components/workflow/nodes/llm/components/config-prompt.tsx +++ b/web/app/components/workflow/nodes/llm/components/config-prompt.tsx @@ -6,7 +6,6 @@ import * as React from 'react' import { useCallback, useMemo, useRef, useState } from 'react' import { useTranslation } from 'react-i18next' import { ReactSortable } from 'react-sortablejs' -import { useStoreApi } from 'reactflow' import { v4 as uuid4 } from 'uuid' import { DragHandle } from '@/app/components/base/icons/src/vender/line/others' import { @@ -18,7 +17,6 @@ import AddButton from '@/app/components/workflow/nodes/_base/components/add-butt import Editor from '@/app/components/workflow/nodes/_base/components/prompt/editor' import VarReferenceVars from '@/app/components/workflow/nodes/_base/components/variable/var-reference-vars' import { cn } from '@/utils/classnames' -import { useWorkflow } from '../../../hooks' import { useStore, useWorkflowStore } from '../../../store' import { BlockEnum, EditionType, isPromptMessageContext, PromptRole, VarType } from '../../../types' import useAvailableVarList from '../../_base/hooks/use-available-var-list' @@ -26,7 +24,6 @@ import ConfigContextItem from './config-context-item' import ConfigPromptItem from './config-prompt-item' const i18nPrefix = 'nodes.llm' - type Props = { readOnly: boolean nodeId: string @@ -66,9 +63,6 @@ const ConfigPrompt: FC = ({ setControlPromptEditorRerenderKey, } = workflowStore.getState() - const store = useStoreApi() - const { getBeforeNodesInSameBranch } = useWorkflow() - const [isContextMenuOpen, setIsContextMenuOpen] = useState(false) const contextMenuTriggerRef = useRef(null) @@ -122,40 +116,21 @@ const ConfigPrompt: FC = ({ return Array.from(merged.values()) }, [availableNodesWithParent, parentAvailableNodes]) - const contextAgentNodes = useMemo(() => { - const agentNodes = mergedAvailableNodesWithParent - .filter(node => node.data.type === BlockEnum.Agent) - - const { getNodes } = store.getState() - const allNodes = getNodes() - const currentNode = allNodes.find(n => n.id === nodeId) - const parentNodeId = currentNode?.parentId - - if (parentNodeId) { - const beforeNodes = getBeforeNodesInSameBranch(parentNodeId) - const parentAgentNodes = beforeNodes - .filter(node => node.data.type === BlockEnum.Agent) - .filter(node => !agentNodes.some(n => n.id === node.id)) - - agentNodes.unshift(...parentAgentNodes) - } - - return agentNodes - }, [mergedAvailableNodesWithParent, nodeId, store, getBeforeNodesInSameBranch]) - const contextVarOptions = useMemo(() => { - return contextAgentNodes.map(node => ({ - nodeId: node.id, - title: node.data.title, - vars: [ - { - variable: 'context', - type: VarType.arrayObject, - schemaType: 'List[promptMessage]', - }, - ], - })) - }, [contextAgentNodes]) + return mergedAvailableNodesWithParent + .filter(node => node.data.type === BlockEnum.Agent || node.data.type === BlockEnum.LLM) + .map(node => ({ + nodeId: node.id, + title: node.data.title, + vars: [ + { + variable: 'context', + type: VarType.arrayObject, + schemaType: 'List[promptMessage]', + }, + ], + })) + }, [mergedAvailableNodesWithParent]) const handleChatModePromptChange = useCallback((index: number) => { return (prompt: string) => { diff --git a/web/app/components/workflow/nodes/tool/components/mixed-variable-text-input/index.tsx b/web/app/components/workflow/nodes/tool/components/mixed-variable-text-input/index.tsx index eec880bd77..f8418cbcbf 100644 --- a/web/app/components/workflow/nodes/tool/components/mixed-variable-text-input/index.tsx +++ b/web/app/components/workflow/nodes/tool/components/mixed-variable-text-input/index.tsx @@ -33,7 +33,6 @@ import Placeholder from './placeholder' * Example: {{@agent-123.context@}} -> captures "agent-123" */ const AGENT_CONTEXT_VAR_PATTERN = /\{\{[@#]([^.@#]+)\.context[@#]\}\}/g - const DEFAULT_MENTION_CONFIG: MentionConfig = { extractor_node_id: '', output_selector: [], @@ -151,6 +150,15 @@ const MixedVariableTextInput = ({ }, {} as Record) }, [availableNodes]) + const contextNodeIds = useMemo(() => { + const ids = new Set() + availableNodes.forEach((node) => { + if (node.data.type === BlockEnum.Agent || node.data.type === BlockEnum.LLM) + ids.add(node.id) + }) + return ids + }, [availableNodes]) + type DetectedAgent = { nodeId: string name: string @@ -165,7 +173,7 @@ const MixedVariableTextInput = ({ const variablePath = match[1] const nodeId = variablePath.split('.')[0] const node = nodesByIdMap[nodeId] - if (node?.data.type === BlockEnum.Agent) { + if (node && contextNodeIds.has(nodeId)) { return { nodeId, name: node.data.title, @@ -173,20 +181,22 @@ const MixedVariableTextInput = ({ } } return null - }, [nodesByIdMap]) + }, [contextNodeIds, nodesByIdMap]) const detectedAgentFromValue: DetectedAgent | null = useMemo(() => { return detectAgentFromText(value) }, [detectAgentFromText, value]) const agentNodes = useMemo(() => { + if (!contextNodeIds.size) + return [] return availableNodes - .filter(node => node.data.type === BlockEnum.Agent) + .filter(node => contextNodeIds.has(node.id)) .map(node => ({ id: node.id, title: node.data.title, })) - }, [availableNodes]) + }, [availableNodes, contextNodeIds]) const syncExtractorPromptFromText = useCallback((text: string) => { if (!toolNodeId || !paramKey) diff --git a/web/app/components/workflow/nodes/tool/components/sub-graph-modal/index.tsx b/web/app/components/workflow/nodes/tool/components/sub-graph-modal/index.tsx index c1d2bcbea4..2d2ce19f76 100644 --- a/web/app/components/workflow/nodes/tool/components/sub-graph-modal/index.tsx +++ b/web/app/components/workflow/nodes/tool/components/sub-graph-modal/index.tsx @@ -49,27 +49,32 @@ const SubGraphModal: FC = ({ const toolParam = (toolNode?.data as ToolNodeType | undefined)?.tool_parameters?.[paramKey] const toolParamValue = toolParam?.value as string | undefined - const parentAgentNodes = useMemo(() => { + const parentBeforeNodes = useMemo(() => { if (!isOpen) return [] - const beforeNodes = getBeforeNodesInSameBranch(toolNodeId, workflowNodes, workflowEdges) - return beforeNodes.filter(node => node.data.type === BlockEnum.Agent) + return getBeforeNodesInSameBranch(toolNodeId, workflowNodes, workflowEdges) }, [getBeforeNodesInSameBranch, isOpen, toolNodeId, workflowEdges, workflowNodes]) - const parentAgentNodeIds = useMemo(() => { - return parentAgentNodes.map(node => node.id) - }, [parentAgentNodes]) + const parentContextNodes = useMemo(() => { + if (!parentBeforeNodes.length) + return [] + return parentBeforeNodes.filter(node => node.data.type === BlockEnum.Agent || node.data.type === BlockEnum.LLM) + }, [parentBeforeNodes]) + + const parentContextNodeIds = useMemo(() => { + return parentContextNodes.map(node => node.id) + }, [parentContextNodes]) const parentAvailableVars = useMemo(() => { - if (!parentAgentNodeIds.length) + if (!parentContextNodeIds.length) return [] const vars = getNodeAvailableVars({ - beforeNodes: parentAgentNodes, + beforeNodes: parentContextNodes, isChatMode, filterVar: () => true, }) - return vars.filter(nodeVar => parentAgentNodeIds.includes(nodeVar.nodeId)) - }, [getNodeAvailableVars, isChatMode, parentAgentNodeIds, parentAgentNodes]) + return vars.filter(nodeVar => parentContextNodeIds.includes(nodeVar.nodeId)) + }, [getNodeAvailableVars, isChatMode, parentContextNodeIds, parentContextNodes]) const mentionConfig = useMemo(() => { const current = toolParam?.mention_config @@ -240,7 +245,7 @@ const SubGraphModal: FC = ({ onMentionConfigChange={handleMentionConfigChange} extractorNode={extractorNode} toolParamValue={toolParamValue} - parentAvailableNodes={parentAgentNodes} + parentAvailableNodes={parentContextNodes} parentAvailableVars={parentAvailableVars} onSave={handleSave} onSyncWorkflowDraft={doSyncWorkflowDraft}