EasyFlow/easyflow-ui-admin/packages/tinyflow-ui/src/components/utils/knowledgeNode.ts

import type { Parameter, SelectItem } from '#types';

export function normalizeKnowledgeIds(data: Record<string, any>): string[] {
  if (Array.isArray(data.knowledgeIds)) {
    const seen = new Set<string>();
    return data.knowledgeIds.flatMap((id: unknown) => {
      if (id === null || id === undefined) return [];
      const value = String(id).trim();
      if (!value || seen.has(value)) return [];
      seen.add(value);
      return [value];
    });
  }
  return data.knowledgeId === null || data.knowledgeId === undefined || String(data.knowledgeId).trim() === ''
    ? []
    : [String(data.knowledgeId).trim()];
}

export function toggleKnowledgeId(
  selectedIds: Array<number | string>,
  value: number | string,
): string[] {
  const key = String(value);
  return selectedIds.some((id) => String(id) === key)
    ? selectedIds.filter((id) => String(id) !== key).map(String)
    : [...selectedIds.map(String), key];
}

export function applyEmbeddingCompatibility(
  items: SelectItem[],
  selectedIds: Array<number | string>,
): SelectItem[] {
  const selectedKeys = new Set(selectedIds.map(String));
  const normalizedItems = items.map((item) => ({ ...item, value: String(item.value) }));
  for (const selectedId of selectedIds.map(String)) {
    if (!normalizedItems.some((item) => String(item.value) === selectedId)) {
      normalizedItems.push({
        value: selectedId,
        label: `已失效知识库（${selectedId}）`,
        selectable: true,
        disabledReason: '知识库不存在或已无权限，仅可移除',
      });
    }
  }
  const firstSelected = normalizedItems.find(
    (item) => String(item.value) === String(selectedIds[0]),
  );
  if (!firstSelected) return normalizedItems;

  const vectorState = (item: SelectItem): 'available' | 'unavailable' | 'unknown' => {
    if (item.vectorStoreEnabled === false) return 'unavailable';
    if (
      item.vectorStoreEnabled === true
      && item.embeddingModelId !== null
      && item.embeddingModelId !== undefined
      && item.embeddingDimension !== null
      && item.embeddingDimension !== undefined
    ) {
      return 'available';
    }
    return 'unknown';
  };

  const firstState = vectorState(firstSelected);
  const isCompatible = (item: SelectItem) => firstState === 'available'
    && vectorState(item) === 'available'
    && firstSelected.embeddingModelId === item.embeddingModelId
    && firstSelected.embeddingDimension === item.embeddingDimension;

  return normalizedItems.map((item) => {
    if (selectedKeys.has(String(item.value))) {
      if (item.disabledReason) return item;
      const itemState = vectorState(item);
      if (itemState === 'unavailable') {
        return {
          ...item,
          selectable: true,
          disabledReason: '未启用向量检索，仅可移除',
        };
      }
      if (
        firstState === 'available'
        && itemState === 'available'
        && !isCompatible(item)
      ) {
        return {
          ...item,
          selectable: true,
          disabledReason: '与首个知识库的 Embedding 配置不一致，仅可移除',
        };
      }
      return { ...item, selectable: true, disabledReason: undefined };
    }
    const itemState = vectorState(item);
    if (firstSelected.disabledReason || firstState === 'unavailable') {
      return {
        ...item,
        selectable: false,
        disabledReason: '请先移除未启用向量检索的知识库',
      };
    }
    if (itemState === 'unavailable') {
      return {
        ...item,
        selectable: false,
        disabledReason: '未启用向量检索',
      };
    }
    if (firstState === 'unknown' || itemState === 'unknown') {
      return { ...item, selectable: true, disabledReason: undefined };
    }
    return isCompatible(item)
      ? { ...item, selectable: true, disabledReason: undefined }
      : {
          ...item,
          selectable: false,
          disabledReason: '与首个知识库的 Embedding 模型或向量维度不一致',
        };
  });
}

export function buildKnowledgeSelectItems(
  items: SelectItem[],
  selectedIds: Array<number | string>,
  reloadItem?: SelectItem,
): SelectItem[] {
  const compatibleItems = applyEmbeddingCompatibility(items, selectedIds);
  return reloadItem
    ? [
        ...compatibleItems,
        { ...reloadItem, selectable: true, disabledReason: undefined },
      ]
    : compatibleItems;
}

const locked = (parameter: Parameter): Parameter => ({
  ...parameter,
  nameDisabled: true,
  dataTypeDisabled: true,
  deleteDisabled: true,
});

const expectedOutputDefs: Parameter[] = [
  locked({
    id: 'knowledge_documents',
    name: 'documents',
    dataType: 'Array',
    addChildDisabled: true,
    children: [
      locked({ id: 'knowledge_documents_title', name: 'title', dataType: 'String' }),
      locked({ id: 'knowledge_documents_content', name: 'content', dataType: 'String' }),
      locked({ id: 'knowledge_documents_document_id', name: 'documentId', dataType: 'Number' }),
      locked({ id: 'knowledge_documents_knowledge_id', name: 'knowledgeId', dataType: 'Number' }),
    ],
  }),
];

function mergeParameter(current: Parameter | undefined, expected: Parameter): Parameter {
  if (!current) {
    return {
      ...expected,
      children: expected.children?.map((child) => ({ ...child })),
    };
  }
  const result = { ...current, ...expected, id: current.id || expected.id };
  if (expected.children) {
    const currentChildren = current.children || [];
    result.children = expected.children.map((child) =>
      mergeParameter(
        currentChildren.find((item) => item.name === child.name),
        child,
      ),
    );
  }
  return result;
}

export function ensureKnowledgeOutputDefs(outputDefs?: Parameter[]): Parameter[] {
  const current = outputDefs || [];
  return expectedOutputDefs.map((expected) =>
    mergeParameter(
      current.find((item) => item.name === expected.name),
      expected,
    ),
  );
}