GraphRAG: Microsoft의 글로벌-로컬 이중 검색 전략

Documents → Entity Extraction → Graph Construction → Community Detection → Hierarchical Summarization

# Microsoft GraphRAG 공식 라이브러리
pip install graphrag

# 추가 의존성
pip install networkx matplotlib pandas numpy
pip install tiktoken openai python-dotenv

SAMPLE_DOCUMENTS = [
    {
        "id": "news_1",
        "type": "news",
        "title": "AI 스타트업 알파테크, 시리즈 B 투자 유치",
        "content": """
        AI 스타트업 알파테크가 VC인 블루벤처스로부터 500억원 규모의 시리즈 B 투자를 유치했다.
        알파테크의 CEO 김철수는 "이번 투자로 RAG 기술 고도화에 집중하겠다"고 밝혔다.
        알파테크는 삼성전자, LG전자와 협력하여 엔터프라이즈 AI 솔루션을 제공하고 있다.
        """
    },
    {
        "id": "news_2",
        "type": "news",
        "title": "삼성전자, AI 반도체 신제품 발표",
        "content": """
        삼성전자가 차세대 AI 반도체 'Exynos AI'를 발표했다.
        이 칩은 알파테크의 RAG 엔진과 호환되며, 현대자동차의 자율주행 시스템에도 탑재될 예정이다.
        """
    },
    # ... 더 많은 문서
]

from dataclasses import dataclass, field
from typing import List

@dataclass
class Entity:
    """추출된 엔티티"""
    name: str
    type: str  # ORGANIZATION, PERSON, TECHNOLOGY, PRODUCT
    description: str = ""
    source_docs: List[str] = field(default_factory=list)

@dataclass
class Relationship:
    """엔티티 간 관계"""
    source: str
    target: str
    relation_type: str  # INVESTED_IN, PARTNERED_WITH, DEVELOPED
    weight: float = 1.0
    source_docs: List[str] = field(default_factory=list)

class EntityExtractor:
    """문서에서 엔티티와 관계를 추출"""

    def __init__(self, entity_definitions: dict):
        self.entity_definitions = entity_definitions

    def extract_entities(self, documents: List[dict]) -> List[Entity]:
        """문서에서 엔티티 추출"""
        entities = {}

        for doc in documents:
            content = doc['content']
            doc_id = doc['id']

            for name, (entity_type, description) in self.entity_definitions.items():
                if name in content:
                    if name not in entities:
                        entities[name] = Entity(
                            name=name,
                            type=entity_type,
                            description=description,
                            source_docs=[doc_id]
                        )
                    else:
                        if doc_id not in entities[name].source_docs:
                            entities[name].source_docs.append(doc_id)

        return list(entities.values())

    def extract_relationships(self, documents: List[dict], entities: List[Entity]) -> List[Relationship]:
        """같은 문장에 등장하는 엔티티 간 관계 추출"""
        relationships = []
        entity_names = {e.name for e in entities}

        for doc in documents:
            sentences = doc['content'].split('.')

            for sentence in sentences:
                # 문장에 등장하는 엔티티들 찾기
                found = [e for e in entities if e.name in sentence]

                # 같은 문장의 엔티티들 간 관계 생성
                for i, e1 in enumerate(found):
                    for e2 in found[i+1:]:
                        relationships.append(Relationship(
                            source=e1.name,
                            target=e2.name,
                            relation_type=self._infer_relation_type(e1, e2),
                            source_docs=[doc['id']]
                        ))

        return self._deduplicate(relationships)

추출된 엔티티: 39개
추출된 관계: 35개

=== 엔티티 유형별 분포 ===
ORGANIZATION: 15개
PERSON: 6개
TECHNOLOGY: 7개
PRODUCT: 5개
TOOL: 5개

import networkx as nx

