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

import type { MutableRefObject } from 'react'
import type { LLMNodeType, StructuredOutput } from '../types'
import { produce } from 'immer'
import {
  useCallback,
  useState,
} from 'react'
import {
  ModelFeatureEnum,
  ModelTypeEnum,
} from '@/app/components/header/account-setting/model-provider-page/declarations'
import { useModelList } from '@/app/components/header/account-setting/model-provider-page/hooks'

type Params = {
  id: string
  model: LLMNodeType['model']
  inputRef: MutableRefObject<LLMNodeType>
  setInputs: (inputs: LLMNodeType) => void
  deleteNodeInspectorVars: (nodeId: string) => void
}

const useLLMStructuredOutputConfig = ({
  id,
  model,
  inputRef,
  setInputs,
  deleteNodeInspectorVars,
}: Params) => {
  const { data: modelList } = useModelList(ModelTypeEnum.textGeneration)
  const isModelSupportStructuredOutput = modelList
    ?.find(providerItem => providerItem.provider === model?.provider)
    ?.models
    .find(modelItem => modelItem.model === model?.name)
    ?.features
    ?.includes(ModelFeatureEnum.StructuredOutput)

  const [structuredOutputCollapsed, setStructuredOutputCollapsed] = useState(true)

  const handleStructureOutputEnableChange = useCallback((enabled: boolean) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.structured_output_enabled = enabled
    })
    setInputs(nextInputs)
    if (enabled)
      setStructuredOutputCollapsed(false)
    deleteNodeInspectorVars(id)
  }, [deleteNodeInspectorVars, id, inputRef, setInputs])

  const handleStructureOutputChange = useCallback((newOutput: StructuredOutput) => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.structured_output = newOutput
    })
    setInputs(nextInputs)
    deleteNodeInspectorVars(id)
  }, [deleteNodeInspectorVars, id, inputRef, setInputs])

  const handleReasoningFormatChange = useCallback((reasoningFormat: 'tagged' | 'separated') => {
    const nextInputs = produce(inputRef.current, (draft) => {
      draft.reasoning_format = reasoningFormat
    })
    setInputs(nextInputs)
  }, [inputRef, setInputs])

  return {
    isModelSupportStructuredOutput,
    structuredOutputCollapsed,
    setStructuredOutputCollapsed,
    handleStructureOutputEnableChange,
    handleStructureOutputChange,
    handleReasoningFormatChange,
  }
}

export default useLLMStructuredOutputConfig