GraphRAG та графи знань: інтеграція LLM з Neo4j та Memgraph

[ Unstructured Documents ]
           |
           v
[ Triplet Extraction (LLM + Zod) ] ---> (Subject)-[:PREDICATE]->(Object)
           |
           v
[ Knowledge Graph (Neo4j / Memgraph) ]
           |
           v
[ Leiden Community Detection Algorithm ]
           |
           +---> Level 0: Микро-кластеры (тесно связанные компоненты)
           +---> Level 1: Мезо-кластеры (функциональные подсистемы)
           +---> Level 2: Макро-кластеры (доменные архитектурные области)
           |
           v
[ Pre-computed Community Summaries ]
           |
+-----------------------------------------------------------------------+
| QUERY PROCESSING DUAL ENGINE                                          |
|                                                                       |
|  [ User Query: Multi-hop question ]    [ User Query: Global summary ] |
|                 |                                     |               |
|                 v                                     v               |
|       [ Local Search Engine ]               [ Global Search Engine ]  |
|   1. Entity Vector Matching             1. Filter Community Summaries |
|   2. K-Hop Graph Traversal (Cypher)     2. Map-Reduce over Clusters   |
|   3. Extract Subgraph Context           3. Synthesize Global Answer   |
+-----------------------------------------------------------------------+

import OpenAI from "openai";
import { z } from "zod";

// Zod-схема графовых сущностей и связей
export const ExtractedEntitySchema = z.object({
  id: z.string().describe("Каноническое нормализованное имя сущности в snake_case"),
  name: z.string().describe("Отображаемое имя сущности"),
  type: z.enum(["SERVICE", "DATABASE", "API_ENDPOINT", "EVENT", "DEVELOPER"]),
  description: z.string().describe("Краткое резюме роли сущности в контексте"),
});

export const ExtractedRelationSchema = z.object({
  sourceEntityId: z.string(),
  targetEntityId: z.string(),
  predicate: z.enum(["CALLS", "READS_FROM", "WRITES_TO", "PUBLISHES", "SUBSCRIBES_TO", "MAINTAINS"]),
  description: z.string().describe("Контекст и детали взаимосвязи"),
  weight: z.number().min(0).max(1).default(1.0),
});

export const GraphExtractionResultSchema = z.object({
  entities: z.array(ExtractedEntitySchema),
  relations: z.array(ExtractedRelationSchema),
});

export type GraphExtractionResult = z.infer<typeof GraphExtractionResultSchema>;

export class TripletExtractor {
  private openai: OpenAI;

  constructor(apiKey: string) {
    this.openai = new OpenAI({ apiKey });
  }

  public async extractFromText(text: string): Promise<GraphExtractionResult> {
    const prompt = `
Извлеки все архитектурные сущности и направленные взаимосвязи между ними из технического описания системы.
Нормализуй идентификаторы сущностей к snake_case (например: "auth_service", "billing_db").

ТЕКСТ:
${text}
`;

    const response = await this.openai.chat.completions.create({
      model: "gpt-4o",
      temperature: 0.0,
      messages: [{ role: "user", content: prompt }],
      response_format: { type: "json_object" },
    });

    const rawJson = JSON.parse(response.choices[0]?.message.content ?? "{}");
    return GraphExtractionResultSchema.parse(rawJson);
  }
}

import neo4j, { Driver, Session } from "neo4j-driver";
import { GraphExtractionResult } from "./triplet-extractor";

export class Neo4jGraphService {
  private driver: Driver;

  constructor(uri: string, user: string, pass: string) {
    this.driver = neo4j.driver(uri, neo4j.auth.basic(user, pass));
  }

  public async close(): Promise<void> {
    await this.driver.close();
  }

