Building GraphRAG with the Gemini API — A Complete Production Guide to Hybrid Knowledge Graph + Vector Retrieval

Step 1: Extract a Knowledge Graph from Documents with Gemini

The first pipeline pulls (subject, predicate, object) triples out of raw documents. The non-negotiable design decision here: enforce the extraction schema with Function Calling. If you ask the model for JSON in a prompt, you will hit JSONDecodeError in production. I've been there more than once.

# pip install google-genai neo4j
import os
import json
from google import genai
from google.genai import types
 
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 
EXTRACT_GRAPH_TOOL = types.Tool(
    function_declarations=[
        types.FunctionDeclaration(
            name="store_knowledge_graph",
            description="Persist knowledge graph triples extracted from a document",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "triples": types.Schema(
                        type=types.Type.ARRAY,
                        items=types.Schema(
                            type=types.Type.OBJECT,
                            properties={
                                "subject": types.Schema(type=types.Type.STRING, description="Subject entity name"),
                                "subject_type": types.Schema(type=types.Type.STRING, description="Person / Product / Tech / Concept etc."),
                                "predicate": types.Schema(type=types.Type.STRING, description="Relation name (uses, depends_on, created_by ...)"),
                                "object": types.Schema(type=types.Type.STRING, description="Object entity name"),
                                "object_type": types.Schema(type=types.Type.STRING),
                                "evidence": types.Schema(type=types.Type.STRING, description="Direct quote from the source supporting the triple"),
                            },
                            required=["subject", "subject_type", "predicate", "object", "object_type", "evidence"],
                        ),
                    )
                },
                required=["triples"],
            ),
        )
    ]
)
 
def extract_triples(document: str, source_id: str) -> list[dict]:
    """Extract knowledge graph triples from a document.
    Returns [{subject, subject_type, predicate, object, object_type, evidence, source_id}, ...]
    """
    prompt = f"""Read the document below and extract triples that capture relationships between entities.
- Normalize entity names so the same entity always uses the same surface form (e.g. "Gemini API" and "the API of Gemini" both become "Gemini API").
- Only extract relationships that are explicitly stated in the document. Do not infer.
- For evidence, quote one or two sentences from the source.
 
Document:
{document}
"""
    response = client.models.generate_content(
        model="gemini-2.5-pro",
        contents=prompt,
        config=types.GenerateContentConfig(
            tools=[EXTRACT_GRAPH_TOOL],
            tool_config=types.ToolConfig(
                function_calling_config=types.FunctionCallingConfig(mode="ANY")
            ),
            temperature=0.0,
        ),
    )
 
    fc = response.candidates[0].content.parts[0].function_call
    if fc is None or fc.name != "store_knowledge_graph":
        return []
    triples = fc.args.get("triples", [])
    for t in triples:
        t["source_id"] = source_id
    return triples
 
 
if __name__ == "__main__":
    sample = """Claude Lab, part of Dolice Labs, is deployed to Cloudflare Workers.
Claude Lab uses Stripe for billing, and its content is authored in MDX.
Claude Lab was built by Masaki."""
    triples = extract_triples(sample, source_id="claudelab-readme")
    print(json.dumps(triples, ensure_ascii=False, indent=2))
    # Expected output:
    # [
    #   {"subject": "Claude Lab", "predicate": "deployed_on", "object": "Cloudflare Workers", ...},
    #   {"subject": "Claude Lab", "predicate": "uses", "object": "Stripe", ...},
    #   {"subject": "Claude Lab", "predicate": "created_by", "object": "Masaki", ...},
    #   ...
    # ]

Two details matter here. First, tool_config with mode="ANY" forces Gemini to return a function call. Without it, you'll occasionally get a free-form response and your downstream parser will choke. Second, temperature=0.0 is non-negotiable for extraction — this isn't a creative task, and any temperature above zero introduces entity-name drift (e.g. "Claude Lab" and "Claudelab" registered as separate nodes), which destroys graph quality over time.

For long documents, split first with something like LangChain's RecursiveCharacterTextSplitter. Cross-chunk relationships will be lost at the boundary, so I run larger overlap (around 500 tokens) than a typical RAG (200 tokens) to recover most of them.

Step 2: Index the Graph and the Vector Store in Parallel

Each extracted triple goes into Neo4j, and the same chunk is embedded and stored in the vector store. The critical design choice here: link the two through a shared source_id. Without this, you can't fetch the original chunk for an entity that the graph traversal returned.

# pip install neo4j
from neo4j import GraphDatabase
 
