dify/web/app/components/workflow/nodes/llm/hooks/use-llm-prompt-config.ts

import type { Draft } from 'immer'
import type { MutableRefObject } from 'react'
import type { LLMNodeType } from '../types'
import type {
  Memory,
  PromptItem,
  ValueSelector,
  Var,
  Variable,
} from '@/app/components/workflow/types'
import { produce } from 'immer'
import {
  useCallback,
  useMemo,
} from 'react'
import {
  checkHasContextBlock,
  checkHasHistoryBlock,
  checkHasQueryBlock,
} from '@/app/components/base/prompt-editor/constants'
import {
  EditionType,
  VarType,
} from '@/app/components/workflow/types'

type Params = {
  inputs: LLMNodeType
  inputRef: MutableRefObject<LLMNodeType>
  isChatMode: boolean
  isChatModel: boolean
  setInputs: (inputs: LLMNodeType) => void
}

const createPromptConfig = () => ({
  jinja2_variables: [] as Variable[],
})

const ensurePromptConfig = (draft: Draft<LLMNodeType>): { jinja2_variables: Variable[] } => {
  if (!draft.prompt_config)
    draft.prompt_config = createPromptConfig()

  if (!draft.prompt_config.jinja2_variables)
    draft.prompt_config.jinja2_variables = []

  return draft.prompt_config as { jinja2_variables: Variable[] }
}

const filterInputVar = (varPayload: Var) => {
  return [
    VarType.number,
    VarType.string,
    VarType.secret,
    VarType.arrayString,
    VarType.arrayNumber,
    VarType.file,
    VarType.arrayFile,
  ].includes(varPayload.type)
}

const filterJinja2InputVar = (varPayload: Var) => {
  return [
    VarType.number,
    VarType.string,
    VarType.secret,
    VarType.arrayString,
    VarType.arrayNumber,
    VarType.arrayBoolean,
    VarType.arrayObject,
    VarType.object,
    VarType.array,
    VarType.boolean,
  ].includes(varPayload.type)
}

const filterMemoryPromptVar = (varPayload: Var) => {
  return [
    VarType.arrayObject,
    VarType.array,
    VarType.number,
    VarType.string,
    VarType.secret,
    VarType.arrayString,
    VarType.arrayNumber,
    VarType.file,
    VarType.arrayFile,
  ].includes(varPayload.type)
}

const useLLMPromptConfig = ({
  inputs,
  inputRef,
  isChatMode,
  isChatModel,
  setInputs,
}: Params) => {
  const hasSetBlockStatus = useMemo(() => {
    const promptTemplate = inputs.prompt_template
    const hasSetContext = isChatModel
      ? (promptTemplate as PromptItem[]).some(item => checkHasContextBlock(item.text))
      : checkHasContextBlock((promptTemplate as PromptItem).text)

    if (!isChatMode) {
      return {
        history: false,
        query: false,
        context: hasSetContext,
      }
    }

    if (isChatModel) {
      return {
        history: false,
        query: (promptTemplate as PromptItem[]).some(item => checkHasQueryBlock(item.text)),
        context: hasSetContext,
      }
    }

    return {
      history: checkHasHistoryBlock((promptTemplate as PromptItem).text),
      query: checkHasQueryBlock((promptTemplate as PromptItem).text),
      context: hasSetContext,
    }
  }, [inputs.prompt_template, isChatMode, isChatModel])

  const shouldShowContextTip = !hasSetBlockStatus.context && inputs.context.enabled

  const isShowVars = useMemo(() => {
    if (isChatModel)
      return (inputs.prompt_template as PromptItem[]).some(item => item.edition_type === EditionType.jinja2)

    return (inputs.prompt_template as PromptItem).edition_type === EditionType.jinja2
  }, [inputs.prompt_template, isChatModel])

  const handleAddEmptyVariable = useCallback(() => {
    const nextInputs = produce(inputRef.current, (draft) => {
      const promptConfig = ensurePromptConfig(draft)
      promptConfig.jinja2_variables.push({
        variable: '',
        value_selector: [],
      })
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleAddVariable = useCallback((payload: Variable) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      const promptConfig = ensurePromptConfig(draft)
      promptConfig.jinja2_variables.push(payload)
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleVarListChange = useCallback((newList: Variable[]) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      const promptConfig = ensurePromptConfig(draft)
      promptConfig.jinja2_variables = newList
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleVarNameChange = useCallback((oldName: string, newName: string) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      if (isChatModel) {
        const promptTemplate = draft.prompt_template as PromptItem[]
        promptTemplate
          .filter(item => item.edition_type === EditionType.jinja2)
          .forEach((item) => {
            item.jinja2_text = (item.jinja2_text || '').replaceAll(`{{ ${oldName} }}`, `{{ ${newName} }}`)
          })
        return
      }

      const promptTemplate = draft.prompt_template as PromptItem
      if (promptTemplate.edition_type !== EditionType.jinja2)
        return

      promptTemplate.jinja2_text = (promptTemplate.jinja2_text || '').replaceAll(`{{ ${oldName} }}`, `{{ ${newName} }}`)
    })
    setInputs(nextInputs)
  }, [inputRef, isChatModel, setInputs])

  const handleContextVarChange = useCallback((newVar: ValueSelector | string) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.context.variable_selector = (newVar as ValueSelector) || []
      draft.context.enabled = !!(newVar && newVar.length > 0)
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handlePromptChange = useCallback((newPrompt: PromptItem[] | PromptItem) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.prompt_template = newPrompt
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleMemoryChange = useCallback((newMemory?: Memory) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.memory = newMemory
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  const handleSyeQueryChange = useCallback((newQuery: string) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      if (!draft.memory) {
        draft.memory = {
          window: {
            enabled: false,
            size: 10,
          },
          query_prompt_template: newQuery,
        }
        return
      }

      draft.memory.query_prompt_template = newQuery
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  return {
    hasSetBlockStatus,
    shouldShowContextTip,
    isShowVars,
    handleAddEmptyVariable,
    handleAddVariable,
    handleVarListChange,
    handleVarNameChange,
    handleContextVarChange,
    handlePromptChange,
    handleMemoryChange,
    handleSyeQueryChange,
    filterInputVar,
    filterJinja2InputVar,
    filterVar: filterMemoryPromptVar,
  }
}

export default useLLMPromptConfig