limit number of entities in the graph (#44)

Co-authored-by: Sorin Peste <[email protected]>

Author: neaorin

src/app/components/GraphDataHandler.tsx

@@ -30,6 +30,7 @@ const GraphDataHandler: React.FC = () => {
   const [includeTextUnits, setIncludeTextUnits] = useState(false);
   const [includeCommunities, setIncludeCommunities] = useState(false);
   const [includeCovariates, setIncludeCovariates] = useState(false);
+  const [maxEntities, setMaxEntities] = useState(500);
 
   const {
     entities,
@@ -54,7 +55,8 @@ const GraphDataHandler: React.FC = () => {
     includeDocuments,
     includeTextUnits,
     includeCommunities,
-    includeCovariates
+    includeCovariates,
+    maxEntities
   );
 
   const hasDocuments = documents.length > 0;
@@ -190,6 +192,9 @@ const GraphDataHandler: React.FC = () => {
             hasTextUnits={hasTextUnits}
             hasCommunities={hasCommunities}
             hasCovariates={hasCovariates}
+            maxEntities={maxEntities}
+            onMaxEntitiesChange={setMaxEntities}
+            totalEntities={entities.length}
           />
         </Box>
       )}

src/app/components/GraphViewer.tsx

@@ -13,6 +13,7 @@ import {
   FormControlLabel,
   FormGroup,
   IconButton,
+  Slider,
   Switch,
   Tooltip,
   Typography,
@@ -60,6 +61,9 @@ interface GraphViewerProps {
   hasTextUnits: boolean;
   hasCommunities: boolean;
   hasCovariates: boolean;
+  maxEntities: number;
+  onMaxEntitiesChange: React.Dispatch<React.SetStateAction<number>>;
+  totalEntities: number;
 }
 
 const NODE_R = 8;
@@ -82,6 +86,9 @@ const GraphViewer: React.FC<GraphViewerProps> = ({
   hasTextUnits,
   hasCommunities,
   hasCovariates,
+  maxEntities,
+  onMaxEntitiesChange,
+  totalEntities,
 }) => {
   const theme = useTheme();
   const [highlightNodes, setHighlightNodes] = useState<Set<CustomNode>>(
@@ -716,6 +723,29 @@ const GraphViewer: React.FC<GraphViewerProps> = ({
             label="Include Covariates"
           />
         </FormGroup>
+
+        {totalEntities > 0 && (
+          <Box sx={{ width: 200, mt: 2 }}>
+            <Typography variant="body2" gutterBottom>
+              Max Entities: {maxEntities === 0 ? "All" : maxEntities}
+              {totalEntities > 0 && ` / ${totalEntities}`}
+            </Typography>
+            <Slider
+              value={maxEntities}
+              onChange={(_, value) => onMaxEntitiesChange(value as number)}
+              min={0}
+              max={Math.max(totalEntities, 1000)}
+              step={50}
+              marks={[
+                { value: 0, label: "All" },
+                { value: 500, label: "500" },
+              ]}
+              valueLabelDisplay="auto"
+              valueLabelFormat={(value) => (value === 0 ? "All" : value)}
+              disabled={apiSearchResults !== null}
+            />
+          </Box>
+        )}
       </Box>
 
       <APISearchDrawer

src/app/hooks/useGraphData.ts

@@ -19,12 +19,32 @@ const useGraphData = (
   includeDocuments: boolean,
   includeTextUnits: boolean,
   includeCommunities: boolean,
-  includeCovariates: boolean
+  includeCovariates: boolean,
+  maxEntities: number = 0 // 0 means no limit
 ) => {
 
   const [graphData, setGraphData] = useState<CustomGraphData>({ nodes: [], links: [] });  
   useEffect(() => {
-    const nodes: CustomNode[] = entities.map((entity) => ({
+    // Calculate relationship count for each entity to determine top N
+    const entityRelationshipCount: { [key: string]: number } = {};
+    relationships.forEach((rel) => {
+      entityRelationshipCount[rel.source] = (entityRelationshipCount[rel.source] || 0) + 1;
+      entityRelationshipCount[rel.target] = (entityRelationshipCount[rel.target] || 0) + 1;
+    });
+
+    // Sort entities by relationship count and limit to maxEntities if specified
+    let filteredEntities = entities;
+    if (maxEntities > 0 && entities.length > maxEntities) {
+      filteredEntities = [...entities]
+        .sort((a, b) => {
+          const countA = entityRelationshipCount[a.title] || 0;
+          const countB = entityRelationshipCount[b.title] || 0;
+          return countB - countA; // Sort descending
+        })
+        .slice(0, maxEntities);
+    }
+
+    const nodes: CustomNode[] = filteredEntities.map((entity) => ({
       uuid: entity.id,
       id: entity.title, // use title as id because relationships use title as source/target
       name: entity.title, // legacy field for old GraphRAG 0.2.x - 0.3.x
@@ -111,13 +131,17 @@ const useGraphData = (
       const textUnitEntityLinks = textunits
         .filter((textunit) => (textunit.entity_ids ?? []).length > 0)
         .flatMap((textunit) =>
-          textunit.entity_ids.map((entityId) => ({
-            source: textunit.id,
-            target: nodes.find((e) => e.uuid === entityId)?.name || "",
-            type: "HAS_ENTITY",
-            id: `${textunit.id}-${entityId}`,
-          }))
-        );
+          textunit.entity_ids.map((entityId) => {
+            const targetName = nodes.find((e) => e.uuid === entityId)?.name;
+            return targetName ? {
+              source: textunit.id,
+              target: targetName,
+              type: "HAS_ENTITY",
+              id: `${textunit.id}-${entityId}`,
+            } : null;
+          })
+        )
+        .filter((link): link is NonNullable<typeof link> => link !== null);
 
       links.push(...textUnitEntityLinks);
     }
@@ -160,7 +184,8 @@ const useGraphData = (
 
             const newLinks = [];
 
-            if (!uniqueLinks.has(sourceLinkId)) {
+            // Only add link if source node exists
+            if (!uniqueLinks.has(sourceLinkId) && nodesMap[relationship.source]) {
               uniqueLinks.add(sourceLinkId);
               newLinks.push({
                 source: relationship.source,
@@ -170,7 +195,8 @@ const useGraphData = (
               });
             }
 
-            if (!uniqueLinks.has(targetLinkId)) {
+            // Only add link if target node exists
+            if (!uniqueLinks.has(targetLinkId) && nodesMap[relationship.target]) {
               uniqueLinks.add(targetLinkId);
               newLinks.push({
                 source: relationship.target,
@@ -242,12 +268,14 @@ const useGraphData = (
       covariateNodes.forEach(node => nodesMap[node.id] = node);
       nodes.push(...covariateNodes);
 
-      const covariateTextUnitLinks = covariates.map((covariate) => ({
-        source: covariate.text_unit_id,
-        target: covariate.id,
-        type: "HAS_COVARIATE",
-        id: `${covariate.text_unit_id}-${covariate.id}`,
-      }));
+      const covariateTextUnitLinks = covariates
+        .filter((covariate) => nodesMap[covariate.text_unit_id] && nodesMap[covariate.id])
+        .map((covariate) => ({
+          source: covariate.text_unit_id,
+          target: covariate.id,
+          type: "HAS_COVARIATE",
+          id: `${covariate.text_unit_id}-${covariate.id}`,
+        }));
 
       links.push(...covariateTextUnitLinks);
     }
@@ -283,6 +311,7 @@ const useGraphData = (
     includeTextUnits,
     includeCommunities,
     includeCovariates,
+    maxEntities,
   ]);
 
   return graphData;