driver = GraphDatabase.driver(
    os.environ["NEO4J_URI"],
    auth=(os.environ["NEO4J_USER"], os.environ["NEO4J_PASSWORD"]),
)
 
UPSERT_QUERY = """
MERGE (s:Entity {name: $subject})
  ON CREATE SET s.type = $subject_type
MERGE (o:Entity {name: $object})
  ON CREATE SET o.type = $object_type
MERGE (s)-[r:RELATES {predicate: $predicate, source_id: $source_id}]->(o)
  ON CREATE SET r.evidence = $evidence, r.created_at = timestamp()
"""
 
def write_triples_to_neo4j(triples: list[dict]) -> None:
    """Idempotently write triples to Neo4j."""
    with driver.session() as session:
        for t in triples:
            try:
                session.run(UPSERT_QUERY, **t)
            except Exception as e:
                # Log and continue so a single bad triple doesn't stop the batch
                print(f"⚠️ failed to upsert triple: {t} — {e}")
 
 
def index_document(document: str, source_id: str, vector_store) -> None:
    """Write a document to both the graph and the vector store."""
    triples = extract_triples(document, source_id=source_id)
    if not triples:
        print(f"⚠️ no triples extracted for {source_id}")
        return
    write_triples_to_neo4j(triples)
 
    embedding = client.models.embed_content(
        model="gemini-embedding-001",
        contents=document,
        config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT"),
    ).embeddings[0].values
    vector_store.upsert(
        id=source_id,
        vector=embedding,
        metadata={"text": document, "triples_count": len(triples)},
    )

MERGE keeps the writes idempotent. Re-ingesting the same document doesn't duplicate entities or relationships, which matters because re-indexing happens routinely in production. Build idempotency in from day one — bolting it on later is painful.

task_type="RETRIEVAL_DOCUMENT" is doing real work in this snippet. On the retrieval side I use RETRIEVAL_QUERY. Gemini's embedding model maps "this is the document being searched" and "this is the query searching for it" into different points in the same space, and you only get the precision benefits if you tag both sides correctly.

Step 3: Route the Question Between Graph Traversal and Vector Search

The retrieval-time design is what determines whether GraphRAG is worth the complexity. My approach starts by classifying the question with Gemini Flash and routing accordingly.

ROUTE_TOOL = types.Tool(
    function_declarations=[
        types.FunctionDeclaration(
            name="route_question",
            description="Classify the question into the optimal retrieval route",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "route": types.Schema(
                        type=types.Type.STRING,
                        enum=["GRAPH_ONLY", "VECTOR_ONLY", "HYBRID"],
                        description="GRAPH_ONLY=explicit relations between entities / VECTOR_ONLY=fuzzy or context-dependent / HYBRID=needs both",
                    ),
                    "entities": types.Schema(
                        type=types.Type.ARRAY,
                        items=types.Schema(type=types.Type.STRING),
                        description="Normalized entity names mentioned in the question",
                    ),
                },
                required=["route", "entities"],
            ),
        )
    ]
)
 
def route_query(question: str) -> dict:
    """Classify the question and return routing metadata."""
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=f"Classify the following question: {question}",
        config=types.GenerateContentConfig(
            tools=[ROUTE_TOOL],
            tool_config=types.ToolConfig(
                function_calling_config=types.FunctionCallingConfig(mode="ANY")
            ),
            temperature=0.0,
        ),
    )
    return response.candidates[0].content.parts[0].function_call.args
 
 
def graph_traverse(entities: list[str], hops: int = 2) -> list[dict]:
    """Fetch the subgraph reachable within N hops from the seed entities."""
    cypher = """
    MATCH path = (start:Entity)-[*1..$hops]-(neighbor:Entity)
    WHERE start.name IN $entities
    RETURN start.name AS start, neighbor.name AS neighbor,
           [rel IN relationships(path) | {predicate: rel.predicate, evidence: rel.evidence, source_id: rel.source_id}] AS path
    LIMIT 50
    """
    with driver.session() as session:
        result = session.run(cypher, entities=entities, hops=hops)
        return [r.data() for r in result]
 
 
def vector_search(question: str, top_k: int = 8, vector_store=None) -> list[dict]:
    """Retrieve chunks via embedding similarity."""
    query_embedding = client.models.embed_content(
        model="gemini-embedding-001",
        contents=question,
        config=types.EmbedContentConfig(task_type="RETRIEVAL_QUERY"),
    ).embeddings[0].values
    return vector_store.query(vector=query_embedding, top_k=top_k, include_metadata=True)

