Generative AI Google Brain

gemini-use-cases-graphrag-graph-rag-spanner-sdk-adk

4 min read

GraphRAG on Google Cloud With Spanner and Vertex AI Agent Engine

Source notebook

Repo path: gemini/use-cases/graphrag/graph_rag_spanner_sdk_adk.ipynb · Open on GitHub · advanced

Builds a GraphRAG Q&A agent with Spanner Graph, Gemini, ADK, and Vertex AI Agent Engine.

Summary

This notebook teaches how to create a knowledge graph from Larry Page text using Gemini through LangChain, store it in Spanner Graph, and query it with a SpannerGraphQAChain. It then adds embeddings and Spanner vector search to rewrite entity mentions such as Larry Page to Lawrence Edward Page before graph querying. Finally, it wraps the workflow as an ADK agent with Google Search fallback and deploys it to Vertex AI Agent Engine.

Key code patterns

Spanner graph schema

request = CreateDatabaseRequest(
    create_statement=f"CREATE DATABASE `{database_id}`",
    extra_statements=["""CREATE TABLE KgNode (
        DocId INT64 NOT NULL,
        Name STRING(1024),
        DOC STRING(1024),
        DocEmbedding ARRAY<FLOAT64>
    ) PRIMARY KEY (DocId)"""],
)

Creates a Spanner table for text chunks, extracted entity names, and vector embeddings.

Graph extraction with Gemini

llm = VertexAI(model_name=MODEL_NAME, project=GCP_PROJECT_ID, location=REGION)
llm_transformer_filtered = LLMGraphTransformer(
    llm=llm,
    allowed_nodes=["Person", "Country", "Organization", "Asset"],
    allowed_relationships=["NATIONALITY", "LOCATED_IN", "WORKED_AT", "SPOUSE", "NET_WORTH", "INVESTMENT", "INFLUENCED_BY"],
)
graph_documents_filtered = llm_transformer_filtered.convert_to_graph_documents(documents)

Converts source text into constrained graph nodes and relationships.

Store graph in Spanner

graph_store = SpannerGraphStore(
    instance_id=SPANNER_INSTANCE_ID,
    database_id=SPANNER_DATABASE_ID,
    graph_name=SPANNER_GRAPH_NAME,
)
graph_store.cleanup()
for graph_document in graph_documents_filtered:
    graph_store.add_graph_documents([graph_document])

Persists LangChain graph documents into Spanner Graph.

Graph Q&A chain

llm = ChatVertexAI(model=MODEL_NAME, temperature=0)
chain = SpannerGraphQAChain.from_llm(
    llm,
    graph=graph_store,
    allow_dangerous_requests=True,
    return_intermediate_steps=True,
)
response = chain.invoke("query=" + question)

Lets Gemini answer questions by generating and running graph queries over Spanner.

Vector-assisted entity rewrite

q_emb = get_embedding(QUESTION, TASK_TYPE, embedding_model)
results = snapshot.execute_sql(
    """SELECT DocId, NAME, Doc FROM KgNode
       ORDER BY COSINE_DISTANCE(DocEmbedding, @q_emb) limit 1""",
    params={"q_emb": q_emb},
)
response = text_model.generate_content(user_prompt_content, generation_config=GenerationConfig(response_mime_type="application/json"))

Uses nearest embedded chunk names to rewrite query entities before graph lookup.

ADK agent with tools

search_agent = Agent(model="gemini-2.0-flash-001", name="SearchAgent", tools=[google_search])
root_agent = Agent(
    name="graph_rag_agent",
    model="gemini-2.0-flash-001",
    tools=[agent_tool.AgentTool(agent=search_agent), ask_graph],
)

Combines graph retrieval and Google Search as callable ADK tools.

Deploy to Agent Engine

app = reasoning_engines.AdkApp(agent=root_agent, enable_tracing=True)
remote_app = agent_engines.create(
    app,
    requirements=["google-cloud-aiplatform==1.91.0", "google-adk==0.5.0", "google-cloud-spanner==3.48.0"],
)

Packages the ADK agent for managed deployment on Vertex AI Agent Engine.

Models & APIs used

  • Models: gemini-2.0-flash-001, text-embedding-004
  • APIs / services: Vertex AI, Cloud Spanner API, Cloud Spanner, Vertex AI Agent Engine, Google Search, Cloud Trace
  • SDKs / libraries: google-cloud-spanner, langchain-core, langchain-google-vertexai, langchain-experimental, langchain-community, langchain-text-splitters, langchain-google-spanner, google-adk, google-cloud-aiplatform, vertexai, google-genai, networkx, json-repair, pydantic, sklearn

When to use this

Use this pattern when a Q&A agent needs graph-structured retrieval, semantic entity matching, and managed deployment.

Gotchas & caveats

  • GCP_PROJECT_ID is left empty and must be set before running gcloud, Spanner, Vertex AI, and ADK code.
  • Cloud Spanner API must be enabled and billing must be active.
  • The tutorial asks for Owner role to complete setup.
  • Spanner instance is created in regional-us-central1 and REGION is set to us-central1.
  • graph_store.cleanup() can remove existing graph data from the database.
  • The %%spanner_graph magic cell only works on Colab.
  • ask_graph redeclares GCP_PROJECT_ID as an empty string, so it must be filled for the tool to work.
  • Agent Engine service account needs roles/spanner.databaseUser on the Spanner database.
  • The staging bucket is assumed to be gs://{GCP_PROJECT_ID}-vertexai-staging.
  • allow_dangerous_requests=True is required by the demonstrated SpannerGraphQAChain setup.

Best practices

  • Constrain graph extraction with allowed_nodes and allowed_relationships.
  • Use temperature=0 for deterministic graph Q&A and entity rewriting.
  • Store embeddings alongside graph-related text in Spanner for semantic matching.
  • Combine vector search with graph search when exact graph entity names may not match user phrasing.
  • Configure ADK instructions to check the graph database before broader Google Search.
  • Ask the user before doing broader search when the graph lacks an answer.
  • Enable tracing on the Agent Engine app for troubleshooting and monitoring.
  • Pin package versions in notebook installs and Agent Engine requirements.

Graph View

Backlinks