Реализация Graph RAG (извлечение из графа знаний)

Реализация Graph RAG (извлечение из графа знаний)

Graph RAG — архитектура, расширяющая стандартный векторный RAG структурой графа знаний. Вместо поиска только по семантически близким чанкам, система может traversить граф: от сущности через связи находить связанные концепции, которые не содержат ключевых слов запроса, но семантически релевантны. Microsoft Research в 2024 году опубликовала GraphRAG — наиболее влиятельную реализацию этого подхода.

Когда нужен Graph RAG

Стандартный RAG не справляется с:

• Вопросами об отношениях между сущностями («Как связаны компания X и контракт Y?»)
• Глобальными суммаризирующими вопросами («Каковы основные темы в корпусе документов?»)
• Multi-hop reasoning («Кто является руководителем отдела, ответственного за договор №123?»)
• Трекингом изменений во времени

Архитектура Microsoft GraphRAG

Документы
    ↓
LLM извлекает сущности и связи
    ↓
Граф знаний (NetworkX/Neo4j)
    ↓
Иерархическое сообщество-обнаружение (Leiden algorithm)
    ↓
Суммари сообществ → Community reports
    ↓
Два режима поиска:
├── Local search: вектор + граф-traversal от точки
└── Global search: суммаризация community reports

Извлечение сущностей и связей через LLM

from openai import OpenAI
import json

client = OpenAI()

ENTITY_EXTRACTION_PROMPT = """Извлеки сущности и связи из следующего текста.
Верни JSON:
{{
  "entities": [
    {{"id": "1", "name": "...", "type": "PERSON|ORG|CONTRACT|REGULATION|CONCEPT", "description": "..."}}
  ],
  "relationships": [
    {{"source": "id1", "target": "id2", "relation": "SIGNED|MANAGES|REFERS_TO|PART_OF", "description": "..."}}
  ]
}}

Текст:
{text}"""

def extract_graph_elements(text: str) -> dict:
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": ENTITY_EXTRACTION_PROMPT.format(text=text)}],
        response_format={"type": "json_object"},
        temperature=0,
    )
    return json.loads(response.choices[0].message.content)

Построение графа знаний с NetworkX

import networkx as nx
from typing import List

class KnowledgeGraph:
    def __init__(self):
        self.graph = nx.DiGraph()
        self.entity_embeddings = {}

    def add_elements(self, elements: dict, source_doc: str):
        # Добавляем сущности
        for entity in elements["entities"]:
            self.graph.add_node(
                entity["id"],
                name=entity["name"],
                type=entity["type"],
                description=entity["description"],
                source=source_doc,
            )

        # Добавляем связи
        for rel in elements["relationships"]:
            self.graph.add_edge(
                rel["source"],
                rel["target"],
                relation=rel["relation"],
                description=rel["description"],
            )

    def get_subgraph(self, entity_id: str, depth: int = 2) -> nx.DiGraph:
        """Возвращает подграф вокруг сущности"""
        nodes = {entity_id}
        for _ in range(depth):
            neighbors = set()
            for node in nodes:
                neighbors.update(self.graph.predecessors(node))
                neighbors.update(self.graph.successors(node))
            nodes.update(neighbors)
        return self.graph.subgraph(nodes)

    def serialize_subgraph(self, subgraph: nx.DiGraph) -> str:
        """Преобразует подграф в текст для контекста LLM"""
        lines = []
        for node in subgraph.nodes(data=True):
            lines.append(f"Сущность: {node[1].get('name')} ({node[1].get('type')})")
            lines.append(f"  Описание: {node[1].get('description', '')}")

        for edge in subgraph.edges(data=True):
            source_name = subgraph.nodes[edge[0]].get("name", edge[0])
            target_name = subgraph.nodes[edge[1]].get("name", edge[1])
            lines.append(f"Связь: {source_name} → {target_name} ({edge[2].get('relation')})")
            lines.append(f"  {edge[2].get('description', '')}")

        return "\n".join(lines)

Local Search: GraphRAG запрос

from langchain_openai import OpenAIEmbeddings
import numpy as np

class GraphRAGRetriever:
    def __init__(self, knowledge_graph: KnowledgeGraph, vectorstore, embeddings):
        self.kg = knowledge_graph
        self.vectorstore = vectorstore
        self.embeddings = embeddings

    def local_search(self, query: str, top_k: int = 5) -> str:
        """
        Local Search: комбинирует векторный поиск
        с graph-traversal от найденных сущностей
        """
        # 1. Векторный поиск чанков
        vector_docs = self.vectorstore.similarity_search(query, k=top_k)

        # 2. Извлечение сущностей из найденных чанков
        mentioned_entities = self._extract_entities_from_docs(vector_docs, query)

        # 3. Graph traversal: расширяем контекст через связанные узлы
        graph_contexts = []
        for entity_id in mentioned_entities[:3]:
            subgraph = self.kg.get_subgraph(entity_id, depth=2)
            graph_context = self.kg.serialize_subgraph(subgraph)
            graph_contexts.append(graph_context)

        # 4. Объединяем текстовый и граф-контекст
        vector_context = "\n\n".join([d.page_content for d in vector_docs])
        graph_context = "\n\n".join(graph_contexts)

        return f"## Текстовый контекст\n{vector_context}\n\n## Контекст из графа знаний\n{graph_context}"

Практический кейс: анализ корпоративной документации

Задача: ассистент юридического отдела для анализа отношений между контрагентами, договорами и сотрудниками (6500 договоров, 12 лет истории).

Вопросы, которые не решал стандартный RAG:

• «Какие поставщики участвовали в тендерах, где победитель впоследствии был признан банкротом?»
• «Какие договоры затронет смена руководителя в компании X?»

Граф: 45 000 сущностей, 180 000 связей (Neo4j).

Результаты:

• Multi-hop вопросы (2+ прыжка): решались в 12% стандартным RAG → 71% Graph RAG
• Глобальные суммаризирующие вопросы: 34% → 82%
• Стандартные вопросы (поиск факта): сопоставимо, незначительный регресс (-3%)
• Время построения графа: 4 дня (GPT-4o для извлечения, $240)

Инструменты для Graph RAG

• Microsoft GraphRAG library: pip install graphrag — полная реализация от Microsoft
• Neo4j + LangChain: Neo4jGraph + GraphCypherQAChain для Cypher-запросов
• LlamaIndex + Knowledge Graph: KnowledgeGraphIndex
• NetworkX: легковесный граф в Python без внешних зависимостей

Сроки

• Разработка extraction pipeline (LLM → граф): 2–3 недели
• Построение графа из существующих документов: 1–4 недели
• Local/Global search реализация: 2 недели
• Тестирование и оценка: 1–2 недели
• Итого: 6–11 недель