Move diagram logic to diagram component

jtmuller5 · Nov 24, 2024 · 91656ec · 91656ec
1 parent 7f825d7
commit 91656ec
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 105 deletions.
diff --git a/dist/sidepanel.js b/dist/sidepanel.js
diff --git a/src/features/side_panel/components/CombinedTab.tsx b/src/features/side_panel/components/CombinedTab.tsx
@@ -23,7 +23,6 @@ import { tab } from "@testing-library/user-event/dist/tab";
 interface SectionData {
   summary: string;
   relationships: Relationship[];
-  diagram: string;
   message: string;
 }
 
@@ -47,7 +46,6 @@ export default function CombinedTab({
   const [showDiagrams, setShowDiagrams] = useState<boolean>(true);
 
   // Combined Diagram and Entities
-  const [combinedDiagram, setCombinedDiagram] = useState<string | null>(null);
   const [combinedRelationships, setCombinedRelationships] = useState<
     Relationship[]
   >([]);
@@ -75,31 +73,12 @@ export default function CombinedTab({
     };
   }, []);
 
-  const renderMermaidDiagram = (relationships: Relationship[]): string => {
-    let mermaidCode = "graph LR\n";
-    relationships.forEach((rel) => {
-      const { entity1, entity2, description } = rel;
-      const safeEntity1 = sanitizeMermaidText(entity1);
-      const safeEntity2 = sanitizeMermaidText(entity2);
-      const safeDescription = sanitizeMermaidLabel(description);
-      mermaidCode += `    ${safeEntity1}[${sanitizeMermaidLabel(
-        entity1
-      )}] -->|${safeDescription}| ${safeEntity2}[${sanitizeMermaidLabel(
-        entity2
-      )}]\n`;
-      mermaidCode += `    style ${safeEntity1} fill:#0077be,color:#fff\n`;
-      mermaidCode += `    style ${safeEntity2} fill:#0077be,color:#fff\n`;
-    });
-    return mermaidCode;
-  };
-
   const analyzeSummaryAndContent = useCallback(async () => {
     try {
       relationshipService.reset();
       setSections([]);
       setError(null);
       setStatus("Analyzing page content...");
-      setCombinedDiagram(null);
       setCombinedRelationships([]);
       setCombinedEntities([]);
       setSelectedEntity(null);
@@ -141,7 +120,6 @@ export default function CombinedTab({
             const sectionData: SectionData = {
               summary: chunkSummary,
               relationships: sectionRelationships,
-              diagram: renderMermaidDiagram(sectionRelationships),
               message: chunkResult,
             };
 
@@ -185,28 +163,18 @@ export default function CombinedTab({
   const handleEntityClick = (entity: string) => {
     if (entity === selectedEntity) {
       setSelectedEntity(null);
-      // Reset diagram to show all relationships
-      const combinedDiagramDefinition = renderMermaidDiagram(
-        combinedRelationships
-      );
-      setCombinedDiagram(combinedDiagramDefinition);
     } else {
       setSelectedEntity(entity);
       // Filter relationships to those involving the selected entity
       const filteredRelationships = combinedRelationships.filter(
         (rel) => rel.entity1 === entity || rel.entity2 === entity
       );
-      const filteredDiagram = renderMermaidDiagram(filteredRelationships);
-      setCombinedDiagram(filteredDiagram);
+      setCombinedRelationships(filteredRelationships);
     }
   };
 
   const handleShowAllRelationships = () => {
     setSelectedEntity(null);
-    const combinedDiagramDefinition = renderMermaidDiagram(
-      combinedRelationships
-    );
-    setCombinedDiagram(combinedDiagramDefinition);
   };
 
   const deleteEntity = (entity: string) => {
@@ -220,11 +188,6 @@ export default function CombinedTab({
     const updatedEntities = combinedEntities.filter((e) => e.name !== entity);
     setCombinedEntities(updatedEntities);
 
-    // Update diagram
-    const combinedDiagramDefinition =
-      renderMermaidDiagram(updatedRelationships);
-    setCombinedDiagram(combinedDiagramDefinition);
-
     // If the deleted entity was selected, reset selection
     if (selectedEntity === entity) {
       setSelectedEntity(null);
@@ -242,8 +205,6 @@ export default function CombinedTab({
       (section) => section.relationships
     );
     // Generate mermaid diagram
-    const combinedDiagramDefinition = renderMermaidDiagram(allRelationships);
-    setCombinedDiagram(combinedDiagramDefinition);
     setCombinedRelationships(allRelationships);
 
     // Collect entities and their counts
@@ -323,7 +284,7 @@ export default function CombinedTab({
         <div className="space-y-6 pt-6">
           <SystemStatus status={status} error={error} />
 
-          {combinedDiagram && (
+          {combinedRelationships && combinedRelationships.length > 0 && (
             <Card className="mb-6">
               <CardContent>
                 <Tabs
@@ -347,7 +308,7 @@ export default function CombinedTab({
                   </div>
                   <TabsContent value="diagram">
                     <DiagramComponent
-                      diagramDefinition={combinedDiagram}
+                      relationships={combinedRelationships}
                       onNodeClick={handleDiagramNodeClick}
                     />
                   </TabsContent>
@@ -386,7 +347,7 @@ export default function CombinedTab({
 
                 {showDiagrams && (
                   <DiagramComponent
-                    diagramDefinition={section.diagram}
+                    relationships={section.relationships}
                     onNodeClick={handleNodeClick}
                   />
                 )}

diff --git a/src/features/side_panel/components/DiagramComponent.tsx b/src/features/side_panel/components/DiagramComponent.tsx
@@ -1,21 +1,47 @@
-import React, { useEffect, useRef } from "react";
+import React, { useEffect, useRef, useState } from "react";
 import mermaid from "mermaid";
+import { Relationship } from "../types/relationship";
+import {
+  sanitizeMermaidText,
+  sanitizeMermaidLabel,
+} from "../utils/diagram_utils";
 
 interface DiagramComponentProps {
-  diagramDefinition: string;
+  relationships: Relationship[];
   onNodeClick?: (entity: string) => void;
 }
 
 const DiagramComponent: React.FC<DiagramComponentProps> = ({
-  diagramDefinition,
+  relationships,
   onNodeClick,
 }) => {
+  const [diagramDefinition, setDiagramDefinition] = useState<string>("");
+
   const diagramRef = useRef<HTMLDivElement>(null);
 
+  const renderMermaidDiagram = (relationships: Relationship[]): string => {
+    let mermaidCode = "graph LR\n";
+    relationships.forEach((rel) => {
+      const { entity1, entity2, description } = rel;
+      const safeEntity1 = sanitizeMermaidText(entity1);
+      const safeEntity2 = sanitizeMermaidText(entity2);
+      const safeDescription = sanitizeMermaidLabel(description);
+      mermaidCode += `    ${safeEntity1}[${sanitizeMermaidLabel(
+        entity1
+      )}] -->|${safeDescription}| ${safeEntity2}[${sanitizeMermaidLabel(
+        entity2
+      )}]\n`;
+      mermaidCode += `    style ${safeEntity1} fill:#0077be,color:#fff\n`;
+      mermaidCode += `    style ${safeEntity2} fill:#0077be,color:#fff\n`;
+    });
+    return mermaidCode;
+  };
+
   useEffect(() => {
     // Initialize Mermaid only once
     mermaid.initialize({ startOnLoad: false });
-  }, []);
+    setDiagramDefinition(renderMermaidDiagram(relationships));
+  }, [relationships]);
 
   useEffect(() => {
     const renderMermaidDiagram = async () => {
@@ -78,6 +104,10 @@ const DiagramComponent: React.FC<DiagramComponentProps> = ({
     renderMermaidDiagram();
   }, [diagramDefinition, onNodeClick]);
 
+  if (!diagramDefinition) {
+    return null;
+  }
+
   return (
     <div className="p-2 transition-opacity duration-300" ref={diagramRef}></div>
   );