Graph RAG: граф знаний для сложных связей между сущностями

from openai import AsyncOpenAI
from pydantic import BaseModel

client = AsyncOpenAI()


class Triplet(BaseModel):
    subject: str
    relation: str
    obj: str       # 'object' зарезервировано в Python
    subject_type: str   # Technology, Person, Team, Service, Concept...
    object_type: str
    confidence: float   # 0.0–1.0


class ExtractionResult(BaseModel):
    triplets: list[Triplet]
    entities: list[str]    # все упомянутые сущности, включая без отношений


EXTRACTION_PROMPT = """Извлеки сущности и отношения из текста в виде троек.

Правила:
- Каждая тройка: (субъект, ТYPEOFRELATION, объект)
- Отношения пиши в UPPER_SNAKE_CASE: USES, DEPENDS_ON, PART_OF, MAINTAINS, DEPLOYED_IN,
  AFFECTS, FIXES, OWNED_BY, INTEGRATES_WITH, REPLACES, COMMUNICATES_WITH
- Субъект и объект — именованные сущности (не местоимения, не «система», не «компонент»)
- Если нет явного отношения — не придумывай
- confidence: 1.0 если явно сказано, 0.7 если следует из контекста

Типы сущностей: Technology, Service, Team, Person, Document, Version, Region,
Vulnerability, Concept, Product, Configuration"""


async def extract_triplets(text: str, source: str = "") -> ExtractionResult:
    """Извлекаем тройки из одного чанка."""
    resp = await client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": EXTRACTION_PROMPT},
            {"role": "user", "content": f"Текст:\n{text}"},
        ],
        response_format=ExtractionResult,
        temperature=0,
    )
    return resp.choices[0].message.parsed


# Пример
async def demo():
    text = """
    Сервис авторизации использует Redis для хранения сессий.
    Redis 7.2.3 развёрнут в AWS eu-west-1.
    Команда Platform обслуживает Redis и отвечает за его доступность.
    CVE-2024-31449 затрагивает Redis версии ниже 7.2.5.
    """

    result = await extract_triplets(text)
    for t in result.triplets:
        print(f"({t.subject}) --[{t.relation}]--> ({t.obj})  conf={t.confidence}")

# Вывод:
# (Сервис авторизации) --[USES]--> (Redis)  conf=1.0
# (Redis 7.2.3) --[DEPLOYED_IN]--> (AWS eu-west-1)  conf=1.0
# (Команда Platform) --[MAINTAINS]--> (Redis)  conf=1.0
# (CVE-2024-31449) --[AFFECTS]--> (Redis < 7.2.5)  conf=1.0

from dataclasses import dataclass, field
from collections import defaultdict

@dataclass
class EntityNormalizer:
    """Группирует варианты написания одной сущности."""
    canonical: dict[str, str] = field(default_factory=dict)
    aliases: dict[str, list[str]] = field(default_factory=lambda: defaultdict(list))

    async def normalize(self, entities: list[str]) -> dict[str, str]:
        """
        Возвращает маппинг alias → canonical name.
        LLM объединяет вариации.
        """
        resp = await client.beta.chat.completions.parse(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    "Получишь список сущностей из технической документации. "
                    "Определи группы синонимов — сущности, обозначающие одно и то же. "
                    "Для каждой группы выбери канонический вариант (наиболее полный и конкретный). "
                    "Верни маппинг всех вариантов на их канонические имена."
                )},
                {"role": "user", "content": "\n".join(entities)},
            ],
            response_format=AliasMapping,  # {mappings: [{alias, canonical}]}
            temperature=0,
        )
        for m in resp.choices[0].message.parsed.mappings:
            self.canonical[m.alias] = m.canonical
        return self.canonical

    def apply(self, triplets: list[Triplet]) -> list[Triplet]:
        """Нормализует субъекты и объекты в тройках."""
        normalized = []
        for t in triplets:
            subj = self.canonical.get(t.subject, t.subject)
            obj  = self.canonical.get(t.obj, t.obj)
            normalized.append(Triplet(
                subject=subj, relation=t.relation, obj=obj,
                subject_type=t.subject_type, object_type=t.object_type,
                confidence=t.confidence,
            ))
        return normalized


class AliasMapping(BaseModel):
    mappings: list[AliasEntry]

class AliasEntry(BaseModel):
    alias: str
    canonical: str

from neo4j import GraphDatabase, AsyncGraphDatabase
from dataclasses import dataclass