class KnowledgeGraph:
    """GraphRAG용 Knowledge Graph"""

    def __init__(self):
        self.graph = nx.Graph()  # 무방향 (커뮤니티 탐지용)
        self.directed_graph = nx.DiGraph()  # 방향 (쿼리용)
        self.entities = {}

    def add_entities(self, entities: List[Entity]):
        for entity in entities:
            self.entities[entity.name] = entity
            self.graph.add_node(
                entity.name,
                type=entity.type,
                description=entity.description
            )

    def add_relationships(self, relationships: List[Relationship]):
        for rel in relationships:
            self.graph.add_edge(
                rel.source, rel.target,
                relation=rel.relation_type,
                weight=rel.weight
            )

degree_centrality = nx.degree_centrality(kg.graph)
top_hubs = sorted(degree_centrality.items(), key=lambda x: x[1], reverse=True)[:5]

print("=== 허브 노드 (연결이 많은 엔티티) ===")
for node, centrality in top_hubs:
    print(f"{node}: {kg.graph.degree(node)}개 연결")

=== 허브 노드 ===
알파테크: 15개 연결
삼성전자: 9개 연결
RAG: 6개 연결
LG전자: 6개 연결
현대자동차: 5개 연결

from networkx.algorithms import community

class CommunityDetector:
    """커뮤니티 탐지 및 계층 구조 생성"""

    def __init__(self, graph: nx.Graph):
        self.graph = graph
        self.communities = []
        self.node_to_community = {}

    def detect_communities(self, resolution: float = 1.0) -> List[set]:
        """
        Louvain 알고리즘으로 커뮤니티 탐지
        (Leiden 알고리즘의 간소화 버전)
        """
        communities = community.louvain_communities(
            self.graph,
            resolution=resolution,
            seed=42
        )

        self.communities = [set(c) for c in communities]

        # 노드 → 커뮤니티 매핑
        for i, comm in enumerate(self.communities):
            for node in comm:
                self.node_to_community[node] = i

        return self.communities

=== 탐지된 커뮤니티 ===

Community 0 (11 members):
  주요 멤버: 알파테크, 블루벤처스, 김철수, 이영희, 박지민
  추정 주제: AI 스타트업 및 투자 생태계

Community 1 (8 members):
  주요 멤버: 삼성전자, 현대자동차, NVIDIA, 테슬라, 웨이모
  추정 주제: 자율주행 및 AI 하드웨어

Community 2 (9 members):
  주요 멤버: LG전자, OpenAI, 구글, 아마존, 한미래
  추정 주제: 스마트홈 및 AI 어시스턴트

Community 3 (6 members):
  주요 멤버: RAG, Knowledge Graph, Vector Store, Embedding
  추정 주제: RAG 및 검색 기술

Community 4 (5 members):
  주요 멤버: LLM, 양자화, TensorRT, vLLM
  추정 주제: LLM 최적화 및 추론

class CommunitySummarizer:
    """커뮤니티별 요약 생성"""

    def __init__(self, graph: nx.Graph, communities: List[set]):
        self.graph = graph
        self.communities = communities
        self.summaries = {}

    def generate_summary(self, community_idx: int) -> str:
        """커뮤니티 요약 생성 (실제로는 LLM 사용)"""
        members = list(self.communities[community_idx])
        subgraph = self.graph.subgraph(members)

        # 엔티티 정보 수집
        entities_info = []
        for node in members[:5]:
            node_data = self.graph.nodes[node]
            entities_info.append({
                'name': node,
                'type': node_data.get('type'),
                'description': node_data.get('description')
            })

        # 관계 정보 수집
        relations_info = []
        for u, v, data in subgraph.edges(data=True):
            relations_info.append({
                'source': u,
                'target': v,
                'type': data.get('relation')
            })

        # 템플릿 기반 요약 생성
        summary = f"""이 커뮤니티는 주로 {self._get_main_types(members)} 엔티티로 구성됩니다.

주요 엔티티:
"""
        for e in entities_info:
            summary += f"- {e['name']} ({e['type']}): {e['description']}\n"

        summary += "\n핵심 관계:\n"
        for r in relations_info[:5]:
            summary += f"- {r['source']} --{r['type']}--> {r['target']}\n"

        return summary

