Generative AI Google Brain

gemini-use-cases-retrieval-augmented-generation-rag-embeddings-and-index-with-alloydb

3 min read

Retrieval Augmented Generation(RAG) with AlloyDB

Source notebook

Repo path: gemini/use-cases/retrieval-augmented-generation/rag_embeddings_and_index_with_alloydb.ipynb · Open on GitHub · intermediate

Builds a RAG workflow over patent abstracts using AlloyDB vector search and Gemini on Vertex AI.

Summary

This notebook teaches how to use AlloyDB as a RAG backend by loading public Google Patents data from BigQuery, storing text fields in AlloyDB, generating embeddings with AlloyDB AI, and indexing them with pgvector IVFFlat. It retrieves semantically similar patent rows for a user query and passes that context into Gemini to generate an answer with supplemental patent information and URLs.

Key code patterns

Initialize Vertex AI

PROJECT_ID = "<your-project-id>"
LOCATION = "us-central1"
 
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets the project and region before using Vertex AI models.

Connect to AlloyDB

connector = Connector()
conn = connector.connect(
    instance_uri=inst_uri,
    driver="pg8000",
    user=user,
    password=password,
    db=db,
    ip_type="PUBLIC",
)
engine = sqlalchemy.create_engine("postgresql+pg8000://", creator=getconn)

Uses the AlloyDB connector with SQLAlchemy pooling to run PostgreSQL queries.

Generate embeddings in AlloyDB

ALTER TABLE google_patents_research
ADD COLUMN embedding vector(768);
 
UPDATE google_patents_research
SET embedding = embedding('text-embedding-005', title || ' ' || abstract);

Creates a vector column and fills it using AlloyDB’s Vertex AI embedding integration.

Create vector index

CREATE INDEX ON google_patents_research
USING ivfflat (embedding vector_cosine_ops);

Adds an IVFFlat pgvector index for approximate nearest neighbor retrieval.

Retrieve similar rows

SELECT publication_number, title, abstract, url
FROM google_patents_research
ORDER BY embedding <-> embedding('text-embedding-005', query)::vector
LIMIT 5;

Ranks patent records by vector similarity to the embedded user query.

Generate grounded answer

response = GenerativeModel(GENERATIVE_MODEL).generate_content(
    contents=input_prompt,
    generation_config=GenerationConfig(temperature=0.6, max_output_tokens=1024),
)
return response.text

Feeds retrieved patent context into Gemini to answer the user’s question.

Models & APIs used

  • Models: text-embedding-005, gemini-2.0-flash
  • APIs / services: Vertex AI, AlloyDB, BigQuery
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-alloydb-connector, vertexai, SQLAlchemy, pg8000, pandas

When to use this

Use this pattern when building a RAG application that stores enterprise text data and embeddings in AlloyDB for PostgreSQL-compatible vector search.

Gotchas & caveats

  • The AlloyDB cluster must reside in us-central1 to generate embeddings with text-embedding-005.
  • Vertex AI and AlloyDB APIs must be enabled before running the workflow.
  • Private services access must be enabled with Google Cloud Platform as the service provider.
  • The AlloyDB service agent needs the roles/aiplatform.user IAM role.
  • The notebook uses public IP for AlloyDB and warns that public IP assignment can take a few minutes.
  • Colab requires package installation followed by runtime restart and authentication.
  • The notebook pins google-cloud-aiplatform1.46.0, google-cloud-alloydb-connector1.0.0, SQLAlchemy2.0.29, and pg80001.31.1.
  • The helper function builds SQL with f-strings, including table names and query text.

Best practices

  • Pin package versions for reproducible notebook setup.
  • Initialize Vertex AI with explicit project and location.
  • Use the AlloyDB connector and SQLAlchemy engine instead of raw unmanaged connections.
  • Create the google_ml_integration and vector extensions before generating and querying embeddings.
  • Concatenate title and abstract before embedding to improve retrieval context.
  • Use an IVFFlat index with cosine distance for high-dimensional embedding search.
  • Pass retrieved rows into a prompt template before calling Gemini.
  • Clean up created AlloyDB instances and clusters after the demo.

Graph View

Backlinks