@dataclass
class GraphStore:
    uri: str      # "bolt://localhost:7687"
    user: str     # "neo4j"
    password: str

    def __post_init__(self):
        self.driver = AsyncGraphDatabase.driver(
            self.uri, auth=(self.user, self.password)
        )

    async def add_triplet(self, triplet: Triplet, source_doc: str = ""):
        """Добавляем тройку в граф — MERGE не создаёт дублей."""
        async with self.driver.session() as session:
            await session.run("""
                MERGE (s:Entity {name: $subject})
                SET s.type = $subject_type

                MERGE (o:Entity {name: $obj})
                SET o.type = $object_type

                MERGE (s)-[r:RELATION {type: $relation}]->(o)
                SET r.confidence = $confidence,
                    r.source = $source
            """,
                subject=triplet.subject,
                subject_type=triplet.subject_type,
                obj=triplet.obj,
                object_type=triplet.object_type,
                relation=triplet.relation,
                confidence=triplet.confidence,
                source=source_doc,
            )

    async def add_triplets_batch(self, triplets: list[Triplet], source: str = ""):
        """Пакетная загрузка через UNWIND — значительно быстрее."""
        data = [
            {
                "subject": t.subject, "subject_type": t.subject_type,
                "obj": t.obj, "object_type": t.object_type,
                "relation": t.relation, "confidence": t.confidence,
                "source": source,
            }
            for t in triplets
        ]
        async with self.driver.session() as session:
            await session.run("""
                UNWIND $rows AS row
                MERGE (s:Entity {name: row.subject})
                SET s.type = row.subject_type
                MERGE (o:Entity {name: row.obj})
                SET o.type = row.object_type
                MERGE (s)-[r:RELATION {type: row.relation}]->(o)
                SET r.confidence = row.confidence, r.source = row.source
            """, rows=data)

    async def close(self):
        await self.driver.close()


# Использование
store = GraphStore("bolt://localhost:7687", "neo4j", "password")
result = await extract_triplets(chunk_text)
await store.add_triplets_batch(result.triplets, source="architecture.md")

import networkx as nx
import json
from pathlib import Path


class NetworkXGraphStore:
    """Легковесный граф в памяти — для прототипов и тестов."""

    def __init__(self):
        self.G = nx.MultiDiGraph()   # MultiDi: несколько рёбер между парой узлов

    def add_triplets(self, triplets: list[Triplet], source: str = ""):
        for t in triplets:
            # Узлы
            if not self.G.has_node(t.subject):
                self.G.add_node(t.subject, entity_type=t.subject_type)
            if not self.G.has_node(t.obj):
                self.G.add_node(t.obj, entity_type=t.object_type)
            # Ребро
            self.G.add_edge(
                t.subject, t.obj,
                relation=t.relation,
                confidence=t.confidence,
                source=source,
            )

    def neighbors(self, entity: str, depth: int = 1) -> list[dict]:
        """Все соседи на расстоянии depth (1-hop по умолчанию)."""
        if entity not in self.G:
            return []
        if depth == 1:
            nodes = list(self.G.predecessors(entity)) + list(self.G.successors(entity))
        else:
            # BFS до depth шагов
            nodes = list(nx.single_source_shortest_path(self.G, entity, cutoff=depth).keys())

        result = []
        for n in set(nodes):
            if n == entity:
                continue
            # Получаем все рёбра между entity и n
            for _, _, data in self.G.edges(n, data=True):
                result.append({"from": n, "rel": data["relation"], "to": entity})
            for _, to, data in self.G.out_edges(entity, data=True):
                if to == n:
                    result.append({"from": entity, "rel": data["relation"], "to": n})
        return result

    def save(self, path: str):
        data = nx.node_link_data(self.G)
        Path(path).write_text(json.dumps(data, ensure_ascii=False, indent=2))

    def load(self, path: str):
        data = json.loads(Path(path).read_text())
        self.G = nx.node_link_graph(data)

from dataclasses import dataclass

@dataclass
class LocalSearchResult:
    entities_found: list[str]
    subgraph: list[dict]   # [{"from": str, "rel": str, "to": str}]
    source_chunks: list[str]
    context_text: str


async def entity_link(question: str) -> list[str]:
    """Извлекаем сущности из вопроса для поиска в графе."""
    resp = await client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": (
                "Извлеки все именованные сущности из вопроса. "
                "Только конкретные имена: сервисы, технологии, команды, продукты. "
                "Не включай общие слова."
            )},
            {"role": "user", "content": question},
        ],
        response_format=EntityList,
        temperature=0,
    )
    return resp.choices[0].message.parsed.entities


