Update app.py

app.py

@@ -19,9 +19,7 @@ from llama_index.llms.openai import OpenAI
 from llama_index.embeddings.openai import OpenAIEmbedding
 from llama_index.core import Settings
 
-import openai
 import os
-from github import Github
 from datetime import datetime
 import gradio as gr
 import pandas as pd
@@ -34,38 +32,49 @@ exec(os.environ.get('context'))
 
 
 ##### Graph start ##########
-import networkx as nx
-import matplotlib.pyplot as plt
-from PIL import Image
-from io import BytesIO
+from pyvis.network import Network
+from io import StringIO
+import html
 
 def draw_graph():
     global kg_data
-    G = nx.DiGraph()
+    
+    # Crear un grafo con opciones visuales básicas
+    net = Network(
+        notebook=True,
+        height="700px",
+        width="100%",
+        select_menu=True,
+        neighborhood_highlight=True,
+        cdn_resources="remote"
+    )
+
+    # Añadir nodos y relaciones desde los datos del grafo
     for source, relation, target in kg_data:
-        G.add_edge(source, target, label=relation)
-
-    # Utilizar spring_layout para mejorar la disposición de los nodos
-    pos = nx.spring_layout(G)
-
-    plt.figure(figsize=(12, 8))
-    # Ajustar el tamaño de los nodos
-    nx.draw(G, pos, with_labels=True, node_color='skyblue', node_size=400, edge_color='k', linewidths=1, font_size=8, font_weight='bold')
-    # Ajustar el tamaño de las flechas y el espaciado entre ellas
-    edge_labels = {}
-    for source, target, data in G.edges(data=True):
-        if 'label' in data:
-            edge_labels[(source, target)] = data['label']
-    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=7, font_weight='normal')
-    plt.title("Graph")
-    plt.axis('off')
-
-    buf = BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
-    plt.close()
-
-    return Image.open(buf)
+        # Agregar nodos
+        net.add_node(source, label=source, color="skyblue")
+        net.add_node(target, label=target, color="skyblue")
+        # Agregar aristas con la relación como etiqueta
+        net.add_edge(source, target, title=relation, label=relation, color="black")
+
+    # Generar el HTML como string
+    html_output = StringIO()
+    net.write_html(html_output)
+    html_output.seek(0)
+    html_contenido = html_output.getvalue()
+
+    # Generar el iframe embebido
+    return generar_iframe_embebido(html_contenido)
+
+def generar_iframe_embebido(html_contenido):
+    # Escapar el contenido HTML para que sea seguro dentro del atributo srcdoc
+    html_escapado = html.escape(html_contenido)
+    
+    # Generar el string del iframe con el contenido embebido
+    iframe = f"""
+    <iframe srcdoc="{html_escapado}" width="100%" height="700px" style="border:none;"></iframe>
+    """
+    return iframe    
 ##### Graph end ##########
 
 
@@ -716,7 +725,7 @@ with gr.Blocks(theme='sudeepshouche/minimalist', css=css) as demo:
 """)
     
     with gr.Row():
-        grafo = gr.Image(label="Grafo", show_share_button=False)
+        grafo = gr.HTML(label="Grafo", show_share_button=False)
 
     with gr.Accordion(elem_classes="accordion", label="Audit trail", open=False):
         with gr.Row():