backend/tools/graph_generator.py

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 "<div style='padding:20px;'>No hay documentos suficientes para generar el grafo.</div>"

        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()