ragflow/web/src/components/structure-graph/adapters.tsx

/*
 *  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

import trim from 'lodash/trim';
import { ClaimBadge } from '@/components/structure-graph/claim-badge';
import { type TreeDataItem } from '@/components/ui/tree-view';
import {
  type IArtifactGraph,
  type IArtifactGraphEntity,
} from '@/interfaces/database/dataset';
import {
  type IClaimEvidence,
  type IStructureGraphEntity,
  type IStructureGraphRelation,
  type IStructureGraphTemplate,
} from '@/interfaces/database/document-structure';
import { type TreeData } from '@antv/g6/lib/types';

declare module '@/components/ui/tree-view' {
  interface TreeDataItem {
    source_chunk_ids?: string[];
    /** page_index fact/conclusion: gate-verified quotes for the detail panel. */
    description?: string;
    evidence?: IClaimEvidence[];
  }
}

export function getEntityDisplayName(entity: IStructureGraphEntity) {
  return trim(entity.name ?? entity.id ?? '');
}

// Exact (case-insensitive) entity name/alias hit for a typed search keyword.
// Returns the canonical display name so callers can treat the keyword like a
// dropdown selection (highlight + dim); '' when the keyword only partially
// matches or names nothing, in which case the server-side keyword subgraph
// search stays the right fallback.
export function findEntityDisplayNameByKeyword(
  entities: IStructureGraphEntity[],
  keyword: string,
) {
  const query = keyword.trim().toLowerCase();
  if (!query) {
    return '';
  }
  const entity = (entities ?? []).find(
    (item) =>
      getEntityDisplayName(item).toLowerCase() === query ||
      (item.aliases ?? []).some(
        (alias) => (alias ?? '').trim().toLowerCase() === query,
      ),
  );
  return entity ? getEntityDisplayName(entity) : '';
}

function normalizeEntity(entity: IStructureGraphEntity) {
  const id = entity.id ?? entity.name ?? '';
  const name = getEntityDisplayName(entity);
  return { ...entity, id, name };
}

function getEntityDescription(entity: IStructureGraphEntity): string {
  return entity.description ?? entity.discription ?? '';
}

function buildTreeDataItems(
  entities: IStructureGraphEntity[],
  relations: IStructureGraphRelation[],
  relationTypes: string[],
  showEntityType = false,
): TreeDataItem[] {
  const normalized = (entities ?? [])
    .map(normalizeEntity)
    .filter((entity) => entity.id);
  const map = new Map<string, TreeDataItem>(
    normalized.map((entity) => [
      entity.id,
      {
        id: entity.id,
        name: entity.name,
        entityType: showEntityType ? entity.type : undefined,
        source_chunk_ids: entity.source_chunk_ids,
        // Leaf clusters only: the tree shows this as a count badge and the
        // claims themselves are fetched on click. The shared TreeView's
        // badge slot is a ReactNode, so the claim domain owns the rendering.
        badge: <ClaimBadge value={entity.claim_count} />,
        // page_index fact/conclusion: gate-verified quotes rendered in the
        // node's detail panel. Without badge/evidence/description the tree and
        // page_index nodes reached the click handler empty, so the detail panel
        // could never open for them.
        evidence: entity.evidence,
        description: getEntityDescription(entity),
      },
    ]),
  );
  const childIds = new Set<string>();

  for (const relation of relations ?? []) {
    if (!relationTypes.includes(relation.type ?? '')) continue;
    // Self-referencing relation (same entity as its own child) creates
    // an infinite recursion in the tree renderer.  This is a backend
    // data integrity issue (duplicate entity names) but we defend
    // against it in the frontend so the UI never hangs.
    if (relation.from === relation.to) continue;

    const parent = map.get(relation.from);
    const child = map.get(relation.to);
    if (!parent || !child) continue;

    childIds.add(child.id);
    parent.children = parent.children ?? [];
    parent.children.push(child);
  }

  return normalized
    .filter((entity) => !childIds.has(entity.id))
    .map((entity) => map.get(entity.id))
    .filter((item): item is TreeDataItem => item !== undefined);
}

