How to Use Google Embeddings with code-review-graph for Semantic Search

Configure the GOOGLE_API_KEY environment variable and initialize EmbeddingStore with provider="google" to leverage Gemini embeddings for semantic code search backed by SQLite vector storage.

The code-review-graph library enables semantic search over codebases by embedding graph nodes into vector space. When you use Google embeddings with code-review-graph, the library instantiates a GoogleEmbeddingProvider that communicates with the Gemini Embedding API via the google-generativeai SDK, storing vectors locally for efficient cosine-similarity queries.

Prerequisites and Configuration

Before generating vectors, you must satisfy two configuration requirements.

API Key Setup

The provider factory get_provider() in code_review_graph/embeddings.py inspects the environment for GOOGLE_API_KEY. If present and non-empty, the function returns an instance of GoogleEmbeddingProvider when provider="google" is requested (lines 1000‑1021). Export this key in your shell or set it programmatically:

import os
os.environ["GOOGLE_API_KEY"] = "YOUR_GEMINI_API_KEY"

Safety Acknowledgment

Because Google embeddings transmit code snippets to a cloud API, the library emits a one-time stderr warning unless you explicitly acknowledge the risk. Set CRG_ACCEPT_CLOUD_EMBEDDINGS=1 to suppress the warning and proceed with remote embedding generation (lines 443‑452).

export CRG_ACCEPT_CLOUD_EMBEDDINGS=1

Provider Architecture and Initialization

The embedding abstraction is defined by the base class EmbeddingProvider. The concrete GoogleEmbeddingProvider implementation initializes the Gemini client in __init__, accepting an optional model name that defaults to gemini-embedding-001 (lines 182‑190).

When you create an EmbeddingStore with provider="google", the constructor invokes get_provider() to obtain the Google implementation and records the provider name in the SQLite embeddings table, isolating Google vectors from other backends (lines 66‑71, 1368‑1380).

from code_review_graph.embeddings import EmbeddingStore

store = EmbeddingStore(
    db_path="code_vectors.db",
    provider="google",
    model="gemini-embedding-001"  # optional; default is gemini-embedding-001

)

Embedding Generation and Storage

Batch Processing and Retry Logic

The GoogleEmbeddingProvider.embed() method batches up to 100 texts per request, invokes models.embed_content, and extracts the values field from each Embedding result (lines 195‑207). To handle transient rate limits or server errors, the provider wraps API calls in _call_with_retry, which implements exponential backoff with jitter (lines 212‑242).

For query-time vectorization, embed_query() performs a single-document request using the "RETRIEVAL_QUERY" task type, optimizing the embedding for search scenarios (lines 249‑259).

Persistent Storage

The EmbeddingStore persists vectors in a local SQLite database. Each row in the embeddings table stores the vector data alongside a provider column, ensuring that searches later filter to the correct embedding backend. When you call embed_nodes(), the store automatically routes text through the selected provider and writes the resulting vectors to disk.

Performing Semantic Search

The EmbeddingStore.search() method executes semantic retrieval by:

1. Embedding the input query via self.provider.embed_query(query).
2. Computing cosine similarity between the query vector and all stored vectors that share the same provider value.
3. Returning the top-k qualified names ordered by similarity score (lines 1470‑1484).

This design guarantees that a Google-embedded query is only compared against Google-embedded codebase vectors, preventing cross-provider mismatches.

Complete Implementation Example

The following Python script demonstrates end-to-end configuration, indexing, and search:

import os
from code_review_graph.embeddings import EmbeddingStore, get_provider

# 1. Configure credentials and safety acknowledgment

os.environ["GOOGLE_API_KEY"] = "YOUR_GEMINI_API_KEY"
os.environ["CRG_ACCEPT_CLOUD_EMBEDDINGS"] = "1"

# 2. Initialize store with Google provider

store = EmbeddingStore(
    db_path="my_project_embeddings.db",
    provider="google",
    model="gemini-embedding-001"
)

# 3. Index graph nodes (assumes list_of_graph_nodes exists)

# store.embed_nodes(list_of_graph_nodes)

# 4. Execute semantic search

results = store.search("how to handle async file parsing", limit=5)
for qualified_name, score in results:
    print(f"{qualified_name}: {score:.4f}")

For command-line usage, export the same environment variables and invoke the CLI:

export GOOGLE_API_KEY="YOUR_GEMINI_API_KEY"
export CRG_ACCEPT_CLOUD_EMBEDDINGS=1

crg embed --provider google --model gemini-embedding-001
crg search "async file processing" --limit 10

Summary

• Configuration: Export GOOGLE_API_KEY and CRG_ACCEPT_CLOUD_EMBEDDINGS=1 to enable Google Gemini embeddings.
• Provider Selection: get_provider() returns GoogleEmbeddingProvider when the Google key is present and requested (lines 1000‑1021).
• Vectorization: The provider batches requests (max 100), applies exponential-backoff retries, and uses gemini-embedding-001 by default (lines 195‑242).
• Storage: EmbeddingStore writes vectors to SQLite with a provider column to segregate backends (lines 1368‑1380).
• Search: Cosine similarity is computed against same-provider vectors using embed_query() with the RETRIEVAL_QUERY task type (lines 1470‑1484).

Frequently Asked Questions

What environment variables are required to use Google embeddings?

You must set GOOGLE_API_KEY to your Gemini API key. Optionally, set CRG_ACCEPT_CLOUD_EMBEDDINGS=1 to suppress the security warning printed to stderr before the first API call (lines 443‑452).

How does code-review-graph handle rate limits or API failures?

The GoogleEmbeddingProvider implements _call_with_retry, which catches transient errors and retries requests with exponential backoff and jitter (lines 212‑242). This protects your indexing workflow from temporary Gemini API outages or quota exceeded errors.

Can I mix Google embeddings with other providers in the same database?

Yes. The embeddings SQLite table includes a provider column that tags each vector with its source (e.g., "google" or "openai"). The search() method automatically filters to vectors matching the current store’s provider, preventing cross-model similarity comparisons (lines 1470‑1484).

What is the difference between embed() and embed_query()?

embed() is optimized for batch document indexing and processes up to 100 texts per request (lines 195‑207). embed_query() is designed for single-query vectorization and passes the "RETRIEVAL_QUERY" task type to the Gemini API, tuning the embedding for question-answering retrieval (lines 249‑259).

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}