Generative AI Google Brain

gemini-use-cases-code-code-retrieval-augmented-generation

3 min read

Use Retrieval Augmented Generation (RAG) with Gemini API

Source notebook

Repo path: gemini/use-cases/code/code_retrieval_augmented_generation.ipynb · Open on GitHub · intermediate

Builds a LangChain RAG pipeline over GitHub code notebooks using Gemini and Vertex AI embeddings.

Summary

This notebook teaches how to augment Gemini API responses with external repository knowledge using a code RAG pattern. It crawls the GoogleCloudPlatform/generative-ai GitHub repository, extracts notebook code, chunks it, embeds it with Vertex AI embeddings, stores it in FAISS, and compares zero-shot Gemini output with RetrievalQA output using retrieved code context.

Key code patterns

Initialize Gemini LLM

vertexai.init(project=PROJECT_ID, location=LOCATION)
code_llm = VertexAI(
    model_name="gemini-2.0-flash",
    max_output_tokens=2048,
    temperature=0.1,
)

Sets up Vertex AI and a deterministic Gemini model wrapper for code generation.

Crawl GitHub code files

response = requests.get(api_url, headers=headers)
response.raise_for_status()
for item in response.json():
    if item["type"] == "file" and item["name"].endswith((".py", ".ipynb")):
        files.append(item["html_url"])
    elif item["type"] == "dir" and not item["name"].startswith("."):
        files.extend(crawl_github_repo(item["url"], True))

Collects repository source files to use as the external knowledge corpus.

Extract notebook code

raw_url = github_url.replace("github.com", "raw.githubusercontent.com").replace("/blob/", "/")
notebook = nbformat.reads(requests.get(raw_url).text, as_version=nbformat.NO_CONVERT)
for cell in notebook.cells:
    if cell.cell_type == "code":
        python_code = cell.source if not python_code else python_code + "\n" + cell.source

Turns notebook code cells into plain text documents for retrieval.

Rate-limited embeddings

class CustomVertexAIEmbeddings(VertexAIEmbeddings):
    def embed_documents(self, texts, batch_size=None):
        limiter = rate_limit(self.requests_per_minute)
        while docs:
            head, docs = docs[:self.num_instances_per_batch], docs[self.num_instances_per_batch:]
            results.extend(self.client.get_embeddings(head))
            next(limiter)
        return [r.values for r in results]

Respects the notebook’s stated batch limit of five documents per embedding request.

Build FAISS retriever

text_splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.PYTHON, chunk_size=2000, chunk_overlap=200
)
texts = text_splitter.split_documents(code_strings)
embeddings = CustomVertexAIEmbeddings(model_name="text-embedding-005", requests_per_minute=100, num_instances_per_batch=5)
db = FAISS.from_documents(texts, embeddings)
retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 5})

Creates a local vector index and retrieves the top five similar code chunks.

Create RetrievalQA chain

qa_chain = RetrievalQA.from_llm(
    llm=code_llm,
    prompt=prompt_RAG_template,
    retriever=retriever,
    return_source_documents=True,
)
results = qa_chain.invoke(input={"query": user_question})

Combines Gemini, a RAG prompt, and retrieved source documents for grounded code generation.

Models & APIs used

  • Models: gemini-2.0-flash, text-embedding-005
  • APIs / services: Vertex AI, Gemini API in Vertex AI, Embeddings for Text API
  • SDKs / libraries: google-cloud-aiplatform, vertexai, langchain, langchain_google_vertexai, langchain-community, faiss-cpu, nbformat, requests

When to use this

Use this pattern when generating code should be grounded in examples from an existing repository or codebase.

Gotchas & caveats

  • Requires an existing Google Cloud project with the Vertex AI API enabled.
  • Colab users must authenticate and restart the runtime after package installation.
  • Requires a GitHub personal access token with repo public_repo scope to crawl the repository.
  • The notebook uses local FAISS for the demo and recommends managed vector stores for production.
  • Embedding requests are batched because the API accepts a maximum of five documents per request in this notebook.

Best practices

  • Use retrieved repository context rather than relying only on zero-shot prompting for code generation.
  • Split Python code with RecursiveCharacterTextSplitter.from_language using chunk overlap.
  • Persist crawled GitHub file URLs to avoid repeatedly downloading the file list.
  • Use low temperature for concise code generation.
  • Return source documents from RetrievalQA for traceability.
  • Use managed vector stores such as Vertex AI Vector Search, AlloyDB for PostgreSQL, or Cloud SQL for PostgreSQL with pgvector for production.

Graph View

Backlinks