dify/web/service/use-workflow.ts

import type { FlowType } from '@/types/common'
import type {
  FetchWorkflowDraftPageParams,
  FetchWorkflowDraftResponse,
  NodeTracing,
  PublishWorkflowParams,
  UpdateWorkflowParams,
  WorkflowConfigResponse,
  WorkflowRunHistoryResponse,
} from '@/types/workflow'
import { useInfiniteQuery, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
import { useInvalid, useReset } from './use-base'
import {
  deleteAllInspectorVars,
  deleteInspectorVar,
  deleteNodeInspectorVars,
  deleteWorkflow,
  editInspectorVar,
  fetchAppWorkflow,
  fetchConversationVarValues,
  fetchLastRun,
  fetchSysVarValues,
  fetchWorkflowConfig,
  fetchWorkflowRunHistory,
  fetchWorkflowVersionHistory,
  publishWorkflow,
  resetConversationVar,
  updateWorkflow,
} from './workflow'

const NAME_SPACE = 'workflow'

export const useAppWorkflow = (appID: string) => {
  return useQuery<FetchWorkflowDraftResponse>({
    enabled: !!appID,
    queryKey: [NAME_SPACE, 'publish', appID],
    queryFn: () => fetchAppWorkflow(appID),
  })
}

export const useWorkflowRunHistory = (url?: string, enabled = true) => {
  return useQuery<WorkflowRunHistoryResponse>({
    queryKey: [NAME_SPACE, 'runHistory', url],
    queryFn: () => fetchWorkflowRunHistory(url as string),
    enabled: !!url && enabled,
  })
}

export const useInvalidateAppWorkflow = () => {
  const queryClient = useQueryClient()
  return (appID: string) => {
    queryClient.invalidateQueries(
      {
        queryKey: [NAME_SPACE, 'publish', appID],
      },
    )
  }
}

export const useWorkflowConfig = <T = WorkflowConfigResponse>(url: string, onSuccess: (v: T) => void) => {
  return useQuery({
    enabled: !!url,
    queryKey: [NAME_SPACE, 'config', url],
    staleTime: 0,
    queryFn: async () => {
      const data = await fetchWorkflowConfig<T>(url)
      onSuccess(data)
      return data
    },
  })
}

const WorkflowVersionHistoryKey = [NAME_SPACE, 'versionHistory']

export const useWorkflowVersionHistory = (params: FetchWorkflowDraftPageParams) => {
  const { url, initialPage, limit, userId, namedOnly } = params
  return useInfiniteQuery({
    enabled: !!url,
    queryKey: [...WorkflowVersionHistoryKey, url, initialPage, limit, userId, namedOnly],
    queryFn: ({ pageParam = 1 }) => fetchWorkflowVersionHistory(url, {
      page: pageParam,
      limit,
      user_id: userId || '',
      named_only: !!namedOnly,
    }),
    getNextPageParam: lastPage => lastPage.has_more ? lastPage.page + 1 : null,
    initialPageParam: initialPage,
  })
}

export const useResetWorkflowVersionHistory = () => {
  return useReset([...WorkflowVersionHistoryKey])
}

export const useUpdateWorkflow = () => {
  return useMutation({
    mutationKey: [NAME_SPACE, 'update'],
    mutationFn: (params: UpdateWorkflowParams) => updateWorkflow(params),
  })
}

export const useDeleteWorkflow = () => {
  return useMutation({
    mutationKey: [NAME_SPACE, 'delete'],
    mutationFn: (url: string) => deleteWorkflow(url),
  })
}

export const usePublishWorkflow = () => {
  return useMutation({
    mutationKey: [NAME_SPACE, 'publish'],
    mutationFn: (params: PublishWorkflowParams) => publishWorkflow(params),
  })
}

const useLastRunKey = [NAME_SPACE, 'last-run']
export const useLastRun = (flowType: FlowType, flowId: string, nodeId: string, enabled: boolean) => {
  return useQuery<NodeTracing>({
    enabled,
    queryKey: [...useLastRunKey, flowType, flowId, nodeId],
    queryFn: async () => {
      return fetchLastRun(flowType, flowId, nodeId)
    },
    retry: 0,
  })
}

export const useInvalidLastRun = (flowType: FlowType, flowId: string, nodeId: string) => {
  return useInvalid([...useLastRunKey, flowType, flowId, nodeId])
}

// Rerun workflow or change the version of workflow
export const useInvalidAllLastRun = (flowType?: FlowType, flowId?: string) => {
  return useInvalid([...useLastRunKey, flowType, flowId])
}

export const useConversationVarValues = (flowType?: FlowType, flowId?: string) => {
  return useQuery({
    enabled: !!flowId,
    queryKey: [NAME_SPACE, flowType, 'conversation var values', flowId],
    queryFn: async () => {
      return fetchConversationVarValues(flowType as FlowType, flowId as string)
    },
  })
}

export const useInvalidateConversationVarValues = (flowType: FlowType, flowId: string) => {
  return useInvalid([NAME_SPACE, flowType, 'conversation var values', flowId])
}

export const useResetConversationVar = (flowType: FlowType, flowId: string) => {
  return useMutation({
    mutationKey: [NAME_SPACE, flowType, 'reset conversation var', flowId],
    mutationFn: async (varId: string) => {
      return resetConversationVar(flowType, flowId, varId)
    },
  })
}

export const useResetToLastRunValue = (flowType: FlowType, flowId: string) => {
  return useMutation({
    mutationKey: [NAME_SPACE, flowType, 'reset to last run value', flowId],
    mutationFn: async (varId: string): Promise<{ value: any }> => {
      return resetConversationVar(flowType, flowId, varId) as Promise<{ value: any }>
    },
  })
}

export const useSysVarValuesKey = [NAME_SPACE, 'sys-variable']
export const useSysVarValues = (flowType?: FlowType, flowId?: string) => {
  return useQuery({
    enabled: !!flowId,
    queryKey: [NAME_SPACE, flowType, 'sys var values', flowId],
    queryFn: async () => {
      return fetchSysVarValues(flowType as FlowType, flowId as string)
    },
  })
}

export const useInvalidateSysVarValues = (flowType: FlowType, flowId: string) => {
  return useInvalid([NAME_SPACE, flowType, 'sys var values', flowId])
}

export const useDeleteAllInspectorVars = (flowType: FlowType, flowId: string) => {
  return useMutation({
    mutationKey: [NAME_SPACE, flowType, 'delete all inspector vars', flowId],
    mutationFn: async () => {
      return deleteAllInspectorVars(flowType, flowId)
    },
  })
}

export const useDeleteNodeInspectorVars = (flowType: FlowType, flowId: string) => {
  return useMutation({
    mutationKey: [NAME_SPACE, flowType, 'delete node inspector vars', flowId],
    mutationFn: async (nodeId: string) => {
      return deleteNodeInspectorVars(flowType, flowId, nodeId)
    },
  })
}

export const useDeleteInspectVar = (flowType: FlowType, flowId: string) => {
  return useMutation({
    mutationKey: [NAME_SPACE, flowType, 'delete inspector var', flowId],
    mutationFn: async (varId: string) => {
      return deleteInspectorVar(flowType, flowId, varId)
    },
  })
}

// edit the name or value of the inspector var
export const useEditInspectorVar = (flowType: FlowType, flowId: string) => {
  return useMutation({
    mutationKey: [NAME_SPACE, flowType, 'edit inspector var', flowId],
    mutationFn: async ({ varId, ...rest }: {
      varId: string
      name?: string
      value?: any
    }) => {
      return editInspectorVar(flowType, flowId, { varId, ...rest })
    },
  })
}