utils/graph_builder.py

"""
Graph Builder cho Medical RAG Pipeline.

Pipeline:
    1. Load MedQuad dataset từ HuggingFace
    2. Parse documents qua LlamaIndex PropertyGraphIndex
    3. Chuyển đổi sang PyTorch Geometric graph
       - Nodes: mỗi text chunk = một node
       - Edges: cosine similarity > threshold (document similarity graph)
    4. Trả về pyg_data, nodes, và raw embeddings

Lý do dùng document similarity graph:
    Hai documents có nội dung liên quan (e.g. cùng nói về diabetes)
    sẽ được kết nối. GNN sẽ propagate thông tin qua các cạnh này,
    enriching embeddings với neighborhood context.
"""

import torch
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.utils import add_self_loops, coalesce

from llama_index.core import PropertyGraphIndex, Document, Settings
from llama_index.core.indices.property_graph import ImplicitPathExtractor
from datasets import load_dataset

from typing import List, Tuple


def load_healthcare_documents(num_samples: int = 250) -> List[Document]:
    """
    Load dataset y tế từ HuggingFace.

    Dataset: keivalya/MedQuad-MedicalQnADataset
    Mỗi Document = 1 câu trả lời y tế, với câu hỏi lưu trong metadata.

    Args:
        num_samples: số lượng documents cần load (giới hạn để tiết kiệm RAM)

    Returns:
        List[Document] - danh sách LlamaIndex Document objects
    """
    print(f"Đang load {num_samples} medical documents từ HuggingFace...")

    dataset = load_dataset(
        "keivalya/MedQuad-MedicalQnADataset",
        split="train"
    )
    num_samples = min(num_samples, len(dataset))
    dataset = dataset.select(range(num_samples))

    docs = []
    for i, row in enumerate(dataset):
        # Truncate text dài để tiết kiệm RAM khi embedding
        answer_text = str(row.get("Answer", "")).strip()[:800]
        question_text = str(row.get("Question", "")).strip()[:200]

        if not answer_text:
            continue

        doc = Document(
            text=answer_text,
            metadata={
                "question": question_text,
                "doc_id": str(i),
                "source": "MedQuad"
            }
        )
        docs.append(doc)

    print(f"✓ Loaded {len(docs)} medical documents")
    return docs


def build_llamaindex_property_graph(
    documents: List[Document],
    embed_model,
) -> Tuple[PropertyGraphIndex, List]:
    """
    Xây dựng LlamaIndex PropertyGraphIndex từ documents.

    Sử dụng ImplicitPathExtractor - không cần LLM,
    phù hợp để deploy trên free tier.

    PropertyGraphIndex lưu trữ:
        - ChunkNodes: các đoạn văn bản từ documents
        - Implicit edges: dựa trên cấu trúc document

    Args:
        documents: danh sách LlamaIndex Document objects
        embed_model: HuggingFace embedding model

    Returns:
        (index, text_nodes) - PropertyGraphIndex và list TextNode objects
    """
    print("Đang xây dựng LlamaIndex PropertyGraphIndex...")

    # Cấu hình Settings - không dùng LLM của LlamaIndex
    # vì sẽ dùng llama-cpp trực tiếp cho generation
    Settings.embed_model = embed_model
    Settings.llm = None

    index = PropertyGraphIndex.from_documents(
        documents,
        embed_model=embed_model,
        kg_extractors=[ImplicitPathExtractor()],  # no LLM required
        show_progress=True,
    )

    # Lấy tất cả text nodes từ docstore
    # index.docstore.docs = {id: TextNode} dict
    text_nodes = list(index.docstore.docs.values())

    print(f"✓ PropertyGraphIndex built: {len(text_nodes)} text nodes")
    return index, text_nodes


def build_pyg_graph(
    text_nodes: List,
    embed_model,
    similarity_threshold: float = 0.70,
) -> Tuple[Data, List, torch.Tensor]:
    """
    Chuyển LlamaIndex text nodes sang PyTorch Geometric graph.

    Strategy xây dựng graph:
        - Mỗi text chunk = 1 node trong graph
        - Edge giữa 2 nodes nếu cosine_similarity > threshold
        - Thêm self-loops cho GCN stability (A_hat = A + I)

    Rationale: Documents về cùng chủ đề y tế (e.g. diabetes, hypertension)
    sẽ có high cosine similarity và được kết nối.
    GNN sẽ propagate thông tin qua các edges này, giúp mỗi node
    "biết" về context của các documents liên quan.

    Args:
        text_nodes: danh sách LlamaIndex TextNode objects
        embed_model: HuggingFace embedding model
        similarity_threshold: ngưỡng cosine similarity để tạo edge

    Returns:
        (pyg_data, nodes, raw_embeddings)
            pyg_data: PyTorch Geometric Data object
            nodes: ordered list of TextNode objects
            raw_embeddings: [N, D] tensor - raw text embeddings
    """
    N = len(text_nodes)
    print(f"Đang embed {N} nodes và xây dựng PyG graph...")

    # ── Step 1: Embed all nodes ──────────────────────────────────────────────
    texts = [node.get_content() for node in text_nodes]
    embeddings = embed_model.get_text_embedding_batch(texts, show_progress=True)
    x = torch.tensor(embeddings, dtype=torch.float)  # [N, D]

    # ── Step 2: Tính cosine similarity matrix ────────────────────────────────
    x_norm = F.normalize(x, dim=-1)            # [N, D] - L2 normalized
    sim_matrix = x_norm @ x_norm.T             # [N, N] - cosine similarity

    # ── Step 3: Tạo edges từ similarity ──────────────────────────────────────
    # Loại bỏ diagonal (self-similarity = 1.0 không có ý nghĩa cấu trúc)
    mask = (sim_matrix > similarity_threshold) & \
           ~torch.eye(N, dtype=torch.bool)

    edge_pairs = mask.nonzero(as_tuple=False)   # [E, 2]

    if len(edge_pairs) > 0:
        edge_index = edge_pairs.T               # [2, E]
        print(f"  Tạo được {edge_index.shape[1]} edges (threshold={similarity_threshold})")
    else:
        # Fallback: nếu không có edges, kết nối mỗi node với top-1 neighbor
        print(f"  Không có edges trên ngưỡng {similarity_threshold}, "
              f"dùng top-1 nearest neighbors...")
        sim_matrix.fill_diagonal_(-1.0)
        top1 = sim_matrix.argmax(dim=1)         # [N]
        src = torch.arange(N)
        edge_index = torch.stack([src, top1])   # [2, N]

    # ── Step 4: Thêm self-loops và loại bỏ duplicate edges ───────────────────
    # Self-loops quan trọng cho GCN:
    # Khi aggregate neighbors, node cũng aggregate chính nó
    edge_index, _ = add_self_loops(edge_index, num_nodes=N)
    edge_index = coalesce(edge_index, num_nodes=N)  # remove duplicates

    # ── Step 5: Build PyG Data object ────────────────────────────────────────
    pyg_data = Data(x=x, edge_index=edge_index, num_nodes=N)

    print(f"✓ PyG graph: {N} nodes, {edge_index.shape[1]} edges")
    print(f"  Average degree: {edge_index.shape[1]/N:.1f}")

    return pyg_data, text_nodes, x