=== Community 0 Summary ===
이 커뮤니티는 주로 조직, 인물 엔티티로 구성됩니다.

주요 엔티티:
- 알파테크 (ORGANIZATION): AI 스타트업, RAG 기술 전문
- 블루벤처스 (ORGANIZATION): 벤처캐피탈
- 김철수 (PERSON): 알파테크 CEO
- 이영희 (PERSON): 알파테크 CTO, 스탠포드 출신
- 박지민 (PERSON): 블루벤처스 파트너

핵심 관계:
- 알파테크 --PARTNERED_WITH--> 블루벤처스
- 알파테크 --EMPLOYS--> 김철수
- 알파테크 --EMPLOYS--> 이영희

class GraphRAGQueryEngine:
    """
    GraphRAG 쿼리 엔진
    - Local Search: 특정 엔티티 관련 질문
    - Global Search: 전체 데이터셋 관련 질문
    """

    def __init__(self, graph, communities, summaries, node_to_community):
        self.graph = graph
        self.communities = communities
        self.summaries = summaries
        self.node_to_community = node_to_community

    def local_search(self, query: str, top_k: int = 5) -> dict:
        """
        Local Search: 특정 엔티티에 대한 질문
        """
        # 1. 쿼리에서 엔티티 찾기
        found_entities = []
        for entity_name in self.graph.entities.keys():
            if entity_name.lower() in query.lower():
                found_entities.append(entity_name)

        if not found_entities:
            return {'mode': 'local', 'context': "관련 엔티티를 찾지 못했습니다."}

        # 2. 관련 노드 수집 (1-hop, 2-hop)
        related_nodes = set()
        for entity in found_entities:
            neighbors = self.graph.get_neighbors(entity)
            related_nodes.update(neighbors)

            for neighbor in neighbors[:3]:
                second_hop = self.graph.get_neighbors(neighbor)
                related_nodes.update(second_hop[:2])

        # 3. 컨텍스트 구성
        context_parts = []
        for node in list(related_nodes)[:top_k]:
            node_info = self.graph.get_node_info(node)
            if node_info:
                context_parts.append(
                    f"- {node} ({node_info.get('type')}): {node_info.get('description')}"
                )

        return {
            'mode': 'local',
            'entities_found': found_entities,
            'context': '\n'.join(context_parts),
            'related_nodes': list(related_nodes)
        }

    def global_search(self, query: str) -> dict:
        """
        Global Search: 전체 데이터셋에 대한 질문
        """
        # 모든 커뮤니티 요약 수집
        all_summaries = []
        for idx, summary in self.summaries.items():
            all_summaries.append(f"[Community {idx}]\n{summary}")

        # 전역 컨텍스트 구성
        global_context = f"""=== 데이터셋 개요 ===
총 {len(self.communities)}개의 커뮤니티, {sum(len(c) for c in self.communities)}개의 엔티티

=== 커뮤니티별 요약 ===

"""
        global_context += '\n\n'.join(all_summaries)

        return {
            'mode': 'global',
            'context': global_context
        }

    def search(self, query: str, mode: str = 'auto') -> dict:
        """통합 검색 인터페이스"""
        if mode == 'local':
            return self.local_search(query)
        elif mode == 'global':
            return self.global_search(query)
        else:
            # 자동 모드 결정
            global_keywords = ['전체', '요약', '주요', '트렌드', '테마', '개요']
            is_global = any(kw in query.lower() for kw in global_keywords)

            if is_global:
                return self.global_search(query)
            else:
                return self.local_search(query)

Query: 알파테크는 어떤 회사와 협력하고 있나요?
Mode: local
Found Entities: ['알파테크']

Context:
- 삼성전자 (ORGANIZATION): 대기업, 반도체/전자
- LG전자 (ORGANIZATION): 대기업, 전자/가전
- 블루벤처스 (ORGANIZATION): 벤처캐피탈
- 알파테크 --PARTNERED_WITH--> 삼성전자
- 알파테크 --PARTNERED_WITH--> LG전자
- 알파테크 --PARTNERED_WITH--> 블루벤처스

