6장: GraphRAG — 그래프 기반 검색 증강 생성

from enum import Enum
 
class SearchMode(Enum):
    GLOBAL = "global"
    LOCAL = "local"
    HYBRID = "hybrid"
 
def route_query(question: str) -> SearchMode:
    """질문 유형에 따라 검색 모드를 결정합니다."""
    # 전역 키워드 감지
    global_indicators = ["전체", "요약", "주요 주제", "동향", "개요", "전반적"]
    if any(indicator in question for indicator in global_indicators):
        return SearchMode.GLOBAL
 
    # 구체적 엔티티 언급 감지
    # 실제로는 LLM이나 NER 모델을 사용하여 더 정교하게 판단
    specific_indicators = ["어떻게", "무엇", "왜", "구체적으로"]
    if any(indicator in question for indicator in specific_indicators):
        return SearchMode.LOCAL
 
    return SearchMode.HYBRID

pip install neo4j-graphrag

from neo4j import GraphDatabase
from neo4j_graphrag.retrievers import VectorCypherRetriever
from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.embeddings import OpenAIEmbeddings
from neo4j_graphrag.generation import GraphRAG
 
# Neo4j 연결
driver = GraphDatabase.driver("bolt://localhost:7687", auth=("neo4j", "password"))
 
# 임베딩 모델
embedder = OpenAIEmbeddings(model="text-embedding-3-small")
 
# 벡터 + Cypher 하이브리드 리트리버
retriever = VectorCypherRetriever(
    driver=driver,
    index_name="document_embeddings",
    retrieval_query="""
        // 벡터 검색으로 찾은 노드에서 그래프를 순회하여 추가 컨텍스트 수집
        MATCH (node)-[:COVERS]->(tech:Technology)
        OPTIONAL MATCH (tech)-[:DEPENDS_ON]->(dep:Technology)
        WITH node, score,
             collect(DISTINCT tech.name) AS technologies,
             collect(DISTINCT dep.name) AS dependencies
        RETURN node.title AS title,
               node.content AS content,
               score,
               technologies,
               dependencies
    """,
    embedder=embedder
)
 
# LLM
llm = OpenAILLM(model_name="gpt-4o")
 
# GraphRAG 파이프라인 생성
rag = GraphRAG(retriever=retriever, llm=llm)
 
# 질의
response = rag.search(query="Neo4j에서 벡터 인덱스를 사용하는 방법은?")
print(response.answer)

from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
from neo4j_graphrag.llm import OpenAILLM
 
# 추출할 엔티티/관계 타입 정의
entities = ["Person", "Technology", "Concept", "Organization"]
relations = ["USES", "DEVELOPED_BY", "DEPENDS_ON", "IMPLEMENTS"]
 
# KG 구축 파이프라인
pipeline = SimpleKGPipeline(
    llm=OpenAILLM(model_name="gpt-4o"),
    driver=driver,
    entities=entities,
    relations=relations,
    on_error="IGNORE",  # 추출 오류 시 건너뛰기
)
 
# 텍스트에서 KG 구축
text = """
Neo4j는 프로퍼티 그래프 데이터베이스로, Cypher 쿼리 언어를 사용합니다.
GraphRAG는 Neo4j의 벡터 인덱스와 그래프 순회를 결합하여
검색 정확도를 향상시킵니다.
"""
 
result = await pipeline.run(text=text)
print(f"생성된 노드: {result.node_count}, 관계: {result.relationship_count}")

from dataclasses import dataclass
 
@dataclass
class SearchResult:
    content: str
    score: float
    source: str  # "vector", "graph", "keyword"
    metadata: dict
 
