Tutorial 09: GraphRAG Knowledge Graphs¶

# examples/graphrag/convert_sources.py
import csv
import os

def csv_to_text(csv_path: str, output_dir: str):
    """Convert a CSV file to a text document for GraphRAG."""
    basename = os.path.splitext(os.path.basename(csv_path))[0]
    with open(csv_path, "r") as f:
        reader = csv.DictReader(f)
        rows = list(reader)

    text = f"# {basename}\n\n"
    text += f"This dataset contains {len(rows)} records.\n\n"
    text += f"## Columns\n\n"
    if rows:
        text += ", ".join(rows[0].keys()) + "\n\n"
        text += "## Sample Records\n\n"
        for row in rows[:10]:
            text += str(dict(row)) + "\n"

    output_path = os.path.join(output_dir, f"{basename}.txt")
    with open(output_path, "w") as f:
        f.write(text)
    print(f"Converted {csv_path} -> {output_path}")

import os
import pandas as pd
from gremlin_python.driver import client, serializer

GREMLIN_URI = os.environ["GREMLIN_ENDPOINT"].replace("https://", "wss://").replace(":443/", ":443/gremlin")
GREMLIN_KEY = os.environ["GREMLIN_KEY"]
DATABASE = "knowledge-graph"
GRAPH = "entities"

gremlin = client.Client(
    GREMLIN_URI, "g",
    username=f"/dbs/{DATABASE}/colls/{GRAPH}",
    password=GREMLIN_KEY,
    message_serializer=serializer.GraphSONSerializersV2d0(),
)

# Load entities from GraphRAG output
entities = pd.read_parquet("output/entities.parquet")
relationships = pd.read_parquet("output/relationships.parquet")

print(f"Uploading {len(entities)} entities and {len(relationships)} relationships...")

# Add vertices
for _, e in entities.iterrows():
    query = (
        f"g.addV('{e['type']}').property('id', '{e['id']}')"
        f".property('name', '{e['name']}')"
        f".property('description', '{e['description'][:200]}')"
        f".property('category', '{e['type']}')"
        f".property('pk', '{e['type']}')"
    )
    gremlin.submit(query)

# Add edges
for _, r in relationships.iterrows():
    query = (
        f"g.V('{r['source']}').addE('{r['type']}')"
        f".to(g.V('{r['target']}'))"
        f".property('description', '{r['description'][:100]}')"
    )
    gremlin.submit(query)

print("Upload complete!")
gremlin.close()

# Query: List all services
result = gremlin.submit("g.V().hasLabel('service').values('name').dedup()")
print("Services:", list(result))

# Query: Find all tables connected to Databricks
result = gremlin.submit(
    "g.V().has('name', 'Databricks').both().hasLabel('table').values('name')"
)
print("Databricks tables:", list(result))

# Query: Count entities by type
result = gremlin.submit("g.V().groupCount().by(label)")
print("Entity counts:", list(result))

import os
from openai import AzureOpenAI
from gremlin_python.driver import client, serializer

oai = AzureOpenAI(
    azure_endpoint=os.environ["AOAI_ENDPOINT"],
    api_key=os.environ["AOAI_KEY"],
    api_version="2025-04-01-preview",
)

GREMLIN_URI = os.environ["GREMLIN_ENDPOINT"].replace("https://", "wss://").replace(":443/", ":443/gremlin")
gremlin = client.Client(
    GREMLIN_URI, "g",
    username="/dbs/knowledge-graph/colls/entities",
    password=os.environ["GREMLIN_KEY"],
    message_serializer=serializer.GraphSONSerializersV2d0(),
)


def graph_context(query: str) -> str:
    """Extract relevant graph context for a query."""
    # Find entities mentioned in the query
    keywords = query.lower().split()
    entities = []
    for kw in keywords:
        if len(kw) > 3:
            result = gremlin.submit(
                f"g.V().has('name', containing('{kw}')).limit(5)"
                f".project('name','type','desc')"
                f".by('name').by(label).by('description')"
            )
            entities.extend(list(result))

    if not entities:
        return "No graph context found."

    # Get relationships for found entities
    context_parts = []
    for e in entities[:5]:
        neighbors = gremlin.submit(
            f"g.V().has('name', '{e['name']}').bothE()"
            f".project('relation','target')"
            f".by(label).by(inV().values('name'))"
        )
        rels = list(neighbors)
        context_parts.append(
            f"Entity: {e['name']} ({e['type']})\n"
            f"  Description: {e['desc']}\n"
            f"  Relationships: {rels[:5]}"
        )

    return "\n\n".join(context_parts)


def hybrid_rag(question: str) -> str:
    """Combine graph context with LLM generation."""
    ctx = graph_context(question)

    messages = [
        {"role": "system", "content": (
            "You answer questions using knowledge graph context. "
            "Reference specific entities and relationships in your answer."
        )},
        {"role": "user", "content": f"Graph Context:\n{ctx}\n\nQuestion: {question}"},
    ]

    resp = oai.chat.completions.create(
        model="gpt-54", messages=messages, temperature=0.2, max_tokens=1024
    )
    return resp.choices[0].message.content


if __name__ == "__main__":
    q = "What services are connected to Databricks and what data do they process?"
    print(f"Question: {q}\n")
    print(hybrid_rag(q))

def impact_analysis(entity_name: str, depth: int = 3) -> list[dict]:
    """Trace downstream impact of changes to an entity."""
    query = (
        f"g.V().has('name', '{entity_name}')"
        f".repeat(outE().inV()).times({depth}).emit()"
        f".path().by('name').by(label)"
    )
    result = gremlin.submit(query)
    paths = list(result)

    print(f"Impact analysis for '{entity_name}' (depth={depth}):")
    for path in paths:
        print(f"  {' -> '.join(str(p) for p in path)}")
    return paths

def lineage_trace(table_name: str) -> list[dict]:
    """Trace upstream lineage of a table."""
    query = (
        f"g.V().has('name', '{table_name}')"
        f".repeat(inE().outV()).times(5).emit()"
        f".path().by('name').by(label)"
    )
    result = gremlin.submit(query)
    paths = list(result)

    print(f"Lineage for '{table_name}':")
    for path in paths:
        print(f"  {' <- '.join(str(p) for p in path)}")
    return paths