Update visualization.py

visualization.py

@@ -5,17 +5,19 @@ from pyvis.network import Network
 
 def extract_key_terms(text: str):
     """
-    A naive approach to extract key terms by matching capitalized words.
+    A simple method to extract key terms by matching capitalized words.
+    In a clinical context, this may highlight important entities.
     """
     return re.findall(r"\b[A-Z][a-zA-Z]+\b", text)
 
 def create_medical_graph(query: str, docs: list) -> str:
     """
-    Builds a PyVis network:
-      - A central "QUERY" node.
-      - A node for each retrieved document.
-      - Sub-nodes for extracted key terms.
-    Returns the full HTML of the generated graph.
+    Creates an interactive knowledge graph using PyVis:
+      - A central node for the query.
+      - One node per abstract.
+      - Sub-nodes for key terms extracted from each abstract.
+    
+    Returns the HTML content of the graph.
     """
     net = Network(height="600px", width="100%", directed=False)
     net.add_node("QUERY", label=f"Query: {query}", color="red", shape="star")
@@ -31,7 +33,7 @@ def create_medical_graph(query: str, docs: list) -> str:
             net.add_node(term_id, label=term, color="green")
             net.add_edge(doc_id, term_id)
 
-    # Write the network HTML to a temporary file and return its content
+    # Write the graph to a temporary HTML file and return its content.
     with tempfile.NamedTemporaryFile(delete=False, suffix=".html") as tmp:
         temp_filename = tmp.name
     net.show(temp_filename)