Generative AI Google Brain

open-models-serving-vertex-ai-ollama-gemma2-rag-agent

4 min read

Running a Gemma 2-based agentic RAG with Ollama on Vertex AI and LangGraph

Source notebook

Repo path: open-models/serving/vertex_ai_ollama_gemma2_rag_agent.ipynb · Open on GitHub · advanced

Deploys a Gemma 2 Ollama container on Vertex AI and uses it in a LangGraph SQL RAG agent.

Summary

The notebook teaches how to package a Gemma 2 LoRA SQL adapter as an Ollama model inside a FastAPI custom serving container, build it with Cloud Build, register it, and deploy it to a Vertex AI endpoint. It then wraps the deployed endpoint for LangGraph, builds a FAISS-backed RAG workflow over BigQuery SQL documentation, and rewrites non-SQL-formatted answers into BigQuery SQL.

Key code patterns

Initialize Vertex AI with staging bucket

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
BUCKET_URI = f"gs://{PROJECT_ID}-bucket"
! gsutil mb -p $PROJECT_ID -l $LOCATION $BUCKET_URI
vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)

Sets project, region, and Cloud Storage staging before using Vertex AI SDK operations.

Download gated Gemma adapter

base_model_id = "google-cloud-partnership/gemma-2-2b-it-lora-sql"
model_dir = MODELS_DIR / "gemma-2-2b-it-lora-sql"
snapshot_download(
    repo_id=base_model_id,
    token=get_token(),
    local_dir=model_dir,
    ignore_patterns=[".gitattributes", ".gitkeep", "*.md"],
)

Fetches the Hugging Face model artifacts needed for the Ollama adapter container.

Create Ollama Modelfile

modelfile = """FROM gemma2:2b
ADAPTER ollama_models/gemma-2-2b-it-lora-sql
"""
with BUILD_DIR.joinpath("gemma-2-2b-it-lora-sql.modelfile").open("w") as f:
    f.write(modelfile)

Combines the Gemma 2 base model with the downloaded SQL LoRA adapter for Ollama serving.

FastAPI Vertex prediction proxy

@app.post(Config.PREDICT_ROUTE, response_model=PredictionResponse)
async def predict(request: PredictionRequest) -> PredictionResponse:
    tasks = []
    for instance in request.instances:
        prompt = instance.get("inputs", "")
        parameters = instance.get("parameters", {})
        tasks.append(ollama_generate(prompt, parameters))
    return PredictionResponse(predictions=await asyncio.gather(*tasks))

Adapts Vertex AI prediction payloads to Ollama generate requests and returns Vertex-compatible predictions.

Upload and deploy custom container

model = Model.upload(
    display_name="google--gemma-2-2b-it-lora-sql-ollama",
    serving_container_image_uri=SERVING_CONTAINER_IMAGE_URI,
    serving_container_ports=[8080],
)
endpoint = Endpoint.create(display_name="google--gemma-2-2b-it-lora-sql-ollama-endpoint")
deployed_model = model.deploy(endpoint=endpoint, machine_type="g2-standard-4", accelerator_type="NVIDIA_L4", accelerator_count=1)

Registers the serving container in Vertex AI Model Registry and deploys it to a GPU endpoint.

Wrap Vertex endpoint for LangGraph

class CustomVertexAIModel:
    def __init__(self, project, location, endpoint_id, **model_params):
        self.endpoint = aiplatform.Endpoint(
            endpoint_name=f"projects/{project}/locations/{location}/endpoints/{endpoint_id}"
        )
        self.model_params = model_params
    def invoke(self, prompt, **kwargs):
        response = self.endpoint.predict([{"inputs": prompt, "parameters": {**self.model_params, **kwargs}}])
        return response.predictions[0]

Creates a minimal interface so a Vertex AI endpoint can be called from the LangGraph workflow.

Build FAISS RAG index from docs

loader = WebBaseLoader(urls)
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(documents)
vectorstore = FAISS.from_documents(splits, embeddings)

Turns BigQuery documentation URLs into a searchable local vector store for retrieval.

Conditional LangGraph rewrite loop

workflow.add_edge("retrieve", "generate")
workflow.add_edge("generate", "should_rewrite")
workflow.add_conditional_edges(
    "should_rewrite", lambda x: x["next"], {"rewrite": "rewrite", "end": END}
)
workflow.add_edge("rewrite", "retrieve")
workflow.set_entry_point("retrieve")

Routes generated answers through a SQL-format check and loops through rewrite when needed.

Models & APIs used

  • Models: google-cloud-partnership/gemma-2-2b-it-lora-sql, gemma2:2b, gemma-2-2b-it-lora-sql-2b, text-embedding-005
  • APIs / services: Vertex AI, Cloud Storage, Artifact Registry, Cloud Build, Compute Engine, Hugging Face Hub
  • SDKs / libraries: google-cloud-aiplatform, vertexai, huggingface_hub, torch, etils, FastAPI, httpx, langchain-community, langchainhub, langchain_google_vertexai, langgraph, faiss-gpu

When to use this

Use this pattern when you need to serve a Hugging Face Gemma 2 Ollama model on Vertex AI and consume it inside a LangGraph RAG workflow.

Gotchas & caveats

  • Requires IAM roles for Artifact Registry, Cloud Build, Vertex AI, service accounts, Service Usage, and Storage.
  • Requires Artifact Registry API, Vertex AI API, and Compute Engine API to be enabled.
  • Hugging Face access is required because the Gemma model is gated and needs accepted usage terms plus a read token.
  • The notebook restarts the kernel after package installation.
  • LocalModel testing requires Docker installed and running locally.
  • Vertex AI deployment can take around 15 to 25 minutes.
  • The custom proxy maps only the Ollama generate API.
  • The deployed endpoint uses g2-standard-4 with one NVIDIA_L4 accelerator.

Best practices

  • Use a Cloud Storage staging bucket when initializing Vertex AI SDK.
  • Use Artifact Registry and Cloud Build to build and store the custom serving image.
  • Expose Vertex AI-compatible health and predict routes in the custom container.
  • Validate prediction requests and return consistent error responses in the FastAPI proxy.
  • Test the serving container locally with Vertex AI LocalModel before deploying when debugging.
  • Use retrieved BigQuery documentation context before generating SQL answers.
  • Use conditional graph routing to rewrite SQL-related answers that are not properly formatted.
  • Clean up endpoints, models, Artifact Registry repositories, and tutorial files when no longer needed.

Graph View

Backlinks