GNN-RAG: Knowledge Graph 위에서 효율적인 LLM 추론을 위한 Graph Neural Network 기반 RAG

# GNN-RAG 파이프라인 개념 코드 (PyG + LangChain 스타일)
import torch
from torch_geometric.nn import RGCNConv

# 1. GNN 리트리버: 질문 엔티티 주변 서브그래프에서 관련 경로 추출
class GNNRetriever(torch.nn.Module):
    def __init__(self, in_channels, hidden, num_relations):
        super().__init__()
        self.conv1 = RGCNConv(in_channels, hidden, num_relations)
        self.conv2 = RGCNConv(hidden, hidden, num_relations)
        self.score = torch.nn.Linear(hidden, 1)

    def forward(self, x, edge_index, edge_type):
        x = self.conv1(x, edge_index, edge_type).relu()
        x = self.conv2(x, edge_index, edge_type).relu()
        return self.score(x).squeeze(-1)  # 각 노드의 관련도 점수

# 2. 점수 높은 경로를 자연어로 변환
def paths_to_text(top_k_paths):
    """[(entity1, relation, entity2), ...] -> str"""
    lines = [f"{e1} --[{rel}]--> {e2}" for e1, rel, e2 in top_k_paths]
    return "\n".join(lines)

# 3. LLM 프롬프트에 경로 삽입
def build_prompt(question, kg_paths_text):
    return f"""다음 Knowledge Graph 경로를 참고해서 질문에 답하세요.

[KG 경로]
{kg_paths_text}

[질문]
{question}

[답변]"""

# 사용 예
# paths = gnn_retriever.get_top_k_paths(question_entity, k=20)
# prompt = build_prompt("영화 Inception의 감독 국적은?", paths_to_text(paths))
# answer = llm(prompt)