  /**
   * Атомарная запись извлеченных сущностей и связей в граф
   */
  public async ingestGraphData(data: GraphExtractionResult): Promise<void> {
    const session: Session = this.driver.session();

    try {
      await session.executeWrite(async (tx) => {
        // 1. Создание или обновление сущностей
        for (const entity of data.entities) {
          await tx.run(
            `
            MERGE (e:Entity {id: $id})
            ON CREATE SET 
              e.name = $name,
              e.type = $type,
              e.description = $description,
              e.createdAt = datetime()
            ON MATCH SET
              e.description = e.description + " | " + $description,
              e.updatedAt = datetime()
            `,
            {
              id: entity.id,
              name: entity.name,
              type: entity.type,
              description: entity.description,
            }
          );
        }

        // 2. Создание связей между сущностями
        for (const rel of data.relations) {
          // Динамический тип связи через аппроксимацию в Cypher
          await tx.run(
            `
            MATCH (source:Entity {id: $sourceId})
            MATCH (target:Entity {id: $targetId})
            MERGE (source)-[r:RELATION {predicate: $predicate}]->(target)
            ON CREATE SET 
              r.description = $description,
              r.weight = $weight,
              r.createdAt = datetime()
            ON MATCH SET
              r.weight = (r.weight + $weight) / 2.0
            `,
            {
              sourceId: rel.sourceEntityId,
              targetId: rel.targetEntityId,
              predicate: rel.predicate,
              description: rel.description,
              weight: rel.weight,
            }
          );
        }
      });
    } finally {
      await session.close();
    }
  }

  /**
   * Выполнение K-Hop поиска окрестностей сущности для Local Search
   */
  public async getEntityNeighborhood(entityId: string, maxHops: number = 2): Promise<string> {
    const session = this.driver.session();
    try {
      const result = await session.run(
        `
        MATCH path = (start:Entity {id: $entityId})-[*1..2]-(neighbor:Entity)
        WITH start, relationships(path) AS rels, nodes(path) AS ns
        UNWIND rels AS r
        UNWIND ns AS n
        RETURN DISTINCT
          start.name AS root,
          start.type AS rootType,
          type(r) AS relationType,
          r.predicate AS predicate,
          r.description AS relDescription,
          n.name AS neighborName,
          n.type AS neighborType
        LIMIT 50
        `,
        { entityId }
      );

      const formattedContext = result.records.map((record) => {
        return `[${record.get("root")}] -(${record.get("predicate")}: ${record.get("relDescription")})-> [${record.get("neighborName")}]`;
      });

      return formattedContext.join("\n");
    } finally {
      await session.close();
    }
  }
}

import OpenAI from "openai";
import { Neo4jGraphService } from "./neo4j-service";

export interface GraphRAGResponse {
  answer: string;
  graphContextUsed: string[];
}

export class HybridGraphRAGPipeline {
  private openai: OpenAI;
  private graphService: Neo4jGraphService;

  constructor(apiKey: string, graphService: Neo4jGraphService) {
    this.openai = new OpenAI({ apiKey });
    this.graphService = graphService;
  }

  /**
   * Поиск ответа с использованием локального графового контекста
   */
  public async executeLocalSearch(
    query: string,
    targetEntityId: string
  ): Promise<GraphRAGResponse> {
    // 1. Извлекаем 2-hop подграф вокруг целевой сущности
    const graphContext = await this.graphService.getEntityNeighborhood(targetEntityId, 2);

    // 2. Формируем промпт с топологическим контекстом графа
    const prompt = `
Ты — ведущий системный архитектор. Ответь на вопрос пользователя, опираясь ИСКЛЮЧИТЕЛЬНО на предоставленную топологию графа знаний.

ТОПОЛОГИЯ ГРАФА СВЯЗЕЙ:
${graphContext || "Связи не обнаружены."}

ВОПРОС:
${query}
`;

    const completion = await this.openai.chat.completions.create({
      model: "gpt-4o",
      temperature: 0.1,
      messages: [{ role: "user", content: prompt }],
    });

    return {
      answer: completion.choices[0]?.message.content ?? "",
      graphContextUsed: graphContext.split("\n").filter(Boolean),
    };
  }
}