dify/web/app/components/workflow/nodes/knowledge-retrieval/hooks/use-knowledge-metadata-config.ts

import type { MutableRefObject } from 'react'
import type {
  HandleAddCondition,
  HandleRemoveCondition,
  HandleToggleConditionLogicalOperator,
  HandleUpdateCondition,
  KnowledgeRetrievalNodeType,
  MetadataFilteringModeEnum,
} from '../types'
import type { Var } from '@/app/components/workflow/types'
import { produce } from 'immer'
import { useCallback } from 'react'
import { v4 as uuid4 } from 'uuid'
import useAvailableVarList from '@/app/components/workflow/nodes/_base/hooks/use-available-var-list'
import { VarType } from '@/app/components/workflow/types'
import { AppModeEnum } from '@/types/app'
import {
  ComparisonOperator,
  LogicalOperator,
  MetadataFilteringVariableType,
} from '../types'

type Params = {
  id: string
  inputRef: MutableRefObject<KnowledgeRetrievalNodeType>
  setInputs: (inputs: KnowledgeRetrievalNodeType) => void
}

const filterStringVar = (varPayload: Var) => {
  return varPayload.type === VarType.string
}

const filterNumberVar = (varPayload: Var) => {
  return varPayload.type === VarType.number
}

const filterFileVar = (varPayload: Var) => {
  return varPayload.type === VarType.file || varPayload.type === VarType.arrayFile
}

const useKnowledgeMetadataConfig = ({
  id,
  inputRef,
  setInputs,
}: Params) => {
  const handleMetadataFilterModeChange = useCallback((newMode: MetadataFilteringModeEnum) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.metadata_filtering_mode = newMode
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleAddCondition = useCallback<HandleAddCondition>(({ id, name, type }) => {
    const comparisonOperator = type === MetadataFilteringVariableType.number
      ? ComparisonOperator.equal
      : ComparisonOperator.is

    const nextInputs = produce(inputRef.current, (draft) => {
      const newCondition = {
        id: uuid4(),
        metadata_id: id,
        name,
        comparison_operator: comparisonOperator,
      }

      if (draft.metadata_filtering_conditions) {
        draft.metadata_filtering_conditions.conditions.push(newCondition)
        return
      }

      draft.metadata_filtering_conditions = {
        logical_operator: LogicalOperator.and,
        conditions: [newCondition],
      }
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleRemoveCondition = useCallback<HandleRemoveCondition>((conditionId) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      const conditions = draft.metadata_filtering_conditions?.conditions || []
      const index = conditions.findIndex(condition => condition.id === conditionId)
      if (index > -1)
        draft.metadata_filtering_conditions?.conditions.splice(index, 1)
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleUpdateCondition = useCallback<HandleUpdateCondition>((conditionId, newCondition) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      const conditions = draft.metadata_filtering_conditions?.conditions || []
      const index = conditions.findIndex(condition => condition.id === conditionId)
      if (index > -1)
        draft.metadata_filtering_conditions!.conditions[index] = newCondition
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleToggleConditionLogicalOperator = useCallback<HandleToggleConditionLogicalOperator>(() => {
    const nextInputs = produce(inputRef.current, (draft) => {
      const currentLogicalOperator = draft.metadata_filtering_conditions?.logical_operator
      draft.metadata_filtering_conditions!.logical_operator = currentLogicalOperator === LogicalOperator.and
        ? LogicalOperator.or
        : LogicalOperator.and
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleMetadataModelChange = useCallback((model: { provider: string, modelId: string, mode?: string }) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.metadata_model_config = {
        provider: model.provider,
        name: model.modelId,
        mode: model.mode || AppModeEnum.CHAT,
        completion_params: draft.metadata_model_config?.completion_params || { temperature: 0.7 },
      }
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleMetadataCompletionParamsChange = useCallback((newParams: Record<string, unknown>) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.metadata_model_config = {
        ...draft.metadata_model_config!,
        completion_params: newParams,
      }
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const {
    availableVars: availableStringVars,
    availableNodesWithParent: availableStringNodesWithParent,
  } = useAvailableVarList(id, {
    onlyLeafNodeVar: false,
    filterVar: filterStringVar,
  })

  const {
    availableVars: availableNumberVars,
    availableNodesWithParent: availableNumberNodesWithParent,
  } = useAvailableVarList(id, {
    onlyLeafNodeVar: false,
    filterVar: filterNumberVar,
  })

  return {
    filterStringVar,
    filterFileVar,
    handleMetadataFilterModeChange,
    handleAddCondition,
    handleRemoveCondition,
    handleUpdateCondition,
    handleToggleConditionLogicalOperator,
    handleMetadataModelChange,
    handleMetadataCompletionParamsChange,
    availableStringVars,
    availableStringNodesWithParent,
    availableNumberVars,
    availableNumberNodesWithParent,
  }
}

export default useKnowledgeMetadataConfig