Building RAG Applications with GKE and Cloud SQL

This article explores a reference architecture for deploying Retrieval Augmented Generation (RAG) applications on Google Cloud, leveraging Google Kubernetes Engine (GKE), Cloud SQL, and popular open-source tools. RAG enhances the quality of Generative AI outputs by grounding them in retrieved knowledge, making it ideal for applications requiring accurate and context-aware responses. This guide is tailored for developers familiar with GKE and Cloud SQL, and who possess a conceptual understanding of AI, Machine Learning (ML), and Large Language Models (LLMs). We'll delve into the architecture's components, data flow, and key considerations for building a robust and efficient RAG system.

The architecture comprises two primary subsystems: the embedding subsystem and the service subsystem. The embedding subsystem is responsible for ingesting data from various sources, transforming it into vector embeddings, and storing these embeddings in a vector database. The service subsystem handles user requests, retrieves relevant information from the vector database using semantic search, and generates responses using an LLM. This separation of concerns allows for efficient data processing and scalable service delivery.

Data from both internal and external sources is uploaded to Cloud Storage. This upload triggers an event that notifies the embedding service. The embedding service then retrieves the data, preprocesses it using Ray Data (which may involve chunking and formatting), and generates vector embeddings using open-source models like intfloat/multilingual-e5-small. These embeddings are then written to a Cloud SQL for PostgreSQL vector database, which is optimized for storing and retrieving high-dimensional vectors.

A user submits a natural language request through a web-based chat interface. The front-end server, running on GKE, uses LangChain to convert the request into an embedding. This embedding is used to perform a semantic search in the vector database, retrieving relevant data. The original request is then combined with the retrieved data to create a contextualized prompt, which is sent to the inference server. The inference server, powered by Hugging Face TGI, uses an open-source LLM (e.g., Mistral-7B-Instruct or Gemma) to generate a response. The response is filtered for safety using Responsible AI (RAI) services before being sent back to the user.

This architecture leverages several key Google Cloud and open-source products. Google Kubernetes Engine (GKE) provides the container orchestration platform. Cloud Storage offers scalable object storage. Cloud SQL for PostgreSQL, enhanced with the pgvector extension, serves as the vector database. Open-source tools include Hugging Face Text Generation Inference (TGI) for LLM serving, Ray for distributed computing, and LangChain for building LLM-powered applications.

RAG is applicable to various scenarios. For personalized product recommendations, a chatbot can leverage historical user data to provide more relevant suggestions. In clinical assistance, doctors can use RAG to access patient history and external knowledge bases for improved diagnoses. In legal research, lawyers can quickly query regulations and case law, enhanced by data from internal contracts and case records.

For a fully managed vector search solution, consider using Vertex AI and Vector Search. Alternatively, you can leverage the vector storage capabilities of other Google Cloud databases like AlloyDB for PostgreSQL. These alternatives offer different trade-offs in terms of management overhead and performance.

Security is paramount. Utilize GKE Autopilot's built-in security features. Implement Identity-Aware Proxy (IAP) for access control. Encrypt data at rest and in transit using Cloud KMS. For Cloud SQL, enforce secure connections using SSL/TLS or the Cloud SQL Auth proxy. Use Sensitive Data Protection to identify and de-identify sensitive data in Cloud Storage. Leverage VPC Service Controls to prevent data exfiltration. Ensure compliance with data residency requirements by specifying the appropriate region for data storage.

Ensure high availability by using regional GKE clusters and configuring Cloud SQL instances with high availability. Utilize Cloud Storage's regional or multi-regional storage options for data redundancy. Consider using reserved capacity for GPU resources to ensure availability during autoscaling events.

Optimize costs by leveraging GKE Autopilot's efficiency. Consider Committed Use Discounts for predictable workloads. Use Spot VMs for GKE nodes to reduce compute costs. For Cloud SQL, use standalone instances if high availability is not required. Utilize Cloud SQL's cost analysis insights to identify over-provisioned instances. Choose the appropriate Cloud Storage class based on data access frequency and retention requirements.

Select the appropriate compute category for GKE pods based on performance requirements. Use GPU machine types for inference servers and embedding services. Optimize Cloud SQL performance by allocating sufficient CPU and memory. Use IVFFlat or HNSW indexes for faster approximate nearest neighbor (ANN) vector search. Utilize Cloud SQL's Query Insights tool to identify and resolve performance bottlenecks. For large file uploads to Cloud Storage, consider using parallel composite uploads.

A sample codebase is available on GitHub for deploying this architecture. This code is intended for experimentation and not production use. It provisions Cloud SQL, deploys Ray, JupyterHub, and Hugging Face TGI to GKE, and deploys a sample chatbot application. Remember to remove any unused resources after experimentation. Explore further by reviewing GKE best practices, investigating Google Cloud's Generative AI grounding options, and studying architectures using Vertex AI and Vector Search or AlloyDB. Consult the Well-Architected Framework for AI and Machine Learning for architectural principles and recommendations.