async def local_search(
    question: str,
    store: GraphStore,
    vector_retriever,           # обычный векторный retriever для чанков
    hops: int = 2,
    max_nodes: int = 20,
) -> LocalSearchResult:
    """Local graph search: от сущностей запроса через граф."""

    # Шаг 1: Entity linking
    entities = await entity_link(question)
    if not entities:
        # Нет конкретных сущностей → fallback на vector search
        chunks = vector_retriever.search(question, k=4)
        return LocalSearchResult(
            entities_found=[],
            subgraph=[],
            source_chunks=[c.page_content for c in chunks],
            context_text="\n\n".join(c.page_content for c in chunks),
        )

    # Шаг 2+3: Обход графа (Cypher: BFS до hops шагов)
    async with store.driver.session() as session:
        # Находим узлы-точки входа (fuzzy match по имени)
        entry_nodes = []
        for entity in entities:
            result = await session.run("""
                MATCH (n:Entity)
                WHERE toLower(n.name) CONTAINS toLower($name)
                RETURN n.name AS name, n.type AS type
                LIMIT 3
            """, name=entity)
            entry_nodes.extend([r["name"] async for r in result])

        if not entry_nodes:
            # Сущности не найдены в графе → vector search fallback
            chunks = vector_retriever.search(question, k=4)
            return LocalSearchResult([], [], [c.page_content for c in chunks],
                                     "\n\n".join(c.page_content for c in chunks))

        # BFS обход на hops шагов
        subgraph_result = await session.run(f"""
            MATCH path = (start:Entity)-[*1..{hops}]-(end:Entity)
            WHERE start.name IN $entry_nodes
            WITH start, end,
                 relationships(path) AS rels,
                 nodes(path) AS nodes_in_path
            LIMIT {max_nodes}
            UNWIND rels AS r
            RETURN
                startNode(r).name AS from_node,
                r.type            AS relation,
                endNode(r).name   AS to_node,
                r.source          AS source_doc
        """, entry_nodes=entry_nodes)

        subgraph = []
        source_docs = set()
        async for row in subgraph_result:
            subgraph.append({
                "from": row["from_node"],
                "rel":  row["relation"],
                "to":   row["to_node"],
            })
            if row["source_doc"]:
                source_docs.add(row["source_doc"])

    # Шаг 4: Достаём исходные чанки из векторной БД по source_doc
    source_chunks = []
    for doc_name in list(source_docs)[:5]:
        chunks = vector_retriever.search(question, k=2, filter={"source": doc_name})
        source_chunks.extend(c.page_content for c in chunks)

    # Шаг 5: Формируем контекст
    context = _build_context(entry_nodes, subgraph, source_chunks)
    return LocalSearchResult(
        entities_found=entry_nodes,
        subgraph=subgraph,
        source_chunks=source_chunks,
        context_text=context,
    )


def _build_context(
    entities: list[str],
    subgraph: list[dict],
    chunks: list[str],
) -> str:
    parts = []

    if subgraph:
        parts.append("## Связи из графа знаний")
        for edge in subgraph[:30]:
            parts.append(f"- {edge['from']} --[{edge['rel']}]--> {edge['to']}")

    if chunks:
        parts.append("\n## Релевантные фрагменты документов")
        for i, chunk in enumerate(chunks[:4], 1):
            parts.append(f"[Документ {i}]\n{chunk}")

    return "\n".join(parts)


class EntityList(BaseModel):
    entities: list[str]

import community as community_louvain  # pip install python-louvain
# Или: pip install leidenalg igraph — для оригинального Leiden
from collections import defaultdict


def detect_communities(G: nx.Graph) -> dict[int, list[str]]:
    """
    Возвращает {community_id: [node1, node2, ...]}.
    Используем Louvain (аппроксимация Leiden, легко ставится).
    """
    # Leiden/Louvain работает с ненаправленным графом
    undirected = G.to_undirected()

    # Убираем изолированные узлы для лучшего кластеризирования
    undirected.remove_nodes_from(list(nx.isolates(undirected)))

    partition = community_louvain.best_partition(undirected, resolution=1.0)
    # partition = {node: community_id}

    communities = defaultdict(list)
    for node, cid in partition.items():
        communities[cid].append(node)

    return dict(communities)