function buildUniqueTreeDataItems(
  entities: IStructureGraphEntity[],
  relations: IStructureGraphRelation[],
  relationTypes: string[],
): TreeDataItem[] {
  const normalized = [
    ...new Map(
      (entities ?? [])
        .map(normalizeEntity)
        .filter((entity) => entity.id)
        .map((entity) => [entity.id, entity]),
    ).values(),
  ];
  const map = new Map<string, TreeDataItem>(
    normalized.map((entity) => [
      entity.id,
      {
        id: entity.id,
        name: entity.name,
        entityType: entity.type,
        source_chunk_ids: entity.source_chunk_ids,
        // Leaf clusters only: the tree UI shows this as a count badge and
        // fetches the claims themselves on expand.
        badge: entity.claim_count,
        // page_index fact/conclusion: gate-verified quotes shown in the
        // node's detail panel.
        evidence: entity.evidence,
      },
    ]),
  );
  const childIds = new Set<string>();
  const parentMap = new Map<string, string>();

  for (const relation of relations ?? []) {
    if (!relationTypes.includes(relation.type ?? '')) continue;

    const parent = map.get(relation.from);
    const child = map.get(relation.to);
    if (!parent || !child) continue;
    if (childIds.has(child.id)) continue;

    // Avoid cycles: do not attach a node under one of its descendants.
    let cursor = relation.from;
    let wouldCycle = false;
    while (parentMap.has(cursor)) {
      const ancestor = parentMap.get(cursor)!;
      if (ancestor === child.id) {
        wouldCycle = true;
        break;
      }
      cursor = ancestor;
    }
    if (wouldCycle) continue;

    childIds.add(child.id);
    parentMap.set(child.id, relation.from);
    parent.children = parent.children ?? [];
    parent.children.push(child);
  }

  return normalized
    .filter((entity) => !childIds.has(entity.id))
    .map((entity) => map.get(entity.id))
    .filter((item): item is TreeDataItem => item !== undefined);
}

// Entity types that are never part of the navigation tree. A claim is an
// atomic, evidence-bearing proposition: it answers "what does this section
// say", which is the middle column's job, not the outline's. Leaving claims in
// the tree buried the heading hierarchy under hundreds of flat leaves.
// ``fact``/``conclusion`` are the pre-rename spellings of the same object, kept
// so templates compiled before the rename render the same way.
const PAGE_INDEX_NON_TREE_TYPES = new Set(['claim', 'fact', 'conclusion']);

export function adaptPageIndexToTreeData(
  template: IStructureGraphTemplate,
): TreeDataItem[] {
  // Filtered before building: relations pointing at a removed entity are
  // dropped by the builder, so headings whose only children were claims simply
  // become leaves instead of dangling.
  const entities = (template.entities ?? []).filter(
    (entity) =>
      !PAGE_INDEX_NON_TREE_TYPES.has(
        String((entity as { type?: string })?.type ?? '')
          .trim()
          .toLowerCase(),
      ),
  );
  return buildTreeDataItems(entities, template.relations, ['include'], true);
}

export function adaptTreeToTreeData(
  template: IStructureGraphTemplate,
): TreeDataItem[] {
  return buildTreeDataItems(template.entities, template.relations, ['child']);
}

