import { AgentGlobals } from '@/constants/agent'; import { useFetchAgent } from '@/hooks/use-agent-request'; import { RAGFlowNodeType } from '@/interfaces/database/flow'; import { buildNodeOutputOptions } from '@/utils/canvas-util'; import { DefaultOptionType } from 'antd/es/select'; import { t } from 'i18next'; import { isEmpty, toLower } from 'lodash'; import get from 'lodash/get'; import { useCallback, useContext, useEffect, useMemo, useState } from 'react'; import { AgentDialogueMode, BeginId, BeginQueryType, JsonSchemaDataType, Operator, VariableType, } from '../constant'; import { AgentFormContext } from '../context'; import { buildBeginInputListFromObject } from '../form/begin-form/utils'; import { BeginQuery } from '../interface'; import OperatorIcon from '../operator-icon'; import useGraphStore from '../store'; import { useFindAgentStructuredOutputLabelByValue, useFindAgentStructuredOutputTypeByValue, } from './use-build-structured-output'; export function useSelectBeginNodeDataInputs() { const getNode = useGraphStore((state) => state.getNode); return buildBeginInputListFromObject( getNode(BeginId)?.data?.form?.inputs ?? {}, ); } export function useIsTaskMode(isTask?: boolean) { const getNode = useGraphStore((state) => state.getNode); return useMemo(() => { if (typeof isTask === 'boolean') { return isTask; } const node = getNode(BeginId); return node?.data?.form?.mode === AgentDialogueMode.Task; }, [getNode, isTask]); } export const useGetBeginNodeDataQuery = () => { const getNode = useGraphStore((state) => state.getNode); const getBeginNodeDataQuery = useCallback(() => { return buildBeginInputListFromObject( get(getNode(BeginId), 'data.form.inputs', {}), ); }, [getNode]); return getBeginNodeDataQuery; }; export const useGetBeginNodeDataInputs = () => { const getNode = useGraphStore((state) => state.getNode); const inputs = get(getNode(BeginId), 'data.form.inputs', {}); const beginNodeDataInputs = useMemo(() => { return buildBeginInputListFromObject(inputs); }, [inputs]); return beginNodeDataInputs; }; export const useGetBeginNodeDataQueryIsSafe = () => { const [isBeginNodeDataQuerySafe, setIsBeginNodeDataQuerySafe] = useState(false); const inputs = useSelectBeginNodeDataInputs(); const nodes = useGraphStore((state) => state.nodes); useEffect(() => { const query: BeginQuery[] = inputs; const isSafe = !query.some((q) => !q.optional && q.type === 'file'); setIsBeginNodeDataQuerySafe(isSafe); }, [inputs, nodes]); return isBeginNodeDataQuerySafe; }; export function useBuildNodeOutputOptions(nodeId?: string) { const nodes = useGraphStore((state) => state.nodes); const edges = useGraphStore((state) => state.edges); return useMemo(() => { return buildNodeOutputOptions({ nodes, edges, nodeId, Icon: ({ name }) => , }); }, [edges, nodeId, nodes]); } // exclude nodes with branches const ExcludedNodes = [ Operator.Categorize, Operator.Relevant, Operator.Begin, Operator.Note, ]; const StringList = [ BeginQueryType.Line, BeginQueryType.Paragraph, BeginQueryType.Options, ]; function transferToVariableType(type: string) { if (StringList.some((x) => x === type)) { return VariableType.String; } return type; } export function useBuildBeginVariableOptions() { const inputs = useSelectBeginNodeDataInputs(); const options = useMemo(() => { return [ { label: {t('flow.beginInput')}, title: t('flow.beginInput'), options: inputs.map((x) => ({ label: x.name, parentLabel: {t('flow.beginInput')}, icon: , value: `begin@${x.key}`, type: transferToVariableType(x.type), })), }, ]; }, [inputs]); return options; } export const useBuildVariableOptions = (nodeId?: string, parentId?: string) => { const nodeOutputOptions = useBuildNodeOutputOptions(nodeId); const parentNodeOutputOptions = useBuildNodeOutputOptions(parentId); const beginOptions = useBuildBeginVariableOptions(); const options = useMemo(() => { return [...beginOptions, ...nodeOutputOptions, ...parentNodeOutputOptions]; }, [beginOptions, nodeOutputOptions, parentNodeOutputOptions]); return options; }; export function useBuildQueryVariableOptions(n?: RAGFlowNodeType) { const { data } = useFetchAgent(); const node = useContext(AgentFormContext) || n; const options = useBuildVariableOptions(node?.id, node?.parentId); const nextOptions = useMemo(() => { const globals = data?.dsl?.globals ?? {}; const globalOptions = Object.entries(globals).map(([key, value]) => ({ label: key, value: key, icon: , parentLabel: {t('flow.beginInput')}, type: Array.isArray(value) ? `${VariableType.Array}${key === AgentGlobals.SysFiles ? '' : ''}` : typeof value, })); return [ { ...options[0], options: [...options[0]?.options, ...globalOptions] }, ...options.slice(1), ]; }, [data.dsl?.globals, options]); return nextOptions; } export function useFilterQueryVariableOptionsByTypes( types?: JsonSchemaDataType[], ) { const nextOptions = useBuildQueryVariableOptions(); const filteredOptions = useMemo(() => { return !isEmpty(types) ? nextOptions.map((x) => { return { ...x, options: x.options.filter( (y) => types?.some((x) => toLower(y.type).includes(x)) || y.type === undefined, // agent structured output ), }; }) : nextOptions; }, [nextOptions, types]); return filteredOptions; } export function useBuildComponentIdOptions(nodeId?: string, parentId?: string) { const nodes = useGraphStore((state) => state.nodes); // Limit the nodes inside iteration to only reference peer nodes with the same parentId and other external nodes other than their parent nodes const filterChildNodesToSameParentOrExternal = useCallback( (node: RAGFlowNodeType) => { // Node inside iteration if (parentId) { return ( (node.parentId === parentId || node.parentId === undefined) && node.id !== parentId ); } return node.parentId === undefined; // The outermost node }, [parentId], ); const componentIdOptions = useMemo(() => { return nodes .filter( (x) => x.id !== nodeId && !ExcludedNodes.some((y) => y === x.data.label) && filterChildNodesToSameParentOrExternal(x), ) .map((x) => ({ label: x.data.name, value: x.id })); }, [nodes, nodeId, filterChildNodesToSameParentOrExternal]); return [ { label: Component Output, title: 'Component Output', options: componentIdOptions, }, ]; } export function useBuildComponentIdAndBeginOptions( nodeId?: string, parentId?: string, ) { const componentIdOptions = useBuildComponentIdOptions(nodeId, parentId); const beginOptions = useBuildBeginVariableOptions(); return [...beginOptions, ...componentIdOptions]; } export const useGetComponentLabelByValue = (nodeId: string) => { const options = useBuildComponentIdAndBeginOptions(nodeId); const flattenOptions = useMemo(() => { return options.reduce((pre, cur) => { return [...pre, ...cur.options]; }, []); }, [options]); const getLabel = useCallback( (val?: string) => { return flattenOptions.find((x) => x.value === val)?.label; }, [flattenOptions], ); return getLabel; }; export function useFlattenQueryVariableOptions(nodeId?: string) { const { getNode } = useGraphStore((state) => state); const nextOptions = useBuildQueryVariableOptions(getNode(nodeId)); const flattenOptions = useMemo(() => { return nextOptions.reduce((pre, cur) => { return [...pre, ...cur.options]; }, []); }, [nextOptions]); return flattenOptions; } export function useGetVariableLabelOrTypeByValue(nodeId?: string) { const flattenOptions = useFlattenQueryVariableOptions(nodeId); const findAgentStructuredOutputTypeByValue = useFindAgentStructuredOutputTypeByValue(); const findAgentStructuredOutputLabel = useFindAgentStructuredOutputLabelByValue(); const getItem = useCallback( (val?: string) => { return flattenOptions.find((x) => x.value === val); }, [flattenOptions], ); const getLabel = useCallback( (val?: string) => { const item = getItem(val); if (item) { return (
{item.parentLabel} / {item.label}
); } return getItem(val)?.label || findAgentStructuredOutputLabel(val); }, [findAgentStructuredOutputLabel, getItem], ); const getType = useCallback( (val?: string) => { return getItem(val)?.type || findAgentStructuredOutputTypeByValue(val); }, [findAgentStructuredOutputTypeByValue, getItem], ); return { getLabel, getType }; }