feat: Enhance context variable handling for Agent and LLM nodes

This commit is contained in:
zhsama
2026-01-15 23:26:19 +08:00
parent f247ebfbe1
commit f43fde5797
8 changed files with 88 additions and 65 deletions

View File

@ -220,7 +220,7 @@ const ComponentPicker = ({
<AgentNodeList
nodes={agentNodes.map(node => ({
...node,
type: BlockEnum.Agent,
type: BlockEnum.Agent || BlockEnum.LLM,
}))}
onSelect={handleSelectAgent}
onClose={handleClose}

View File

@ -67,7 +67,8 @@ const WorkflowVariableBlockComponent = ({
)()
const [localWorkflowNodesMap, setLocalWorkflowNodesMap] = useState<WorkflowNodesMap>(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 <span className="hidden" ref={ref} />
const Item = (

View File

@ -123,8 +123,9 @@ export class WorkflowVariableBlockNode extends DecoratorNode<React.JSX.Element>
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}}}`
}
}

View File

@ -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,

View File

@ -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,
}
}

View File

@ -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<Props> = ({
setControlPromptEditorRerenderKey,
} = workflowStore.getState()
const store = useStoreApi()
const { getBeforeNodesInSameBranch } = useWorkflow()
const [isContextMenuOpen, setIsContextMenuOpen] = useState(false)
const contextMenuTriggerRef = useRef<HTMLDivElement>(null)
@ -122,40 +116,21 @@ const ConfigPrompt: FC<Props> = ({
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<NodeOutPutVar[]>(() => {
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) => {

View File

@ -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<string, Node>)
}, [availableNodes])
const contextNodeIds = useMemo(() => {
const ids = new Set<string>()
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)

View File

@ -49,27 +49,32 @@ const SubGraphModal: FC<SubGraphModalProps> = ({
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<MentionConfig>(() => {
const current = toolParam?.mention_config
@ -240,7 +245,7 @@ const SubGraphModal: FC<SubGraphModalProps> = ({
onMentionConfigChange={handleMentionConfigChange}
extractorNode={extractorNode}
toolParamValue={toolParamValue}
parentAvailableNodes={parentAgentNodes}
parentAvailableNodes={parentContextNodes}
parentAvailableVars={parentAvailableVars}
onSave={handleSave}
onSyncWorkflowDraft={doSyncWorkflowDraft}