export function adaptKnowledgeGraphToForceGraph(
  template: IStructureGraphTemplate,
): IArtifactGraph {
  const entities: IArtifactGraphEntity[] = (template.entities ?? []).map(
    (entity) => {
      const normalized = normalizeEntity(entity);
      return {
        slug: normalized.id,
        name: normalized.name,
        aliases: normalized.aliases ?? [],
        description: getEntityDescription(normalized),
        type: normalized.type ?? '',
        weight: normalized.mention_count ?? 0,
        source_chunk_ids: normalized.source_chunk_ids,
      };
    },
  );

  const entityNames = new Set(entities.map((entity) => entity.slug));

  return {
    entities,
    relations: (template.relations ?? [])
      .filter(
        (relation) =>
          // Only keep relations whose source and target entities both exist in the graph.
          entityNames.has(relation.from) && entityNames.has(relation.to),
      )
      .map((relation) => ({
        from: relation.from,
        to: relation.to,
        type: relation.type ?? '',
      })),
  };
}

function treeDataItemToG6TreeData(item: TreeDataItem): TreeData {
  const node: TreeData = {
    id: item.id,
    data: {
      name: item.name,
      source_chunk_ids: item.source_chunk_ids,
    },
  };

  if (item.children && item.children.length > 0) {
    node.children = item.children.map(treeDataItemToG6TreeData);
  }

  return node;
}

export interface TimelineX6NodeData {
  id: string;
  shape: 'rect' | 'circle';
  width: number;
  height: number;
  label: string;
  data?: IStructureGraphEntity;
  attrs?: Record<string, unknown>;
}

export interface TimelineX6EdgeData {
  id: string;
  shape: 'edge';
  source: string;
  target: string;
  attrs?: Record<string, unknown>;
}

export function adaptTimelineToX6Data(template: IStructureGraphTemplate): {
  nodes: TimelineX6NodeData[];
  edges: TimelineX6EdgeData[];
} {
  const normalized = template.entities
    .map(normalizeEntity)
    .filter((entity) => entity.id);
  const entityIds = new Set(normalized.map((entity) => entity.id));

  const nodes = normalized.map((entity) => {
    const isTimestamp = entity.type === 'timestamp';
    return {
      id: entity.id,
      shape: isTimestamp ? ('circle' as const) : ('rect' as const),
      width: isTimestamp ? 40 : 200,
      height: isTimestamp ? 40 : 80,
      label: entity.name,
      data: entity,
      attrs: {
        body: {
          fill: isTimestamp ? 'rgb(var(--accent-primary))' : 'transparent',
          stroke: 'rgb(var(--accent-primary))',
          rx: isTimestamp ? 48 : 8,
          ry: isTimestamp ? 48 : 8,
        },
        label: {
          fill: isTimestamp
            ? 'rgb(var(--accent-primary))'
            : 'rgb(var(--text-primary))',
          fontSize: 12,
          ...(isTimestamp && {
            refX: '50%',
            refY: 0,
            refY2: -8,
            textAnchor: 'middle',
            textVerticalAnchor: 'bottom',
          }),
          textWrap: {
            width: isTimestamp ? 80 : 180,
            height: isTimestamp ? 80 : 64,
            ellipsis: true,
          },
        },
      },
    };
  });

  const edges = (template.relations ?? [])
    .filter(
      (relation) => entityIds.has(relation.from) && entityIds.has(relation.to),
    )
    .map((relation, index) => ({
      id: `timeline-edge-${index}`,
      shape: 'edge' as const,
      source: relation.from,
      target: relation.to,
      attrs: {
        line: {
          stroke: '#8c8c8c',
          strokeWidth: 1,
          targetMarker: 'classic',
        },
      },
    }));

  return { nodes, edges };
}

export function adaptMindMapToIndentedTree(
  template: IStructureGraphTemplate,
): TreeData {
  const roots = buildUniqueTreeDataItems(
    template.entities,
    template.relations,
    // Python's mindmap structure-graph projection stores parent-child edges
    // as generic "related" relations. Keep the branch spellings for older
    // data, but accept the canonical Python form as well.
    ['related', 'has_branch', 'has_sub_branch'],
  );

  const g6Roots = roots.map(treeDataItemToG6TreeData);

  if (g6Roots.length === 1) {
    return g6Roots[0]!;
  }

  return {
    id: 'mindmap-root',
    children: g6Roots,
  };
}