dify/web/app/components/base/prompt-editor/plugins/agent-output-block/utils.ts

import type {
  DeclaredOutputConfig,
  DeclaredOutputType,
} from '@dify/contracts/api/console/apps/types.gen'

export type AgentOutputTypeOptionValue =
  | DeclaredOutputType
  | 'array[boolean]'
  | 'array[file]'
  | 'array[number]'
  | 'array[object]'
  | 'array[string]'

export type AgentOutputTypeOption = {
  label: string
  type: DeclaredOutputType
  value: AgentOutputTypeOptionValue
  arrayItemType?: DeclaredOutputType
}

const AGENT_OUTPUT_TOKEN_REGEX = /\[§output:([^:§\]]+):([^:§\]]+)§\]/
const LEGACY_AGENT_OUTPUT_TOKEN_REGEX = /§output:([^:§\]]+):([^:§\]]+)§/
export const AGENT_OUTPUT_NAME_PATTERN = /^[a-z_]\w*$/i

export function getAgentOutputToken(name: string) {
  return `[§output:${name}:${name}§]`
}

export function parseAgentOutputToken(text: string) {
  const match = AGENT_OUTPUT_TOKEN_REGEX.exec(text) ?? LEGACY_AGENT_OUTPUT_TOKEN_REGEX.exec(text)
  if (!match) return null

  return {
    name: match[1]!,
    start: match.index,
    end: match.index + match[0].length,
  }
}

export function extractAgentOutputNames(text: string) {
  const names = new Set<string>()
  const bracketedRegex = new RegExp(AGENT_OUTPUT_TOKEN_REGEX.source, 'g')
  const legacyRegex = new RegExp(LEGACY_AGENT_OUTPUT_TOKEN_REGEX.source, 'g')

  for (const match of text.matchAll(bracketedRegex)) names.add(match[1]!)

  for (const match of text.matchAll(legacyRegex)) names.add(match[1]!)

  return names
}

export function replaceAgentOutputName(text: string, oldName: string, nextName: string) {
  const replaceTokenName = (match: string, name: string, _mirrorName: string) => {
    if (name !== oldName) return match

    return match.startsWith('[')
      ? getAgentOutputToken(nextName)
      : `§output:${nextName}:${nextName}§`
  }

  return text
    .replace(new RegExp(AGENT_OUTPUT_TOKEN_REGEX.source, 'g'), replaceTokenName)
    .replace(new RegExp(LEGACY_AGENT_OUTPUT_TOKEN_REGEX.source, 'g'), replaceTokenName)
}

export const AGENT_OUTPUT_TYPE_OPTIONS: AgentOutputTypeOption[] = [
  { value: 'string', label: 'string', type: 'string' },
  { value: 'number', label: 'number', type: 'number' },
  { value: 'boolean', label: 'boolean', type: 'boolean' },
  { value: 'object', label: 'object', type: 'object' },
  { value: 'file', label: 'file', type: 'file' },
  { value: 'array[string]', label: 'array[string]', type: 'array', arrayItemType: 'string' },
  { value: 'array[number]', label: 'array[number]', type: 'array', arrayItemType: 'number' },
  { value: 'array[boolean]', label: 'array[boolean]', type: 'array', arrayItemType: 'boolean' },
  { value: 'array[object]', label: 'array[object]', type: 'array', arrayItemType: 'object' },
  { value: 'array[file]', label: 'array[file]', type: 'array', arrayItemType: 'file' },
]

export function getAgentOutputTypeOption(value: AgentOutputTypeOptionValue) {
  return (
    AGENT_OUTPUT_TYPE_OPTIONS.find((option) => option.value === value) ||
    AGENT_OUTPUT_TYPE_OPTIONS[0]!
  )
}

export function getAgentOutputTypeOptionValue(
  output: DeclaredOutputConfig,
): AgentOutputTypeOptionValue {
  if (output.type !== 'array') return output.type

  return `array[${output.array_item?.type || 'object'}]` as AgentOutputTypeOptionValue
}

export function createAgentOutputConfig(
  name: string,
  type: AgentOutputTypeOptionValue,
): DeclaredOutputConfig {
  const option = getAgentOutputTypeOption(type)
  const output: DeclaredOutputConfig = {
    name: name.trim(),
    type: option.type,
    required: false,
  }

  if (option.type === 'array') {
    output.array_item = {
      type: option.arrayItemType || 'object',
    }
  }

  if (option.type === 'file') {
    output.file = {
      extensions: [],
      mime_types: [],
    }
  }

  return output
}

export function getUniqueAgentOutputName(outputs: DeclaredOutputConfig[]) {
  const outputNames = new Set(outputs.map((output) => output.name))
  const baseName = 'output'
  if (!outputNames.has(baseName)) return baseName

  let index = 1
  while (outputNames.has(`${baseName}_${index}`)) index += 1

  return `${baseName}_${index}`
}