Query: 이 데이터셋의 주요 테마와 트렌드는 무엇인가요?
Mode: global

Context:
=== 데이터셋 개요 ===
총 5개의 커뮤니티, 39개의 엔티티

=== 커뮤니티별 요약 ===

[Community 0]
AI 스타트업 및 투자 생태계...

[Community 1]
자율주행 및 AI 하드웨어...

# 1. 프로젝트 디렉토리 생성
mkdir -p ./my_graphrag/input

# 2. 입력 문서 저장 (.txt 파일들)
cp my_documents/*.txt ./my_graphrag/input/

# 3. 초기화
graphrag init --root ./my_graphrag

# 4. API 키 설정 (.env 파일)
echo "GRAPHRAG_API_KEY=your-openai-api-key" > ./my_graphrag/.env

# 5. 인덱싱 실행 (시간 소요)
graphrag index --root ./my_graphrag

# 6. Global 검색
graphrag query --root ./my_graphrag --method global \
  --query "이 문서들의 주요 테마는?"

# 7. Local 검색
graphrag query --root ./my_graphrag --method local \
  --query "알파테크에 대해 알려줘"

import asyncio
from graphrag.query.indexer_adapters import (
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.structured_search.global_search.community_context import GlobalCommunityContext
from graphrag.query.structured_search.global_search.search import GlobalSearch

# LLM 설정
llm = ChatOpenAI(
    api_key="your-api-key",
    model="gpt-4o-mini",
    api_type=OpenaiApiType.OpenAI,
)

# 인덱스 데이터 로드
INPUT_DIR = "./my_graphrag/output/artifacts"

entities = read_indexer_entities(INPUT_DIR)
relationships = read_indexer_relationships(INPUT_DIR)
reports = read_indexer_reports(INPUT_DIR)
text_units = read_indexer_text_units(INPUT_DIR)

# Global Search 설정
context_builder = GlobalCommunityContext(
    community_reports=reports,
    entities=entities,
    token_encoder=token_encoder,
)

global_search = GlobalSearch(
    llm=llm,
    context_builder=context_builder,
    token_encoder=token_encoder,
)

# 쿼리 실행
result = await global_search.asearch("이 데이터셋의 주요 테마는?")
print(result.response)

청크 1: AI 스타트업 알파테크가 VC인 블루벤처스로부터 500억원 규모...
청크 2: 삼성전자가 차세대 AI 반도체 'Exynos AI'를 발표했다...
청크 3: 현대자동차가 자율주행 레벨4 기술을 달성했다고 발표했다...

→ 문제: 개별 청크만 보여주고 "전체 테마"에 대한 답변 불가

=== 데이터셋 개요 ===
총 5개의 커뮤니티, 39개의 엔티티

핵심 테마:
1. AI 스타트업 생태계 (알파테크, 블루벤처스, 김철수, 이영희)
2. 자율주행/반도체 (삼성전자, 현대자동차, NVIDIA)
3. 스마트홈 AI (LG전자, OpenAI, 한미래)
4. RAG/검색 기술 (RAG, Knowledge Graph, Vector Store)
5. LLM 최적화 (LLM, 양자화, TensorRT)

→ 해결: 커뮤니티 요약을 통해 "숲"을 볼 수 있음

graphrag prompt-tune --root ./my_graphrag \
  --config ./settings.yaml \
  --no-entity-types

def hybrid_search(query: str):
    # 1. 질문 유형 분류
    if is_global_question(query):
        return graphrag.global_search(query)
    elif contains_entity(query):
        return graphrag.local_search(query)
    else:
        return traditional_rag.search(query)

# 커뮤니티 요약 캐시 (Redis 등)
community_summaries = cache.get("community_summaries")
if not community_summaries:
    community_summaries = generate_all_summaries()
    cache.set("community_summaries", community_summaries, ttl=3600)