hops=2 is empirical. One hop is too shallow; three or more rapidly expands into noise. On my benchmark set, accuracy on relational questions plateaus at two hops, and three hops actually drags answer quality down because Gemini gets distracted by tangential edges.

Step 4: Fuse the Context and Answer with Gemini 2.5 Pro

The synthesis step is where GraphRAG earns its keep. The subgraph from traversal gets serialized into a human-readable form, lined up alongside the vector hits, and handed to Gemini 2.5 Pro.

def subgraph_to_text(subgraph: list[dict]) -> str:
    """Convert a subgraph into a text format Gemini can read fluently."""
    lines = []
    for row in subgraph:
        path_str = " → ".join(
            f"[{p['predicate']}] (source: {p['source_id']})"
            for p in row["path"]
        )
        lines.append(f"{row['start']} {path_str} {row['neighbor']}")
    return "\n".join(lines)
 
 
def answer_with_graphrag(question: str, vector_store) -> str:
    """Answer a question using the GraphRAG pipeline."""
    routing = route_query(question)
    route = routing["route"]
    entities = routing["entities"]
 
    graph_context = ""
    vector_context = ""
 
    if route in ("GRAPH_ONLY", "HYBRID") and entities:
        subgraph = graph_traverse(entities, hops=2)
        graph_context = subgraph_to_text(subgraph)
 
    if route in ("VECTOR_ONLY", "HYBRID"):
        chunks = vector_search(question, top_k=8, vector_store=vector_store)
        vector_context = "\n---\n".join(c["metadata"]["text"] for c in chunks)
 
    prompt = f"""Use the context below to answer the question. Always cite the source_id in [square brackets] in your answer.
Do not speculate. If the context does not contain the answer, reply with "Unknown".
 
[Knowledge graph (relationship-based)]
{graph_context or 'none'}
 
[Document chunks (semantic similarity-based)]
{vector_context or 'none'}
 
[Question]
{question}
"""
    response = client.models.generate_content(
        model="gemini-2.5-pro",
        contents=prompt,
        config=types.GenerateContentConfig(temperature=0.2),
    )
    return response.text

The "reply with Unknown" instruction is essential for hallucination control. Gemini 2.5 Pro takes that directive seriously when stated explicitly — leave it out and the model will sometimes fill gaps with plausible-sounding fabrications.

Entity Normalization Done Right — The Quiet Killer of Long-Lived Graphs

If you take only one thing from this article, take this: invest in entity normalization before your graph crosses the 10,000-node mark. The fix is cheap upfront and ruinous to retrofit.

Why does this matter so much? In raw extraction, Gemini will faithfully record whatever surface form appears in the source. So a corpus that mentions "Cloudflare Workers," "CF Workers," "Workers (Cloudflare)," and "cloudflare-workers" will produce four distinct nodes in your graph, all referring to the same product. From a vector-search perspective this is fine — embeddings smooth over surface variation. But for a graph, every alias creates a structural fork: relationships that should converge on one node spread across several, and traversal cannot bridge them.

My production approach is a two-pass strategy. The first pass runs during extraction: I include a normalization instruction in the system prompt and ask Gemini Flash to canonicalize entities against a small registry of known names. The second pass runs as a nightly batch job using Levenshtein distance and embedding similarity to detect new aliases and propose merges, which I review weekly.

# pip install python-Levenshtein
import Levenshtein
 
CANONICAL_REGISTRY = {
    "Gemini API": ["Gemini's API", "the API of Gemini", "gemini-api", "Gemini API service"],
    "Cloudflare Workers": ["CF Workers", "Workers (Cloudflare)", "cloudflare-workers"],
    "Stripe": ["stripe", "Stripe.com", "Stripe Payments"],
}
 
ALIAS_TO_CANONICAL = {
    alias.lower(): canonical
    for canonical, aliases in CANONICAL_REGISTRY.items()
    for alias in [canonical] + aliases
}
 
def normalize_entity(name: str) -> str:
    # Return the canonical form for an entity, falling back to the input.
    key = name.lower().strip()
    if key in ALIAS_TO_CANONICAL:
        return ALIAS_TO_CANONICAL[key]
    # Fuzzy match against known canonicals when no exact alias exists
    for canonical in CANONICAL_REGISTRY:
        if Levenshtein.ratio(key, canonical.lower()) > 0.92:
            return canonical
    return name
 
 
