Fixing weird behaviors

app.py

@@ -31,11 +31,14 @@ CUSTOM_CSS = """
 """
 
 # Cache directory and file paths
-CACHE_DIR = "cache"
+CACHE_DIR = "./cache"
+WORKING_DIR = "./sample"
 EXAMPLE_CACHE_FILE = os.path.join(CACHE_DIR, "first_example_cache.pkl")
+GRAPHML_FILE = WORKING_DIR + "/graph_chunk_entity_relation.graphml"
 
 # Create cache directory if it doesn't exist
 os.makedirs(CACHE_DIR, exist_ok=True)
+os.makedirs(WORKING_DIR, exist_ok=True)
 
 # Initialize the LLMGraph model
 model = LLMGraph()
@@ -62,7 +65,7 @@ def handle_text(text=""):
     
     return " ".join(text.split())
 
-def extract_kg(text="", model_name=None):
+def extract_kg(text="", model_name=MODEL_LIST[0]):
     """
     Extract knowledge graph from text
     """
@@ -73,7 +76,10 @@ def extract_kg(text="", model_name=None):
     try:
         result = model.extract(text, model_name)
 
-        return rapidjson.loads(result)
+        if isinstance(result, dict):
+            return result
+        else: # convert string to dict
+            return rapidjson.loads(result)
     except Exception as e:
         raise gr.Error(f"❌ Extraction error: {str(e)}")
 
@@ -108,7 +114,7 @@ def find_token_indices(doc, substring, text):
 
     return result
 
-def create_custom_entity_viz(data, full_text):
+def create_custom_entity_viz(data, full_text, type_col="type"):
     """
     Create custom entity visualization using spaCy's displacy
     """
@@ -130,7 +136,7 @@ def create_custom_entity_viz(data, full_text):
                 overlapping = any(s.start < end and start < s.end for s in spans)
 
                 if not overlapping:                
-                    node_type = node.get("type", "Entity")
+                    node_type = node.get(type_col, "Entity")
                     span = Span(doc, start, end, label=node_type)
                     spans.append(span)
 
@@ -156,30 +162,33 @@ def create_custom_entity_viz(data, full_text):
 
     return styled_html
 
-def create_graph(json_data):
+def create_graph(json_data, model_name=MODEL_LIST[0]):
     """
     Create interactive knowledge graph using pyvis
     """
 
-    G = nx.Graph()
+    if model_name == MODEL_LIST[0]:
+        G = nx.Graph()
 
-    # Add nodes with tooltips and error handling for missing keys
-    for node in json_data['nodes']:
-        # Get node type with fallback
-        type = node.get("type", "Entity")
+        # Add nodes with tooltips and error handling for missing keys
+        for node in json_data['nodes']:
+            # Get node type with fallback
+            type = node.get("type", "Entity")
 
-        # Get detailed type with fallback
-        detailed_type = node.get("detailed_type", type)
-        
-        # Use node ID and type info for the tooltip
-        G.add_node(node['id'], title=f"{type}: {detailed_type}")
-
-    # Add edges with labels
-    for edge in json_data['edges']:
-        # Check if the required keys exist
-        if 'from' in edge and 'to' in edge:
-            label = edge.get('label', 'related')
-            G.add_edge(edge['from'], edge['to'], title=label, label=label)
+            # Get detailed type with fallback
+            detailed_type = node.get("detailed_type", type)
+            
+            # Use node ID and type info for the tooltip
+            G.add_node(node['id'], title=f"{type}: {detailed_type}")
+
+        # Add edges with labels
+        for edge in json_data['edges']:
+            # Check if the required keys exist
+            if 'from' in edge and 'to' in edge:
+                label = edge.get('label', 'related')
+                G.add_edge(edge['from'], edge['to'], title=label, label=label)
+    else:
+        G = nx.read_graphml(GRAPHML_FILE)
 
     # Create network visualization
     network = Network(
@@ -193,17 +202,20 @@ def create_graph(json_data):
     
     # Configure network display
     network.from_nx(G)
-    # network.barnes_hut(
-    #     gravity=-3000,
-    #     central_gravity=0.3,
-    #     spring_length=50,
-    #     spring_strength=0.001,
-    #     damping=0.09,
-    #     overlap=0,
-    # )
+    network.barnes_hut(
+        gravity=-3000,
+        central_gravity=0.3,
+        spring_length=50,
+        spring_strength=0.001,
+        damping=0.09,
+        overlap=0,
+    )
 
     # Customize node appearance
     for node in network.nodes:
+        if "description" in node:
+            node["title"] = node["description"]
+
         node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}}
         node['font'] = {'size': 14, 'color': '#1e293b'}
         node['shape'] = 'dot'
