dify/web/app/components/workflow/nodes/knowledge-retrieval/hooks/use-knowledge-input-manager.ts

import type { MutableRefObject } from 'react'
import type { KnowledgeRetrievalNodeType } from '../types'
import type { ValueSelector } from '@/app/components/workflow/types'
import { produce } from 'immer'
import { useCallback, useEffect, useRef } from 'react'
import { RETRIEVE_TYPE } from '@/types/app'

type Params = {
  inputs: KnowledgeRetrievalNodeType
  doSetInputs: (inputs: KnowledgeRetrievalNodeType) => void
}

const normalizeInputs = (nextInputs: KnowledgeRetrievalNodeType) => {
  return produce(nextInputs, (draft) => {
    if (nextInputs.retrieval_mode === RETRIEVE_TYPE.multiWay) delete draft.single_retrieval_config
    else delete draft.multiple_retrieval_config
  })
}

const useKnowledgeInputManager = ({ inputs, doSetInputs }: Params) => {
  const inputRef = useRef(inputs)

  useEffect(() => {
    inputRef.current = inputs
  }, [inputs])

  const setInputs = useCallback(
    (nextInputs: KnowledgeRetrievalNodeType) => {
      const normalizedInputs = normalizeInputs(nextInputs)
      doSetInputs(normalizedInputs)
      inputRef.current = normalizedInputs
    },
    [doSetInputs],
  )

  const handleQueryVarChange = useCallback(
    (newVar: ValueSelector | string) => {
      const nextInputs = produce(inputRef.current, (draft) => {
        draft.query_variable_selector = newVar as ValueSelector
      })
      setInputs(nextInputs)
    },
    [setInputs],
  )

  const handleQueryAttachmentChange = useCallback(
    (newVar: ValueSelector | string) => {
      const nextInputs = produce(inputRef.current, (draft) => {
        draft.query_attachment_selector = newVar as ValueSelector
      })
      setInputs(nextInputs)
    },
    [setInputs],
  )

  return {
    inputRef: inputRef as MutableRefObject<KnowledgeRetrievalNodeType>,
    setInputs,
    handleQueryVarChange,
    handleQueryAttachmentChange,
  }
}

export default useKnowledgeInputManager