Knowledge Graph RAG - Structured Reasoning with Entity Graphs

import networkx as nx
from sentence_transformers import SentenceTransformer

graph = nx.DiGraph()
embedder = SentenceTransformer("all-MiniLM-L6-v2")

def add_entity(name: str, entity_type: str, chunk_ids: list[str]):
    """Add a node with its type and associated document chunks."""
    graph.add_node(name, type=entity_type, chunks=chunk_ids)

def add_relation(source: str, relation: str, target: str):
    """Add a directed edge between two entities."""
    graph.add_edge(source, target, relation=relation)

# Populate from structured data or NER extraction
add_entity("Sarah Chen", "Person", ["chunk_12", "chunk_45"])
add_entity("Payments Team", "Team", ["chunk_88", "chunk_91"])
add_entity("Falcon Gateway", "Product", ["chunk_42", "chunk_87"])

add_relation("Sarah Chen", "manages", "Payments Team")
add_relation("Payments Team", "built", "Falcon Gateway")

def extract_entities(query: str, llm) -> list[str]:
    """Use an LLM to identify named entities in the query."""
    prompt = f"""Extract all named entities (people, products, teams, orgs)
from the following query. Return a JSON list of strings.

Query: {query}

Entities:"""
    result = llm.invoke(prompt, schema={"type": "array", "items": {"type": "string"}})
    return result

def traverse_graph(seed_entities: list[str], max_hops: int = 2) -> set[str]:
    """
    Starting from seed entities, collect all reachable nodes within max_hops.
    Returns a set of chunk IDs associated with reached nodes.
    """
    visited_nodes = set()
    chunk_ids = set()
    frontier = set(seed_entities)

    for _ in range(max_hops):
        next_frontier = set()
        for node in frontier:
            if node not in graph or node in visited_nodes:
                continue
            visited_nodes.add(node)
            # Collect document chunks linked to this node
            chunk_ids.update(graph.nodes[node].get("chunks", []))
            # Expand to neighbours
            next_frontier.update(graph.successors(node))
            next_frontier.update(graph.predecessors(node))
        frontier = next_frontier - visited_nodes

    return chunk_ids


def knowledge_graph_rag(query: str) -> tuple[str, dict]:
    # Step 1: Extract entities mentioned in the query
    entities = extract_entities(query, llm)

    # Step 2: Traverse graph from those entities
    graph_chunk_ids = traverse_graph(entities, max_hops=2)

    # Step 3: Vector search for semantic similarity
    semantic_chunks = vector_store.search(query, top_k=5)

    # Step 4: Fetch graph-referenced chunks from the document store
    graph_chunks = [doc_store.get(cid) for cid in graph_chunk_ids if doc_store.has(cid)]

    # Step 5: Merge and deduplicate
    seen = set()
    merged = []
    for chunk in graph_chunks + semantic_chunks:
        if chunk["id"] not in seen:
            seen.add(chunk["id"])
            merged.append(chunk)

    context = "\n\n".join(c["content"] for c in merged[:8])
    response = llm.generate(f"Context: {context}\n\nQuestion: {query}")

    return response, {
        "entities_found": entities,
        "graph_chunks": len(graph_chunk_ids),
        "semantic_chunks": len(semantic_chunks),
        "merged_total": len(merged),
    }