@@ -211,6 +223,9 @@ def create_graph(json_data):
     
     # Customize edge appearance
     for edge in network.edges:
+        if "description" in edge:
+            edge["title"] = edge["description"]
+
         edge['width'] = 4
         # edge['arrows'] = {'to': {'enabled': False, 'type': 'arrow'}}
         edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
@@ -236,20 +251,20 @@ def process_and_visualize(text, model_name, progress=gr.Progress()):
     
     # Check if we're processing the first example for caching
     is_first_example = text == EXAMPLES[0][0]
-
-    asyncio.run(model.initialize_rag())  # Ensure RAG is initialized
+    
+    # Ensure RAG is initialized
+    asyncio.run(model.initialize_rag()) 
 
     # Try to load from cache if it's the first example
-    if is_first_example and os.path.exists(EXAMPLE_CACHE_FILE):
+    if is_first_example and model_name == MODEL_LIST[0] and os.path.exists(EXAMPLE_CACHE_FILE):
         try:
             progress(0.3, desc="Loading from cache...")
             with open(EXAMPLE_CACHE_FILE, 'rb') as f:
-                cache_data = pickle.load(f)
+                cached_data = pickle.load(f)
                 
             progress(1.0, desc="Loaded from cache!")
-            return cache_data["graph_html"], cache_data["entities_viz"], cache_data["json_data"], cache_data["stats"]
+            return cached_data["graph_html"], cached_data["entities_viz"], cached_data["json_data"], cached_data["stats"]
         except Exception as e:
-            # print(f"Cache loading error: {str(e)}")
             logging.error(f"Cache loading error: {str(e)}")
 
     # Continue with normal processing if cache fails
@@ -257,28 +272,30 @@ def process_and_visualize(text, model_name, progress=gr.Progress()):
     json_data = extract_kg(text, model_name)
     
     progress(0.5, desc="Creating entity visualization...")
-    entities_viz = create_custom_entity_viz(json_data, text)
+    if model_name == MODEL_LIST[0]:
+        entities_viz = create_custom_entity_viz(json_data, text, type_col="type")
+    else:
+        entities_viz = create_custom_entity_viz(json_data, text, type_col="entity_type")
     
     progress(0.8, desc="Building knowledge graph...")
-    graph_html = create_graph(json_data)
+    graph_html = create_graph(json_data, model_name)
     
     node_count = len(json_data["nodes"])
     edge_count = len(json_data["edges"])
     stats = f"📊 Extracted {node_count} entities and {edge_count} relationships"
     
     # Save to cache if it's the first example
-    if is_first_example:
+    if is_first_example and model_name == MODEL_LIST[0]:
         try:
