import networkx as nx
from pyvis.network import Network
import os
import uuid
# Directorio para guardar grafos temporales
GRAPH_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "assets", "graphs")
os.makedirs(GRAPH_DIR, exist_ok=True)
class GraphGenerator:
def __init__(self):
pass
def generate_research_graph(self, docs: list) -> str:
"""
Toma una lista de documentos y genera un grafo interactivo de relaciones.
Devuelve el código HTML del grafo (o la ruta al archivo generado).
"""
if not docs:
return "
No hay documentos suficientes para generar el grafo.
"
G = nx.Graph()
# Añadir un nodo central para la consulta (opcional, en este caso los agruparemos por autores y fuentes)
# Recorrer documentos
for doc in docs:
title = doc.get("title", "Desconocido")[:30] + "..."
doc_id = doc.get("doi") or doc.get("pdfUrl") or title
source = doc.get("source", "Fuente Desconocida")
# Añadir nodo del documento principal
G.add_node(doc_id, label=title, title=doc.get("title", ""), color="#8b5cf6", shape="dot", size=20)
# Nodo de la fuente
G.add_node(source, label=source, color="#10b981", shape="square", size=25)
G.add_edge(doc_id, source)
# Nodos de autores
authors = doc.get("authors", [])
if isinstance(authors, list):
for author in authors[:3]: # Solo los primeros 3 autores para no saturar
author_name = str(author).strip()
if author_name:
G.add_node(author_name, label=author_name, color="#f59e0b", shape="triangle", size=15)
G.add_edge(doc_id, author_name)
# Generar con pyvis
net = Network(height="600px", width="100%", bgcolor="#0f172a", font_color="white", select_menu=True)
# Opciones físicas para un buen layout
net.force_atlas_2based(gravity=-50, central_gravity=0.01, spring_length=100, spring_strength=0.08, damping=0.4, overlap=0)
net.from_nx(G)
filename = f"graph_{uuid.uuid4().hex[:8]}.html"
filepath = os.path.join(GRAPH_DIR, filename)
net.write_html(filepath)
# Leer el contenido HTML generado
with open(filepath, "r", encoding="utf-8") as f:
html_content = f.read()
return html_content
generator = GraphGenerator()