Feat: Adjust the color and size of the graph based on the data. (#16868)

web/src/components/artifact-force-graph/index.tsx

@@ -1,11 +1,17 @@
 import { type IArtifactGraphEntity } from '@/interfaces/database/dataset';
 import { cn } from '@/lib/utils';
-import isEmpty from 'lodash/isEmpty';
-import { memo, useCallback, useEffect, useMemo, useRef } from 'react';
+import { memo, useCallback, useEffect, useRef } from 'react';
 import ForceGraph2D, { type ForceGraphMethods } from 'react-force-graph-2d';
-import { type ArtifactForceGraphProps } from './types';
+import { renderNodeLabel } from './node-label';
+import {
+  getNodeColor as defaultGetNodeColor,
+  getNodeRadius as defaultGetNodeRadius,
+  MinNodeRadius,
+} from './node-style';
+import { type ArtifactForceGraphProps, type ArtifactGraphNode } from './types';
+import { useArtifactGraphData } from './use-artifact-graph-data';
 import { useContainerDimensions } from './use-container-dimensions';
-import { defaultMapNodeToValue, renderNodeLabel } from './utils';
+import { defaultMapNodeToValue } from './utils';
 
 function ArtifactForceGraph<TNodeValue = IArtifactGraphEntity>({
   data,
@@ -15,31 +21,22 @@ function ArtifactForceGraph<TNodeValue = IArtifactGraphEntity>({
     node: IArtifactGraphEntity,
   ) => TNodeValue,
   getNodeId = (node) => node.slug,
+  getNodeColor = defaultGetNodeColor,
+  getNodeRadius = defaultGetNodeRadius,
 }: ArtifactForceGraphProps<TNodeValue>) {
   const containerRef = useRef<HTMLDivElement>(null);
-  const fgRef = useRef<ForceGraphMethods<IArtifactGraphEntity> | undefined>(
+  const fgRef = useRef<ForceGraphMethods<ArtifactGraphNode> | undefined>(
     undefined,
   );
   const hasFittedRef = useRef(false);
   const dimensions = useContainerDimensions(containerRef, show);
 
-  const graphData = useMemo(() => {
-    if (isEmpty(data) || !data) {
-      return { nodes: [], links: [] };
-    }
-
-    const nodes = (data.entities || []).map((entity) => ({
-      ...entity,
-      id: getNodeId(entity),
-    }));
-
-    const links = (data.relations || []).map((relation) => ({
-      source: relation.from,
-      target: relation.to,
-    }));
-
-    return { nodes, links };
-  }, [data, getNodeId]);
+  const graphData = useArtifactGraphData({
+    data,
+    getNodeId,
+    getNodeColor,
+    getNodeRadius,
+  });
 
   useEffect(() => {
     hasFittedRef.current = false;
@@ -69,6 +66,13 @@ function ArtifactForceGraph<TNodeValue = IArtifactGraphEntity>({
       .trim();
   }, []);
 
+  const nodeColor = useCallback((node: ArtifactGraphNode) => node.__color, []);
+
+  const nodeVal = useCallback(
+    (node: ArtifactGraphNode) => node.__radius ?? MinNodeRadius,
+    [],
+  );
+
   return (
     <div
       ref={containerRef}
@@ -80,7 +84,9 @@ function ArtifactForceGraph<TNodeValue = IArtifactGraphEntity>({
           width={dimensions.width}
           height={dimensions.height}
           graphData={graphData}
-          nodeAutoColorBy="type"
+          nodeRelSize={1}
+          nodeColor={nodeColor}
+          nodeVal={nodeVal}
           cooldownTicks={100}
           nodeLabel={''}
           onEngineStop={handleEngineStop}

web/src/components/artifact-force-graph/node-label.ts

@@ -0,0 +1,25 @@
+import { type ComponentProps } from 'react';
+import ForceGraph2D from 'react-force-graph-2d';
+import { MinNodeRadius } from './node-style';
+import { type ArtifactGraphNode } from './types';
+
+export const renderNodeLabel: NonNullable<
+  ComponentProps<typeof ForceGraph2D>['nodeCanvasObject']
+> = (node, ctx) => {
+  const graphNode = node as ArtifactGraphNode;
+  const label = graphNode.name;
+  const radius = graphNode.__radius ?? MinNodeRadius;
+  const fontSize = 2;
+  ctx.font = `${fontSize}px Sans-Serif`;
+  ctx.textAlign = 'center';
+  ctx.textBaseline = 'top';
+
+  const textSecondary = getComputedStyle(document.documentElement)
+    .getPropertyValue('--text-secondary')
+    .trim();
+  ctx.fillStyle = `rgb(${textSecondary})`;
+
+  if (typeof graphNode.x === 'number' && typeof graphNode.y === 'number') {
+    ctx.fillText(label, graphNode.x, graphNode.y + radius - 2);
+  }
+};

web/src/components/artifact-force-graph/node-style.ts

@@ -0,0 +1,25 @@
+import { type IArtifactGraphEntity } from '@/interfaces/database/dataset';
+
+export const EntityNodeColor = '#00BEB4';
+export const ConceptNodeColor = '#4CACFF';
+
+export const MinNodeRadius = 4;
+export const MaxNodeRadius = 14;
+
+export const getNodeColor = (node: IArtifactGraphEntity): string => {
+  if (node.type === 'entity') return EntityNodeColor;
+  if (node.type === 'concept') return ConceptNodeColor;
+  return EntityNodeColor;
+};
+
+export const getNodeRadius = (
+  node: IArtifactGraphEntity,
+  minWeight: number,
+  maxWeight: number,
+): number => {
+  const weight = node.weight ?? 0;
+  if (maxWeight <= minWeight) return MinNodeRadius;
+  const clamped = Math.max(minWeight, Math.min(maxWeight, weight));
+  const t = (clamped - minWeight) / (maxWeight - minWeight);
+  return MinNodeRadius + t * (MaxNodeRadius - MinNodeRadius);
+};

web/src/components/artifact-force-graph/types.ts

@@ -2,6 +2,13 @@ import {
   type IArtifactGraph,
   type IArtifactGraphEntity,
 } from '@/interfaces/database/dataset';
+import { type NodeObject } from 'react-force-graph-2d';
+
+export type ArtifactGraphNode = NodeObject<IArtifactGraphEntity> & {
+  id: string;
+  __color: string;
+  __radius: number;
+};
 
 export interface ArtifactForceGraphProps<TNodeValue = IArtifactGraphEntity> {
   data?: IArtifactGraph;
@@ -9,4 +16,10 @@ export interface ArtifactForceGraphProps<TNodeValue = IArtifactGraphEntity> {
   onNodeClick?: (node: TNodeValue) => void;
   mapNodeToValue?: (node: IArtifactGraphEntity) => TNodeValue;
   getNodeId?: (node: IArtifactGraphEntity) => string;
+  getNodeColor?: (node: IArtifactGraphEntity) => string;
+  getNodeRadius?: (
+    node: IArtifactGraphEntity,
+    minWeight: number,
+    maxWeight: number,
+  ) => number;
 }

web/src/components/artifact-force-graph/use-artifact-graph-data.ts

@@ -0,0 +1,57 @@
+import {
+  type IArtifactGraph,
+  type IArtifactGraphEntity,
+} from '@/interfaces/database/dataset';
+import isEmpty from 'lodash/isEmpty';
+import { useMemo } from 'react';
+import { type GraphData } from 'react-force-graph-2d';
+import {
+  getNodeColor as defaultGetNodeColor,
+  getNodeRadius as defaultGetNodeRadius,
+} from './node-style';
+import { type ArtifactGraphNode } from './types';
+
+export interface UseArtifactGraphDataOptions {
+  data?: IArtifactGraph;
+  getNodeId?: (node: IArtifactGraphEntity) => string;
+  getNodeColor?: (node: IArtifactGraphEntity) => string;
+  getNodeRadius?: (
+    node: IArtifactGraphEntity,
+    minWeight: number,
+    maxWeight: number,
+  ) => number;
+}
+
+export const useArtifactGraphData = ({
+  data,
+  getNodeId = (node) => node.slug,
+  getNodeColor = defaultGetNodeColor,
+  getNodeRadius = defaultGetNodeRadius,
+}: UseArtifactGraphDataOptions) => {
+  return useMemo<
+    GraphData<ArtifactGraphNode, { source: string; target: string }>
+  >(() => {
+    if (isEmpty(data) || !data) {
+      return { nodes: [], links: [] };
+    }
+
+    const entities = data.entities || [];
+    const weights = entities.map((entity) => entity.weight ?? 0);
+    const minWeight = Math.min(0, ...weights);
+    const maxWeight = Math.max(0, ...weights);
+
+    const nodes = entities.map((entity) => ({
+      ...entity,
+      id: getNodeId(entity),
+      __color: getNodeColor(entity),
+      __radius: getNodeRadius(entity, minWeight, maxWeight),
+    }));
+
+    const links = (data.relations || []).map((relation) => ({
+      source: relation.from,
+      target: relation.to,
+    }));
+
+    return { nodes, links };
+  }, [data, getNodeColor, getNodeId, getNodeRadius]);
+};

web/src/components/artifact-force-graph/utils.ts

@@ -1,25 +1,4 @@
 import { type IArtifactGraphEntity } from '@/interfaces/database/dataset';
-import { type ComponentProps } from 'react';
-import ForceGraph2D from 'react-force-graph-2d';
-
-export const renderNodeLabel: NonNullable<
-  ComponentProps<typeof ForceGraph2D>['nodeCanvasObject']
-> = (node, ctx) => {
-  const label = node.name;
-  const fontSize = 2;
-  ctx.font = `${fontSize}px Sans-Serif`;
-  ctx.textAlign = 'center';
-  ctx.textBaseline = 'top';
-
-  const textSecondary = getComputedStyle(document.documentElement)
-    .getPropertyValue('--text-secondary')
-    .trim();
-  ctx.fillStyle = `rgb(${textSecondary})`;
-
-  if (node.x && node.y) {
-    ctx.fillText(label, node.x, node.y + 5);
-  }
-};
 
 export const defaultMapNodeToValue = <TNode extends IArtifactGraphEntity>(
   node: TNode,

web/src/pages/agent/utils.ts

@@ -21,7 +21,6 @@ import {
   omit,
   sample,
 } from 'lodash';
-import pipe from 'lodash/fp/pipe';
 import isObject from 'lodash/isObject';
 import {
   AgentDialogueMode,
@@ -164,10 +163,7 @@ function buildCategorize(edges: Edge[], nodes: Node[], nodeId: string) {
 }
 
 const buildOperatorParams = (operatorName: string) =>
-  pipe(
-    removeUselessDataInTheOperator(operatorName),
-    // initializeOperatorParams(operatorName), // Final processing, for guarantee
-  );
+  removeUselessDataInTheOperator(operatorName);
 
 const ExcludeOperators = [Operator.Note, Operator.Tool, Operator.Placeholder];
 

web/src/pages/agents/index.tsx

@@ -101,7 +101,7 @@ export default function Agents() {
               value={filterValue}
             >
               <DropdownMenu>
-                <DropdownMenuTrigger data-testid="create-agent">
+                <DropdownMenuTrigger data-testid="create-agent" asChild>
                   <Button>
                     <Plus className="size-[1em]" />
                     {t('flow.createGraph')}