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
const AGENT_OUTPUT_FILE_NAME_PATTERN = /^[^.][^/\\:*?"<>|\n\r]*\.([a-z0-9]{1,16})$/i
const AGENT_OUTPUT_FILE_EXTENSION_WHITELIST = new Set([
  'amr',
  'csv',
  'doc',
  'docx',
  'eml',
  'epub',
  'gif',
  'html',
  'jpeg',
  'jpg',
  'm4a',
  'markdown',
  'md',
  'mdx',
  'mov',
  'mp3',
  'mp4',
  'mpeg',
  'mpga',
  'msg',
  'pdf',
  'png',
  'ppt',
  'pptx',
  'svg',
  'txt',
  'wav',
  'webm',
  'webp',
  'xls',
  'xlsx',
  'xml',
])

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.replace(`${name}:${mirrorName}`, `${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]!
}

function isSupportedAgentOutputFileName(name: string) {
  const match = AGENT_OUTPUT_FILE_NAME_PATTERN.exec(name.trim())
  if (!match)
    return false

  return AGENT_OUTPUT_FILE_EXTENSION_WHITELIST.has(match[1]!.toLowerCase())
}

export function inferAgentOutputType(name: string, fallbackType: AgentOutputTypeOptionValue): AgentOutputTypeOptionValue {
  if (isSupportedAgentOutputFileName(name))
    return 'file'

  return fallbackType
}

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}`
}