def detect_alias_candidates(graph_entities: list[str]) -> list[tuple[str, str, float]]:
    # Find pairs of entity nodes that may actually be aliases of each other.
    candidates = []
    for i, a in enumerate(graph_entities):
        for b in graph_entities[i + 1:]:
            ratio = Levenshtein.ratio(a.lower(), b.lower())
            if 0.85 < ratio < 1.0:
                candidates.append((a, b, ratio))
    return sorted(candidates, key=lambda x: -x[2])

The cost of this discipline is a few hours of weekly review. The cost of skipping it is a graph that quietly degrades for months and eventually forces a full re-index. I've paid the second cost. You don't have to. As your domain stabilizes, the canonical registry plateaus in size and the weekly review takes minutes — but only because the early investment kept entity drift from compounding.

A second normalization concern is predicate consistency. The Function Calling schema discussed earlier locks predicates to a fixed enum, but Gemini will still occasionally try to bend a relationship into the closest available predicate when the document's actual relationship doesn't fit cleanly. I monitor this by sampling 50 random triples per week and labelling each as "predicate accurate" or "predicate forced." If accuracy drops below 90%, that's a signal to add another predicate to the enum.

Evaluating GraphRAG Quality — Build the Eval Set Before You Need It

A surprising number of teams launch GraphRAG without an evaluation harness, then can't tell whether changes are improving or regressing the system. Don't be that team. Build a small eval set on day one and run it on every meaningful change.

My harness has three components: a labelled question set, an automated judge, and a per-route accuracy breakdown.

The labelled question set is 200 questions hand-written from real user logs and synthetic edge cases. Each question is annotated with the expected answer, the expected route (GRAPH_ONLY, VECTOR_ONLY, or HYBRID), and the source documents that should appear in the citation. I biased the set toward relational questions because that's the failure mode I wanted to fix — about 60% of my questions are relational, 25% are semantic, and 15% are aggregations. The exact mix should reflect your own user traffic; mine is skewed toward the cases that hurt before GraphRAG existed, and that's deliberate.

The automated judge is Gemini 2.5 Pro running an LLM-as-judge prompt that scores each generated answer on three axes: factual correctness, citation accuracy, and completeness. Scores are calibrated against a 30-question subset that I scored manually, and I re-calibrate every quarter. Auto-scoring lets me run the eval on every code change without spending a half-day reading 200 answers. The risk with LLM-as-judge is that the judge develops blind spots, which is why the manual calibration matters.

The per-route accuracy breakdown is where you find what to fix next. If GRAPH_ONLY accuracy is low while VECTOR_ONLY is high, the issue is upstream — extraction quality or normalization. If HYBRID underperforms both individual routes, you have a context-fusion problem in the synthesis prompt. If VECTOR_ONLY regresses after an embedding model change, it's a chunking or retrieval-tuning issue. Each diagnostic points at a different layer of the system, which is exactly what makes GraphRAG worth maintaining: failures localize.

def evaluate_eval_set(eval_set: list[dict], answer_fn) -> dict:
    # Run the eval set and return per-route accuracy.
    results = {"GRAPH_ONLY": [], "VECTOR_ONLY": [], "HYBRID": []}
    for item in eval_set:
        answer = answer_fn(item["question"])
        score = judge_answer(item["question"], item["expected"], answer)
        route = item["expected_route"]
        results[route].append(score)
    return {
        route: {"avg": sum(scores) / len(scores) if scores else 0.0, "n": len(scores)}
        for route, scores in results.items()
    }

Even at 200 questions, this evaluation pipeline pays for itself within a couple of weeks. I have caught regressions from prompt tweaks, embedding-model upgrades, and ontology changes that I would have missed by eyeballing. A nice side effect: when stakeholders ask "is the new model better?" you can answer with numbers, not vibes.

Five Production Pitfalls You'll Almost Certainly Hit

After running this stack for half a year, here are the failure modes that bit me repeatedly.

1. Skip entity normalization and your graph turns into a desert. "Gemini API," "the API of Gemini," and "gemini-api" each become separate nodes. I underestimated this early on, and three months later the graph had ballooned to 100k nodes, queries slowed to a crawl, and I had to re-index. Either run a normalization pass through Gemini Flash during extraction, or post-process with Levenshtein-based clustering. This isn't optional.

2. A loose Function Calling schema produces inconsistent extractions. Letting predicate accept free text means you'll see uses, utilizes, and is_using register as separate predicates. I locked predicates to a seven-value enum: uses, depends_on, created_by, deployed_on, written_in, owned_by, related_to. Start narrow and widen later. A constrained ontology compounds in value over time.

