dify/web/features/agent-v2/agent-detail/configure/hooks.ts

'use client'

import type {
  AgentConfigSnapshotDetailResponse,
  AgentSoulConfig,
} from '@dify/contracts/api/console/agent/types.gen'
import { skipToken, useQuery } from '@tanstack/react-query'
import { useAtom, useAtomValue } from 'jotai'
import { ModelTypeEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
import {
  useDefaultModel,
  useTextGenerationCurrentProviderAndModelAndModelList,
} from '@/app/components/header/account-setting/model-provider-page/hooks'
import { agentComposerAppFeaturesAtom } from '@/features/agent-v2/agent-composer/store-modules/app-features'
import { agentComposerModelAtom } from '@/features/agent-v2/agent-composer/store-modules/model'
import { consoleQuery } from '@/service/client'

export function useAgentConfigureData(agentId: string, selectedVersionId: string | null) {
  const agentQuery = useQuery(
    consoleQuery.agent.byAgentId.get.queryOptions({
      input: {
        params: {
          agent_id: agentId,
        },
      },
    }),
  )
  const composerQuery = useQuery(
    consoleQuery.agent.byAgentId.composer.get.queryOptions({
      input: {
        params: {
          agent_id: agentId,
        },
      },
    }),
  )
  const publishedVersionId = composerQuery.data?.active_config_snapshot?.id
  const shouldLoadPublishedVersion = !selectedVersionId && !composerQuery.data?.agent_soul
  const versionIdToLoad =
    selectedVersionId ?? (shouldLoadPublishedVersion ? publishedVersionId : undefined)
  const shouldLoadVersion = !!versionIdToLoad
  const versionQuery = useQuery(
    consoleQuery.agent.byAgentId.versions.byVersionId.get.queryOptions({
      input: versionIdToLoad
        ? {
            params: {
              agent_id: agentId,
              version_id: versionIdToLoad,
            },
          }
        : skipToken,
    }),
  )
  const versionDetail = versionQuery.data as AgentConfigSnapshotDetailResponse | undefined
  const activeVersionId =
    selectedVersionId ?? (shouldLoadPublishedVersion ? publishedVersionId : null)
  const activeConfigSnapshot = selectedVersionId
    ? versionDetail
    : (composerQuery.data?.active_config_snapshot ?? versionDetail)
  const agentSoulConfig = selectedVersionId
    ? versionDetail?.config_snapshot
    : (composerQuery.data?.agent_soul ?? versionDetail?.config_snapshot)
  const isPending =
    agentQuery.isPending || composerQuery.isPending || (shouldLoadVersion && versionQuery.isPending)

  return {
    agentQuery,
    composerQuery,
    versionQuery,
    shouldLoadVersion,
    selectedVersionId,
    activeVersionId,
    activeConfigSnapshot,
    agentSoulConfig,
    isPending,
  }
}

export function useAgentConfigureModelOptions() {
  const [model, setModel] = useAtom(agentComposerModelAtom)
  const { data: defaultTextGenerationModel } = useDefaultModel(ModelTypeEnum.textGeneration)
  const defaultModel = defaultTextGenerationModel
    ? {
        provider: defaultTextGenerationModel.provider.provider,
        model: defaultTextGenerationModel.model,
      }
    : undefined
  const currentModel = model ?? defaultModel
  const { textGenerationModelList } =
    useTextGenerationCurrentProviderAndModelAndModelList(currentModel)

  return {
    currentModel,
    setConfigureModel: setModel,
    textGenerationModelList,
  }
}

export function useAgentPreviewSoulConfig(agentSoulConfig: AgentSoulConfig | undefined) {
  const draftAppFeatures = useAtomValue(agentComposerAppFeaturesAtom)

  if (!agentSoulConfig || !draftAppFeatures) return agentSoulConfig

  return {
    ...agentSoulConfig,
    app_features: draftAppFeatures,
  }
}