class HybridRetriever:
    """벡터 + 그래프 + 키워드 하이브리드 검색을 수행합니다."""
 
    def __init__(self, driver, embedder):
        self.driver = driver
        self.embedder = embedder
 
    def vector_search(self, query: str, top_k: int = 5) -> list[SearchResult]:
        """벡터 유사도 기반 검색을 수행합니다."""
        embedding = self.embedder.embed_query(query)
        records, _, _ = self.driver.execute_query("""
            CALL db.index.vector.queryNodes("document_embeddings", $topK, $embedding)
            YIELD node, score
            RETURN node.title AS title, node.content AS content, score
        """, topK=top_k, embedding=embedding)
 
        return [
            SearchResult(
                content=r["content"], score=r["score"],
                source="vector", metadata={"title": r["title"]}
            ) for r in records
        ]
 
    def graph_search(self, entities: list[str], depth: int = 2) -> list[SearchResult]:
        """엔티티에서 출발하여 그래프를 순회합니다."""
        records, _, _ = self.driver.execute_query("""
            UNWIND $entities AS entityName
            MATCH (e {name: entityName})
            CALL apoc.path.subgraphAll(e, {maxLevel: $depth})
            YIELD nodes, relationships
            UNWIND nodes AS n
            WHERE n.content IS NOT NULL
            RETURN DISTINCT n.title AS title, n.content AS content,
                   1.0 / (1 + size(shortestPath((e)-[*]-(n)))) AS score
        """, entities=entities, depth=depth)
 
        return [
            SearchResult(
                content=r["content"], score=r["score"],
                source="graph", metadata={"title": r["title"]}
            ) for r in records
        ]
 
    def keyword_search(self, query: str, top_k: int = 5) -> list[SearchResult]:
        """전문 검색 인덱스를 활용한 키워드 검색을 수행합니다."""
        records, _, _ = self.driver.execute_query("""
            CALL db.index.fulltext.queryNodes("document_fulltext", $query)
            YIELD node, score
            RETURN node.title AS title, node.content AS content, score
            LIMIT $topK
        """, query=query, topK=top_k)
 
        return [
            SearchResult(
                content=r["content"], score=r["score"],
                source="keyword", metadata={"title": r["title"]}
            ) for r in records
        ]
 
    def hybrid_search(self, query: str, entities: list[str],
                      weights: dict = None) -> list[SearchResult]:
        """세 가지 검색을 결합합니다."""
        if weights is None:
            weights = {"vector": 0.4, "graph": 0.4, "keyword": 0.2}
 
        # 각 검색 수행
        vector_results = self.vector_search(query)
        graph_results = self.graph_search(entities) if entities else []
        keyword_results = self.keyword_search(query)
 
        # 가중치 적용
        all_results = {}
        for result in vector_results:
            key = result.metadata.get("title", result.content[:50])
            all_results[key] = result
            all_results[key].score *= weights["vector"]
 
        for result in graph_results:
            key = result.metadata.get("title", result.content[:50])
            if key in all_results:
                all_results[key].score += result.score * weights["graph"]
            else:
                result.score *= weights["graph"]
                all_results[key] = result
 
        for result in keyword_results:
            key = result.metadata.get("title", result.content[:50])
            if key in all_results:
                all_results[key].score += result.score * weights["keyword"]
            else:
                result.score *= weights["keyword"]
                all_results[key] = result
 
        # 점수 순 정렬
        sorted_results = sorted(all_results.values(), key=lambda r: r.score, reverse=True)
        return sorted_results

class GraphRAGPipeline:
    """GraphRAG 전체 파이프라인입니다."""
 
    def __init__(self, retriever: HybridRetriever, llm):
        self.retriever = retriever
        self.llm = llm
 
    def answer(self, question: str) -> str:
        """질문에 대한 답변을 생성합니다."""
        # 1. 쿼리 라우팅
        mode = route_query(question)
 
        # 2. 검색
        if mode == SearchMode.GLOBAL:
            context = self._global_search(question)
        elif mode == SearchMode.LOCAL:
            entities = self._extract_entities(question)
            results = self.retriever.hybrid_search(question, entities)
            context = "\n\n".join([r.content for r in results[:5]])
        else:
            entities = self._extract_entities(question)
            results = self.retriever.hybrid_search(question, entities)
            context = "\n\n".join([r.content for r in results[:5]])
 
        # 3. LLM 응답 생성
        prompt = f"""다음 컨텍스트를 기반으로 질문에 답변하세요.
 
컨텍스트:
{context}
 
질문: {question}
 
답변:"""
        return self.llm.generate(prompt)
 
    def _extract_entities(self, question: str) -> list[str]:
        """질문에서 핵심 엔티티를 추출합니다."""
        # 실제 구현에서는 NER 모델이나 LLM을 사용
        # 여기서는 간략화
        return []
 
    def _global_search(self, question: str) -> str:
        """커뮤니티 요약을 활용한 글로벌 검색을 수행합니다."""
        records, _, _ = self.retriever.driver.execute_query("""
            MATCH (c:Community)
            RETURN c.summary AS summary, c.title AS title
        """)
        summaries = [f"[{r['title']}] {r['summary']}" for r in records]
        return "\n\n".join(summaries)