async def generate_community_summary(
    community_id: int,
    nodes: list[str],
    G: nx.Graph,
) -> str:
    """LLM генерирует текстовое описание сообщества."""
    # Собираем все рёбра внутри сообщества
    node_set = set(nodes)
    edges_text = []
    for u, v, data in G.edges(data=True):
        if u in node_set and v in node_set:
            edges_text.append(f"{u} --[{data.get('relation','связан с')}]--> {v}")

    if not edges_text:
        edges_text = [f"Сущности: {', '.join(nodes[:10])}"]

    resp = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": (
                "Напиши краткое (2–4 предложения) описание группы связанных сущностей "
                "и их взаимоотношений. Текст для поиска — должен отвечать на вопросы "
                "о теме этой группы."
            )},
            {"role": "user", "content": (
                f"Сущности и связи:\n" + "\n".join(edges_text[:30])
            )},
        ],
        temperature=0.2,
    )
    return resp.choices[0].message.content


async def build_community_index(
    G: nx.Graph,
    vector_store,   # любой vectorstore с .add_texts()
) -> dict[int, str]:
    """Строим и индексируем все community summaries."""
    communities = detect_communities(G)
    summaries = {}

    for cid, nodes in communities.items():
        if len(nodes) < 2:
            continue    # одиночные узлы не суммаризируем
        summary = await generate_community_summary(cid, nodes, G)
        summaries[cid] = summary

    # Индексируем в векторную БД
    texts = list(summaries.values())
    metadatas = [{"community_id": cid, "nodes": ",".join(communities[cid][:20])}
                 for cid in summaries]
    vector_store.add_texts(texts, metadatas=metadatas)

    print(f"Построено {len(summaries)} community summaries")
    return summaries

async def global_search(
    question: str,
    community_vector_store,
    top_k: int = 5,
) -> str:
    """
    Global search: ищем по community summaries.
    Хорош для широких/обзорных вопросов.
    """
    # Поиск релевантных сообществ
    relevant_communities = community_vector_store.similarity_search(question, k=top_k)

    if not relevant_communities:
        return "Не удалось найти релевантные разделы в базе знаний."

    # Собираем контекст из summaries
    context_parts = []
    for i, doc in enumerate(relevant_communities, 1):
        cid = doc.metadata.get("community_id", i)
        context_parts.append(f"[Кластер {cid}]\n{doc.page_content}")

    context = "\n\n".join(context_parts)

    # Генерация ответа
    resp = await client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": (
                "Ты отвечаешь на вопросы на основе структурированных данных о связях "
                "между сущностями. Используй предоставленные кластеры знаний."
            )},
            {"role": "user", "content": (
                f"Контекст:\n{context}\n\nВопрос: {question}"
            )},
        ],
        temperature=0.1,
    )
    return resp.choices[0].message.content

from enum import StrEnum
from dataclasses import dataclass


class SearchMode(StrEnum):
    LOCAL  = "local"
    GLOBAL = "global"
    HYBRID = "hybrid"   # local + global, объединяем контексты


@dataclass
class GraphRAGConfig:
    hops: int = 2
    max_nodes: int = 30
    community_top_k: int = 5
    local_chunks_k: int = 3
    model: str = "gpt-4o-mini"
    reasoning_model: str = "gpt-4o"