3. Drift between graph and vector store leads to contradictory answers. I once shipped an update path that re-embedded documents but didn't re-extract triples. The graph stayed stale while vectors moved forward, and Gemini received two contradictory views of the same fact in one prompt. Its answer reflected the contradiction. Treat index updates as transactional — both stores update together or neither does, with rollback on failure.

4. Forgetting LIMIT on graph traversal will time out the API. Hot entities (in my domain, "Gemini") can fan out to thousands of neighbors at two hops. Even with Gemini's large context window, dumping that whole expansion in degrades answer quality more than it helps. Always cap with LIMIT and, when possible, filter by entity type.

5. Caching the Flash routing call halves p95 latency. The classification call in Step 3 still costs 200–400 ms on Flash. For business systems that see the same question template repeatedly (FAQ, support, analytics), caching the routing decision keyed on a normalized hash cut my p95 latency by close to half. I store these in Redis. The pattern is covered in more depth in building a Redis-backed semantic cache for the Gemini API in production.

One last operational note: monitor your graph size and edge density weekly, not just total nodes. If edge density per node spikes suddenly, an extraction prompt change probably regressed and is over-extracting noise relationships. If it drops, extraction may have become too conservative. I plot both metrics in a small Grafana dashboard alongside per-route accuracy, and the correlation between graph health and answer quality is tight enough that I treat it as an early-warning signal.

Performance, Cost, and What to Optimize First

Real numbers from my production system (cross-site search across the four Dolice Labs sites, ~12,000 documents):

• Indexing cost: about $14 per 1,000 documents on Gemini 2.5 Pro for triple extraction, plus $0.13 on gemini-embedding-001. Pro extraction dominates.
• Per-query cost: averaging $0.018 (Flash routing + Pro synthesis). Higher than vector-only RAG ($0.011), but accuracy on relational questions improved from 47% to 81%.
• Latency: p50 = 1.8 s, p95 = 3.4 s end-to-end (graph traversal + vector search + Pro synthesis combined).
• Quality: on a 200-question internal eval set, hybrid mode consistently outperformed both GRAPH_ONLY and VECTOR_ONLY.

The single biggest cost optimization: split indexing into "initial batch" and "incremental updates," and use Gemini 2.5 Flash for the incremental side. Flash extracts triples slightly less accurately than Pro, but for small deltas the gap is acceptable. Pay the Pro premium once on the initial 12,000-document load to establish quality, then run cheap incrementals afterward.

If you want to go deeper on the underlying theory, the graph-traversal patterns in this article become much more intuitive after working through a graph theory primer. For implementation, refining the vector side using building a semantic search engine with Gemini API and pgvector gives you a tidy PostgreSQL-centered alternative to running Pinecone.

A practical follow-up worth thinking about: how do you decide whether a question should be HYBRID or one of the single routes? My routing prompt nudges Flash toward HYBRID when the question contains both named entities and adjectives ("the Stripe-using projects deployed to Cloudflare," for example). Pure entity questions ("what does Claude Lab depend on?") map to GRAPH_ONLY, and pure semantic questions ("articles about the experience of solo development") map to VECTOR_ONLY. Edge cases — questions that look entity-heavy but actually need fuzzy matching — are the place where I see the highest classification error rate, around 12%. I don't try to fix this in the classifier; I let HYBRID act as the fallback whenever Flash is unsure. The cost of running both paths and discarding one is small compared to the cost of routing a question wrong and missing the answer.

Another nuance worth flagging: graph traversal scoring. By default, my Cypher query returns up to 50 paths with no inherent ranking. In production this turns out to matter. I now annotate edges with evidence_count (how many source documents support a given relationship) and a confidence score from the extraction step. The traversal then orders results by an aggregate score along the path, which keeps the top-N hits clean even when the subgraph fans out. Without scoring, popular entities flood the context with weakly-supported relationships and Gemini's answer quality erodes.

Closing — The Smallest Useful First Step

GraphRAG is not a silver bullet. Implementation complexity is roughly 2–3× a vanilla vector RAG, and operational load increases proportionally. That said, when relational questions matter to your business, the return on that investment is real.

A concrete next step: pick a small document set (around 100 documents) and run only Step 1 — the triple extraction. Neo4j Desktop is free, and just looking at the extracted triples is enough to develop intuition for whether your domain benefits from GraphRAG. If the output excites you, the remaining steps are about a week of focused work.

For me, the moment GraphRAG started running in production was the moment I stopped reaching for the phrase "vector search isn't enough" in design discussions. If your retrieval system has hit a similar plateau, this architecture is worth the experiment.