-            cache_data = {
+            cached_data = {
                 "graph_html": graph_html,
                 "entities_viz": entities_viz,
                 "json_data": json_data,
                 "stats": stats
             }
             with open(EXAMPLE_CACHE_FILE, 'wb') as f:
-                pickle.dump(cache_data, f)
+                pickle.dump(cached_data, f)
         except Exception as e:
-            # print(f"Cache saving error: {str(e)}")
             logging.error(f"Cache saving error: {str(e)}")
     
     progress(1.0, desc="Complete!")
@@ -312,7 +329,6 @@ def generate_first_example():
     """
 
     if not os.path.exists(EXAMPLE_CACHE_FILE):
-        # print("Generating cache for first example...")
         logging.info("Generating cache for first example...")
 
         try:
@@ -338,15 +354,12 @@ def generate_first_example():
 
             with open(EXAMPLE_CACHE_FILE, 'wb') as f:
                 pickle.dump(cached_data, f)
-            # print("First example cache generated successfully")
             logging.info("First example cache generated successfully")
 
             return cached_data
         except Exception as e:
-            # print(f"Error generating first example cache: {str(e)}")
             logging.error(f"Error generating first example cache: {str(e)}")
     else:
-        # print("First example cache already exists")
         logging.info("First example cache already exists")
 
         # Load existing cache
@@ -354,7 +367,6 @@ def generate_first_example():
             with open(EXAMPLE_CACHE_FILE, 'rb') as f:
                 return pickle.load(f)
         except Exception as e:
-            # print(f"Error loading existing cache: {str(e)}")
             logging.error(f"Error loading existing cache: {str(e)}")
     
     return None

knowledge_graph.html

@@ -88,8 +88,8 @@
                   
 
                   // parsing and collecting nodes and edges from the python
-                  nodes = new vis.DataSet([{"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756577670, "description": "Aerosmith is a legendary rock band that has been active for 54 years and has officially announced their retirement from touring.", "entity_id": "Aerosmith", "entity_type": "organization", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Aerosmith", "label": "Aerosmith", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Aerosmith is a legendary rock band that has been active for 54 years and has officially announced their retirement from touring."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756577670, "description": "Steven Tyler is the lead singer of Aerosmith who suffered an unrecoverable vocal cord injury, leading to the band\u0027s retirement from touring.", "entity_id": "Steven Tyler", "entity_type": "person", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Steven Tyler", "label": "Steven Tyler", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler is the lead singer of Aerosmith who suffered an unrecoverable vocal cord injury, leading to the band\u0027s retirement from touring."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756577670, "description": "Vocal cord injury refers to the unrecoverable injury suffered by Steven Tyler that caused Aerosmith to retire from touring.", "entity_id": "Vocal Cord Injury", "entity_type": "category", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Vocal Cord Injury", "label": "Vocal Cord Injury", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Vocal cord injury refers to the unrecoverable injury suffered by Steven Tyler that caused Aerosmith to retire from touring."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756577670, "description": "Retirement from touring is the event announced by Aerosmith after 54 years, prompted by Steven Tyler\u0027s vocal cord injury.", "entity_id": "Retirement from Touring", "entity_type": "event", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Retirement from Touring", "label": "Retirement from Touring", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Retirement from touring is the event announced by Aerosmith after 54 years, prompted by Steven Tyler\u0027s vocal cord injury."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756577670, "description": "September 2023 is the time when Steven Tyler suffered a fractured larynx.", "entity_id": "September 2023", "entity_type": "event", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "September 2023", "label": "September 2023", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "September 2023 is the time when Steven Tyler suffered a fractured larynx."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756577670, "description": "Fractured larynx is the specific injury Steven Tyler suffered in September 2023, which was unsuccessfully treated.", "entity_id": "Fractured Larynx", "entity_type": "category", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Fractured Larynx", "label": "Fractured Larynx", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Fractured larynx is the specific injury Steven Tyler suffered in September 2023, which was unsuccessfully treated."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756577670, "description": "Unsuccessful treatment refers to the medical efforts to heal Steven Tyler\u0027s fractured larynx that did not result in recovery.", "entity_id": "Unsuccessful Treatment", "entity_type": "category", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Unsuccessful Treatment", "label": "Unsuccessful Treatment", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Unsuccessful treatment refers to the medical efforts to heal Steven Tyler\u0027s fractured larynx that did not result in recovery."}]);
-                  edges = new vis.DataSet([{"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577684, "description": "Steven Tyler is the lead singer of Aerosmith, whose vocal injury led to the band\u0027s retirement from touring.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Aerosmith", "keywords": "band membership,cause of retirement", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler is the lead singer of Aerosmith, whose vocal injury led to the band\u0027s retirement from touring.", "to": "Steven Tyler", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577685, "description": "Aerosmith\u0027s retirement from touring is due to Steven Tyler\u0027s unrecoverable vocal cord injury.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Aerosmith", "keywords": "cause of retirement,health impact", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Aerosmith\u0027s retirement from touring is due to Steven Tyler\u0027s unrecoverable vocal cord injury.", "to": "Vocal Cord Injury", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577687, "description": "Aerosmith officially announced their retirement from touring after 54 years.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Aerosmith", "keywords": "band decision,career milestone", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Aerosmith officially announced their retirement from touring after 54 years.", "to": "Retirement from Touring", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577685, "description": "Steven Tyler suffered a fractured larynx in September 2023.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Steven Tyler", "keywords": "injury timing,medical event", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler suffered a fractured larynx in September 2023.", "to": "September 2023", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577687, "description": "The vocal cord injury is the medical condition affecting Steven Tyler that caused Aerosmith\u0027s retirement.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Steven Tyler", "keywords": "cause-effect,medical condition", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "The vocal cord injury is the medical condition affecting Steven Tyler that caused Aerosmith\u0027s retirement.", "to": "Vocal Cord Injury", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577690, "description": "The fractured larynx is the specific injury Steven Tyler suffered, leading to unsuccessful treatment and vocal cord damage.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Steven Tyler", "keywords": "injury detail,medical diagnosis", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "The fractured larynx is the specific injury Steven Tyler suffered, leading to unsuccessful treatment and vocal cord damage.", "to": "Fractured Larynx", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577691, "description": "Steven Tyler underwent months of unsuccessful treatment for his fractured larynx.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Steven Tyler", "keywords": "health outcome,medical treatment", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler underwent months of unsuccessful treatment for his fractured larynx.", "to": "Unsuccessful Treatment", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577691, "description": "The vocal cord injury is a result of the fractured larynx suffered by Steven Tyler.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Vocal Cord Injury", "keywords": "injury relationship,medical causation", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "The vocal cord injury is a result of the fractured larynx suffered by Steven Tyler.", "to": "Fractured Larynx", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756577693, "description": "The unsuccessful treatment was aimed at healing the fractured larynx suffered by Steven Tyler.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Fractured Larynx", "keywords": "injury focus,medical intervention", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "The unsuccessful treatment was aimed at healing the fractured larynx suffered by Steven Tyler.", "to": "Unsuccessful Treatment", "width": 4}]);
+                  nodes = new vis.DataSet([{"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756651432, "description": "Aerosmith is a legendary rock band that has announced their retirement from touring after 54 years.", "entity_id": "Aerosmith", "entity_type": "organization", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Aerosmith", "label": "Aerosmith", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Aerosmith is a legendary rock band that has announced their retirement from touring after 54 years."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756651432, "description": "Steven Tyler is the lead singer of Aerosmith who suffered an unrecoverable vocal cord injury leading to the band\u0027s retirement from touring.", "entity_id": "Steven Tyler", "entity_type": "person", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Steven Tyler", "label": "Steven Tyler", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler is the lead singer of Aerosmith who suffered an unrecoverable vocal cord injury leading to the band\u0027s retirement from touring."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756651432, "description": "Vocal cord injury refers to the unrecoverable injury suffered by Steven Tyler that caused Aerosmith to retire from touring.", "entity_id": "Vocal Cord Injury", "entity_type": "category", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Vocal Cord Injury", "label": "Vocal Cord Injury", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Vocal cord injury refers to the unrecoverable injury suffered by Steven Tyler that caused Aerosmith to retire from touring."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756651432, "description": "Touring refers to the activity of performing live concerts in various locations, which Aerosmith has retired from after 54 years.", "entity_id": "Touring", "entity_type": "category", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Touring", "label": "Touring", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Touring refers to the activity of performing live concerts in various locations, which Aerosmith has retired from after 54 years."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756651432, "description": "Fractured larynx is the specific injury Steven Tyler suffered in September 2023, which was unsuccessfully treated.", "entity_id": "Fractured Larynx", "entity_type": "category", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Fractured Larynx", "label": "Fractured Larynx", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Fractured larynx is the specific injury Steven Tyler suffered in September 2023, which was unsuccessfully treated."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756651432, "description": "Unsuccessful treatment refers to the medical efforts to heal Steven Tyler\u0027s fractured larynx that did not result in recovery.", "entity_id": "Unsuccessful Treatment", "entity_type": "category", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "Unsuccessful Treatment", "label": "Unsuccessful Treatment", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Unsuccessful treatment refers to the medical efforts to heal Steven Tyler\u0027s fractured larynx that did not result in recovery."}, {"color": {"background": "#e0e7ff", "border": "#6366f1", "highlight": {"background": "#c7d2fe", "border": "#4f46e5"}}, "created_at": 1756651432, "description": "September 2023 is the time when Steven Tyler suffered a fractured larynx.", "entity_id": "September 2023", "entity_type": "event", "file_path": "unknown_source", "font": {"color": "#1e293b", "size": 14}, "id": "September 2023", "label": "September 2023", "shape": "dot", "size": 20, "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "September 2023 is the time when Steven Tyler suffered a fractured larynx."}]);
+                  edges = new vis.DataSet([{"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756651440, "description": "Steven Tyler is the lead singer of Aerosmith, whose vocal injury led to the band\u0027s retirement from touring.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Aerosmith", "keywords": "band membership,cause of retirement", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler is the lead singer of Aerosmith, whose vocal injury led to the band\u0027s retirement from touring.", "to": "Steven Tyler", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756651442, "description": "The vocal cord injury to Steven Tyler led to Aerosmith\u0027s retirement from touring.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Aerosmith", "keywords": "band decision,cause and effect", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "The vocal cord injury to Steven Tyler led to Aerosmith\u0027s retirement from touring.", "to": "Vocal Cord Injury", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756651444, "description": "Aerosmith has officially retired from touring after 54 years.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Aerosmith", "keywords": "career activity,retirement", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Aerosmith has officially retired from touring after 54 years.", "to": "Touring", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756651441, "description": "Steven Tyler suffered an unrecoverable vocal cord injury that affected his ability to perform.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Steven Tyler", "keywords": "career impact,health issue", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler suffered an unrecoverable vocal cord injury that affected his ability to perform.", "to": "Vocal Cord Injury", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756651442, "description": "Steven Tyler\u0027s fractured larynx in September 2023 was the cause of his vocal cord injury.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Steven Tyler", "keywords": "health event,injury cause", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler\u0027s fractured larynx in September 2023 was the cause of his vocal cord injury.", "to": "Fractured Larynx", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756651444, "description": "Steven Tyler underwent unsuccessful treatment for his fractured larynx.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Steven Tyler", "keywords": "health outcome,medical treatment", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "Steven Tyler underwent unsuccessful treatment for his fractured larynx.", "to": "Unsuccessful Treatment", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756651444, "description": "The fractured larynx injury occurred in September 2023.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "September 2023", "keywords": "injury timing,temporal occurrence", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "The fractured larynx injury occurred in September 2023.", "to": "Fractured Larynx", "width": 4}, {"color": {"color": "#6366f1", "highlight": "#4f46e5"}, "created_at": 1756651447, "description": "The unsuccessful treatment was aimed at healing the fractured larynx.", "file_path": "unknown_source", "font": {"color": "#4b5563", "face": "Arial", "size": 12}, "from": "Fractured Larynx", "keywords": "injury management,medical intervention", "source_id": "chunk-150cfba3862e116efcee671d872955be", "title": "The unsuccessful treatment was aimed at healing the fractured larynx.", "to": "Unsuccessful Treatment", "width": 4}]);
 
                   nodeColors = {};
                   allNodes = nodes.get({ returnType: "Object" });

llm_graph.py

@@ -1,6 +1,8 @@
 import os
-import asyncio
+import time
+import shutil
 import numpy as np
+import networkx as nx
 
 from textwrap import dedent
 from dotenv import load_dotenv
@@ -25,7 +27,8 @@ AZURE_OPENAI_ENDPOINT = os.environ["AZURE_OPENAI_ENDPOINT"]
 AZURE_EMBEDDING_DEPLOYMENT = os.environ["AZURE_EMBEDDING_DEPLOYMENT"]
 AZURE_EMBEDDING_API_VERSION = os.environ["AZURE_EMBEDDING_API_VERSION"]
 
-WORKING_DIR = "./cache"
+WORKING_DIR = "./sample"
+GRAPHML_FILE = WORKING_DIR + "/graph_chunk_entity_relation.graphml"
 
 MODEL_LIST = [
   "EmergentMethods/Phi-3-mini-128k-instruct-graph",
@@ -52,9 +55,9 @@ class LLMGraph:
                     func=self._embedding_func,
                 ),
             )
-
-            await self.rag.initialize_storages()
-            await initialize_pipeline_status()
+        # TODO: Check if this works as expected
+        await self.rag.initialize_storages()
+        await initialize_pipeline_status()
 
     # async def test_responses(self):
     #     """
@@ -151,24 +154,37 @@ class LLMGraph:
 
         return messages
     
-    def extract(self, text, model_name=MODEL_LIST[0]) -> str:
+    def extract(self, text, model_name=MODEL_LIST[0]):
         """
-        Extract knowledge graph from text
+        Extract knowledge graph in structured format from text.
         """
-
-        generated_text = "This is a placeholder response."
         
         if model_name == MODEL_LIST[0]:
             # Use Hugging Face Inference API with Phi-3-mini-128k-instruct-graph
             messages = self._get_messages(text)
-            generated_text = self._generate(messages)
+
+            json_graph = self._generate(messages)
+            return json_graph
         else:
+            if os.path.exists(WORKING_DIR):
+                shutil.rmtree(WORKING_DIR)
+            os.makedirs(WORKING_DIR, exist_ok=True)
+            
             # Use LightRAG with Azure OpenAI
+            # TODO: Clear all the previous inserted texts first
             self.rag.insert(text) # Insert the text into the RAG storage
-            # TODO: Extract JSON format of the knowledge graph
+
+            # Wait for GRAPHML_FILE to be created
+            while not os.path.exists(GRAPHML_FILE):
+                time.sleep(0.1) # Sleep for 100ms before checking again
             
-        return generated_text
+            # Extract dict format of the knowledge graph
+            G = nx.read_graphml(GRAPHML_FILE)
 
+            # Convert the graph to node-link data format
+            dict_graph = nx.node_link_data(G, edges="edges")
+            return dict_graph
+        
     async def _llm_model_func(self, prompt, system_prompt=None, history_messages=[], **kwargs) -> str:
         """
         Call the Azure OpenAI chat completion endpoint with the given prompt and optional system prompt and history messages.

sample/kv_store_doc_status.json

@@ -0,0 +1,35 @@
+{
+  "doc-605403c35618c5288c57c562f8eca566": {
+    "status": "processed",
+    "chunks_count": 1,
+    "chunks_list": [
+      "chunk-605403c35618c5288c57c562f8eca566"
+    ],
+    "content_summary": "The family of Azerbaijan President Ilham Aliyev leads a charmed, glamorous life, thanks in part to financial interests in almost every sector of the economy. His wife, Mehriban, comes from the privileged and powerful Pashayev family that owns banks, ...",
+    "content_length": 1074,
+    "created_at": "2025-08-31T15:50:59.506391+00:00",
+    "updated_at": "2025-08-31T15:52:26.018288+00:00",
+    "file_path": "unknown_source",
+    "track_id": "insert_20250831_235059_6946ff78",
+    "metadata": {
+      "processing_start_time": 1756655459,
+      "processing_end_time": 1756655546
+    }
+  },
+  "doc-eea199eb7feea197ebb82e9333a2d2f2": {
+    "status": "processing",
+    "chunks_count": 1,
+    "chunks_list": [
+      "chunk-eea199eb7feea197ebb82e9333a2d2f2"
+    ],
+    "content_summary": "Les jardins du Luxembourg, situés au cœur du sixième arrondissement de Paris, offrent un véritable havre de paix aux citadins pressés. Créés au début du dix-septième siècle sur l'initiative de Marie de Médicis, ces jardins à la française s'étendent s...",
+    "content_length": 697,
+    "created_at": "2025-08-31T15:54:38.060638+00:00",
+    "updated_at": "2025-08-31T15:54:38.068349+00:00",
+    "file_path": "unknown_source",
+    "track_id": "insert_20250831_235438_22d326d7",
+    "metadata": {
+      "processing_start_time": 1756655678
+    }
+  }
+}

visualize.py

@@ -1,5 +1,4 @@
 import networkx as nx
-import rapidjson
 import warnings
 import os
 
@@ -13,98 +12,41 @@ file_path = "./cache/graph_chunk_entity_relation.graphml"
 assert os.path.exists(file_path), f"File {file_path} does not exist."
 G = nx.read_graphml(file_path)
 
-def create_graph(json_data):
-    """
-    Create interactive knowledge graph using pyvis
-    """
-
-    G = nx.Graph()
-
-    # Add nodes with tooltips and error handling for missing keys
-    for node in json_data['nodes']:
-        # Get node type with fallback
-        type = node.get("type", "Entity")
-
-        # Get detailed type with fallback
-        detailed_type = node.get("detailed_type", type)
-        
-        # Use node ID and type info for the tooltip
-        G.add_node(node['id'], title=f"{type}: {detailed_type}")
-
-    # Add edges with labels
-    for edge in json_data['edges']:
-        # Check if the required keys exist
-        if 'from' in edge and 'to' in edge:
-            label = edge.get('label', 'related')
-            G.add_edge(edge['from'], edge['to'], title=label, label=label)
-
-    # Create network visualization
-    network = Network(
-        width="100%",
-        height="100vh",
-        notebook=False,
-        bgcolor="#f8fafc", 
-        font_color="#1e293b"
-    )
-    
-    # Configure network display
-    network.from_nx(G)
-    
-    # Customize node appearance
-    for node in network.nodes:
-        node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}}
-        node['font'] = {'size': 14, 'color': '#1e293b'}
-        node['shape'] = 'dot'
-        node['size'] = 20
-    
-    # Customize edge appearance
-    for edge in network.edges:
-        edge['width'] = 4
-        edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
-        edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'}
-    
-    # Save and display the network
-    filename_out = "knowledge_graph.html"
-    network.show(filename_out)
-    print(f"Knowledge graph saved to {filename_out}")
-
 # Convert the graph to node-link data format
-js_graph = nx.node_link_data(G)
-js_data = rapidjson.loads(rapidjson.dumps(js_graph))
-# print(js_data)
-
-create_graph(js_data)
-
-# # Create a Pyvis network
-# network = Network(width="100%",
-#                   height="100vh",
-#                   notebook=True,
-#                   bgcolor="#f8fafc",
-#                   font_color="#1e293b")
-
-# # Convert NetworkX graph to Pyvis network
-# network.from_nx(G)
-
-# # Add colors and title to nodes
-# for node in network.nodes:
-#     if "description" in node:
-#         node["title"] = node["description"]
-
-#     node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}}
-#     node['font'] = {'size': 14, 'color': '#1e293b'}
-#     node['shape'] = 'dot'
-#     node['size'] = 20
-
-# # Add title to edges
-# for edge in network.edges:
-#     if "description" in edge:
-#         edge["title"] = edge["description"]
-
-#     edge['width'] = 4
-#     edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
-#     edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'}
-
-# # Save and display the network
-# filename_out = "knowledge_graph.html"
-# network.show(filename_out)
-# print(f"Knowledge graph saved to {filename_out}")
+dict_graph = nx.node_link_data(G)
+print("Number of nodes:", len(dict_graph['nodes']))
+print("Number of edges:", len(dict_graph['links']))
+
+# Create a Pyvis network
+network = Network(width="100%",
+                  height="100vh",
+                  notebook=True,
+                  bgcolor="#f8fafc",
+                  font_color="#1e293b")
+
+# Convert NetworkX graph to Pyvis network
+network.from_nx(G)
+
+# Add colors and title to nodes
+for node in network.nodes:
+    if "description" in node:
+        node["title"] = node["description"]
+
+    node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}}
+    node['font'] = {'size': 14, 'color': '#1e293b'}
+    node['shape'] = 'dot'
+    node['size'] = 20
+
+# Add title to edges
+for edge in network.edges:
+    if "description" in edge:
+        edge["title"] = edge["description"]
+
+    edge['width'] = 4
+    edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
+    edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'}
+
+# Save and display the network
+filename_out = "knowledge_graph.html"
+network.show(filename_out)
+print(f"Knowledge graph saved to {filename_out}")