class GraphRAG:
    def __init__(
        self,
        graph_store: GraphStore,
        community_store,    # vectorstore с summaries
        vector_retriever,   # обычный RAG retriever
        config: GraphRAGConfig = None,
    ):
        self.graph  = graph_store
        self.comms  = community_store
        self.retriever = vector_retriever
        self.cfg    = config or GraphRAGConfig()

    async def answer(
        self,
        question: str,
        mode: SearchMode = SearchMode.HYBRID,
    ) -> str:
        if mode == SearchMode.LOCAL:
            return await self._local_answer(question)
        if mode == SearchMode.GLOBAL:
            return await global_search(question, self.comms, self.cfg.community_top_k)
        # HYBRID: пробуем local, если мало контекста — добавляем global
        return await self._hybrid_answer(question)

    async def _local_answer(self, question: str) -> str:
        result = await local_search(
            question, self.graph, self.retriever,
            hops=self.cfg.hops, max_nodes=self.cfg.max_nodes,
        )
        if not result.subgraph and not result.source_chunks:
            # Граф не помог → global fallback
            return await global_search(question, self.comms)

        resp = await client.chat.completions.create(
            model=self.cfg.reasoning_model,
            messages=[
                {"role": "system", "content": (
                    "Ты отвечаешь на вопросы, используя граф знаний и фрагменты документов. "
                    "Ссылайся на конкретные сущности из графа."
                )},
                {"role": "user", "content": (
                    f"Контекст из графа:\n{result.context_text}\n\n"
                    f"Вопрос: {question}"
                )},
            ],
            temperature=0.1,
        )
        return resp.choices[0].message.content

    async def _hybrid_answer(self, question: str) -> str:
        import asyncio
        # Параллельно: local search + community search
        local_result, global_docs = await asyncio.gather(
            local_search(question, self.graph, self.retriever, self.cfg.hops),
            self.comms.asimilarity_search(question, k=3),
        )

        # Объединяем контексты
        context_parts = []
        if local_result.subgraph:
            context_parts.append("## Граф: прямые связи\n" + local_result.context_text)
        if global_docs:
            summaries = "\n\n".join(
                f"[Кластер]\n{d.page_content}" for d in global_docs
            )
            context_parts.append("## Обзор по теме\n" + summaries)
        if local_result.source_chunks:
            chunks_text = "\n\n".join(local_result.source_chunks[:3])
            context_parts.append("## Исходные документы\n" + chunks_text)

        context = "\n\n".join(context_parts) or "Информация не найдена."

        resp = await client.chat.completions.create(
            model=self.cfg.reasoning_model,
            messages=[
                {"role": "system", "content": (
                    "Ты отвечаешь на вопросы используя граф знаний. "
                    "Сначала используй прямые связи из графа, затем — обзорные данные."
                )},
                {"role": "user", "content": f"Контекст:\n{context}\n\nВопрос: {question}"},
            ],
            temperature=0.1,
        )
        return resp.choices[0].message.content

import asyncio
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document


async def index_documents(
    documents: list[Document],
    graph_store: GraphStore | NetworkXGraphStore,
    community_vector_store,
    chunk_size: int = 800,
    concurrency: int = 5,
) -> dict:
    """Полный цикл: документы → граф → community summaries."""

    # 1. Chunking
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size, chunk_overlap=100
    )
    chunks = splitter.split_documents(documents)
    print(f"Получено {len(chunks)} чанков из {len(documents)} документов")

    # 2. Параллельное извлечение троек
    semaphore = asyncio.Semaphore(concurrency)
    all_triplets: list[tuple[list[Triplet], str]] = []

    async def process_chunk(chunk: Document):
        async with semaphore:
            result = await extract_triplets(chunk.page_content)
            source = chunk.metadata.get("source", "unknown")
            return result.triplets, source

    results = await asyncio.gather(*[process_chunk(c) for c in chunks])

    # 3. Нормализация сущностей
    all_entities = []
    for triplets, _ in results:
        all_entities.extend([t.subject for t in triplets])
        all_entities.extend([t.obj for t in triplets])

    normalizer = EntityNormalizer()
    await normalizer.normalize(list(set(all_entities)))

    # 4. Загрузка в граф
    for triplets, source in results:
        normalized = normalizer.apply(triplets)
        if isinstance(graph_store, GraphStore):
            await graph_store.add_triplets_batch(normalized, source)
        else:
            graph_store.add_triplets(normalized, source)

    total_triplets = sum(len(t) for t, _ in results)
    print(f"Загружено {total_triplets} троек в граф")

    # 5. Community detection + summaries
    if isinstance(graph_store, NetworkXGraphStore):
        G = graph_store.G
    else:
        G = await _load_graph_from_neo4j(graph_store)  # экспорт в nx для Leiden

    summaries = await build_community_index(G, community_vector_store)

    return {
        "chunks": len(chunks),
        "triplets": total_triplets,
        "communities": len(summaries),
    }


# Запуск
async def main():
    from langchain_community.document_loaders import DirectoryLoader

    loader = DirectoryLoader("./docs", glob="**/*.md")
    documents = loader.load()

    stats = await index_documents(
        documents,
        graph_store=NetworkXGraphStore(),
        community_vector_store=chroma_db,
    )
    print(f"Indexed: {stats}")

import networkx as nx

G = nx.MultiDiGraph()

# Индексирование
triplets = await extract_triplets(chunk_text)
for t in triplets.triplets:
    G.add_edge(t.subject, t.obj, relation=t.relation)

# Local search
entity = "Redis"
neighbors = list(G.predecessors(entity)) + list(G.successors(entity))
context = "\n".join(
    f"{u} --[{d['relation']}]--> {v}"
    for u, v, d in G.edges(data=True)
    if u in neighbors or v in neighbors
)
answer = await llm(question, context)