Feat: Add Gemini Embedding Support to LightRAG

[danielaskdd]

Feat: 🎯 Add Gemini Embedding Support to LightRAG

Summary

This PR adds comprehensive support for Google Gemini embeddings as a new embedding binding option in LightRAG, enabling users to leverage Gemini's gemini-embedding-001 model with advanced features including dynamic dimension control and automatic normalization.

🚀 Features

1. Gemini Embedding Function (lightrag/llm/gemini.py)

• ✅ Default dimension: 1536 - Optimal balance between performance and quality
• ✅ Dynamic dimension support - Range: 128-3072 (via EMBEDDING_DIM environment variable)
• ✅ Automatic L2 normalization - For dimensions < 3072 to ensure optimal semantic similarity
• ✅ Task type optimization - Support for 8 different task types (RETRIEVAL_DOCUMENT, RETRIEVAL_QUERY, SEMANTIC_SIMILARITY, etc.)
• ✅ Retry mechanism - Exponential backoff with 3 attempts for robust error handling
• ✅ Token tracking - Full integration with LightRAG's token usage monitoring
• ✅ Async execution - Thread pool executor for optimal performance

2. Configuration Options (lightrag/llm/binding_options.py)

• Added GeminiEmbeddingOptions class with task_type parameter
• Follows existing binding options pattern for consistency
• Automatic CLI and environment variable generation

3. Server Integration (lightrag/api/lightrag_server.py)

• Added Gemini to supported embedding bindings list
• Implemented LLMConfigCache support for Gemini embedding options
• Created optimized embedding function with configuration caching

4. CLI/Environment Variable Support (lightrag/api/config.py)

• Added "gemini" to --embedding-binding choices
• Automatic parameter exposure via CLI: --gemini-embedding-task-type
• Environment variable support: GEMINI_EMBEDDING_TASK_TYPE

📝 Usage Examples

Environment Variables

# Basic configuration
export EMBEDDING_BINDING=gemini
export EMBEDDING_MODEL=gemini-embedding-001
export GEMINI_API_KEY=your_api_key_here
export EMBEDDING_DIM=1536

# Optional: Task type optimization
export GEMINI_EMBEDDING_TASK_TYPE=RETRIEVAL_QUERY

Command Line

python -m lightrag.api.lightrag_server \
  --embedding-binding gemini \
  --embedding-model gemini-embedding-001 \
  --embedding-dim 1536 \
  --gemini-embedding-task-type RETRIEVAL_DOCUMENT

🔧 Technical Details

Normalization Strategy

Gemini API behavior:

• Dimension 3072: Already normalized by API
• Dimensions < 3072: Require L2 normalization for optimal cosine similarity

Implementation automatically handles this:

if embedding_dim and embedding_dim < 3072:
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    embeddings = embeddings / norms

Task Type Options

Task Type	Description	Use Case
RETRIEVAL_DOCUMENT	Optimize for document indexing	RAG document storage
RETRIEVAL_QUERY	Optimize for search queries	RAG query processing
SEMANTIC_SIMILARITY	Text similarity comparison	Duplicate detection
CLASSIFICATION	Text classification tasks	Sentiment analysis
CLUSTERING	Group similar texts	Document organization
CODE_RETRIEVAL_QUERY	Code search optimization	Code suggestions
QUESTION_ANSWERING	Q&A system optimization	Chatbot applications
FACT_VERIFICATION	Fact-checking optimization	Automated verification

🧪 Testing Checklist

• Embedding function with different dimensions (768, 1536, 3072)
• Normalization verification for dimensions < 3072
• Task type variations
• CLI parameter parsing
• Environment variable configuration
• Server integration
• Configuration caching

📄 Files Changed

• lightrag/llm/gemini.py - Added gemini_embed() function (+129 lines)
• lightrag/llm/binding_options.py - Added GeminiEmbeddingOptions class (+13 lines)
• lightrag/api/lightrag_server.py - Added Gemini embedding integration (+37 lines)
• lightrag/api/config.py - Exposed Gemini embedding CLI/ENV options (+22 lines)

🔄 Migration Guide

For existing users, no breaking changes. New feature is opt-in:

✅ Checklist

• Code follows project style guidelines
• Documentation added/updated
• All imports properly organized
• Error handling implemented
• Logging added for debugging
• Follows existing binding patterns
• No breaking changes introduced

[danielaskdd]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

LightRAG/lightrag/api/lightrag_server.py

Lines 751 to 768 in de4ed73

	# Get embedding_send_dim from centralized configuration
	embedding_send_dim = args.embedding_send_dim
	# Check if the function signature has embedding_dim parameter
	# Note: Since optimized_embedding_func is an async function, inspect its signature
	sig = inspect.signature(optimized_embedding_func)
	has_embedding_dim_param = "embedding_dim" in sig.parameters
	# Determine send_dimensions value based on binding type
	# Jina REQUIRES dimension parameter (forced to True)
	# OpenAI and others: controlled by EMBEDDING_SEND_DIM environment variable
	if args.embedding_binding == "jina":
	# Jina API requires dimension parameter - always send it
	send_dimensions = has_embedding_dim_param
	dimension_control = "forced by Jina API"
	else:
	# For OpenAI and other bindings, respect EMBEDDING_SEND_DIM setting
	send_dimensions = embedding_send_dim and has_embedding_dim_param

Gemini embeddings ignore configured dimensions

The new gemini_embed implementation only applies output_dimensionality and L2 normalisation when the embedding_dim argument is non‑None (gemini.py lines 501‑539), but the server never forwards a value for that parameter unless the global EMBEDDING_SEND_DIM flag is manually enabled. In create_app the send_dimensions flag remains False for all non‑Jina bindings by default (lightrag_server.py lines 751‑768), so the optimized embedding function is invoked with embedding_dim=None, which means the Gemini API is always called with its default dimension and the normalisation branch never runs. As a result, neither the advertised default of 1536 dimensions nor any user‑supplied EMBEDDING_DIM or --gemini-embedding-task-type settings have any effect out of the box, yielding unnormalised embeddings of whatever size the API chooses. This makes the dynamic dimension and normalisation feature effectively non‑functional for Gemini unless users know to set an unrelated EMBEDDING_SEND_DIM override.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[danielaskdd]

Treat Gemini the same as Jina, requiring embedding dimension parameter.

[danielaskdd]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Keep them coming!

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".