Feat: Add Optional Embedding Dimension Control with OpenAI API

[danielaskdd]

Add Optional Embedding Dimension Control with OpenAI API

🎯 Overview

This PR implements proper dimension parameter handling for OpenAI and Jina embedding APIs, with mandatory dimension sending for Jina (API requirement) and optional control for OpenAI.

📋 Changes Summary

1. Embedding Function Updates

OpenAI (lightrag/llm/openai.py)

• Added optional embedding_dim parameter to openai_embed()
• Parameter is conditionally passed to API based on EMBEDDING_SEND_DIM configuration
• Enables API-level dimension reduction when needed
• Maintains backward compatibility when parameter is not provided

Jina (lightrag/llm/jina.py)

• Renamed dimensions parameter to embedding_dim for consistency
• CRITICAL: Jina API requires dimension parameter - always enabled
• Aligned parameter naming with OpenAI implementation
• Updated documentation to reflect auto-injection behavior

2. Smart Dimension Injection (lightrag/utils.py)

Enhanced EmbeddingFunc class with intelligent dimension management:

class EmbeddingFunc:
    embedding_dim: int
    func: callable
    max_token_size: int | None = None
    send_dimensions: bool = False  # NEW: Control dimension parameter sending
    
    async def __call__(self, *args, **kwargs) -> np.ndarray:
        if self.send_dimensions:
            # Auto-inject embedding_dim with validation
            kwargs['embedding_dim'] = self.embedding_dim
        return await self.func(*args, **kwargs)

Features:

• Automatic dimension parameter injection when enabled
• Parameter validation with warning system
• Detects and warns about user-provided conflicting values
• Enforces dimension consistency from decorator configuration

3. Provider-Specific Logic (lightrag/api/lightrag_server.py)

Implemented binding-aware dimension parameter control:

if args.embedding_binding == 'jina':
    # Jina API requires dimension parameter - always send it
    send_dimensions = has_embedding_dim_param
    dimension_control = "forced (Jina API requirement)"
else:
    # For OpenAI and other bindings, respect EMBEDDING_SEND_DIM setting
    send_dimensions = embedding_send_dim and has_embedding_dim_param
    dimension_control = f"env_var={embedding_send_dim}"

Logic Flow:

• Jina: send_dimensions forced to True (API requirement)
• OpenAI: send_dimensions controlled by EMBEDDING_SEND_DIM env var
• Others: send_dimensions set to False (not supported)

4. Environment Configuration (env.example)

Updated documentation with clear guidance:

### Control whether to send embedding_dim parameter to embedding API
### IMPORTANT: Jina ALWAYS sends dimension parameter (API requirement) - this setting is ignored for Jina
### For OpenAI: Set to 'true' to enable dynamic dimension adjustment
### For OpenAI: Set to 'false' (default) to disable sending dimension parameter
### Note: Automatically ignored for backends that don't support dimension parameter (e.g., Ollama)
# EMBEDDING_SEND_DIM=false

5. Prohibit direct access to internal functions of EmbeddingFunc in query stage

Update embedding in query stage to use embedding_func_config([query]) instead of embedding_func_config.func([query]), ensuring the EmbeddingFunc wrapper properly applies dimension configuration from EMBEDDING_SEND_DIM=true.

Fixed Locations:

• _perform_kg_search function__ (line ~1510): Pre-computing query embedding for all vector operations
• _find_related_text_unit_from_entities function__ (line ~1960): Entity-related chunk selection by vector similarity
• _find_related_text_unit_from_relations function__ (line ~2180): Relation-related chunk selection by vector similarity

🔧 Technical Details

Dimension Sending Decision Matrix

Backend	EMBEDDING_SEND_DIM	send_dimensions	Reason
Jina	true	✅ True	Required by API
Jina	false	✅ True	Required by API
OpenAI	true	✅ True	User enabled
OpenAI	false	❌ False	User disabled (default)
Ollama	true	❌ False	Not supported
Ollama	false	❌ False	Not supported

Parameter Injection Flow

graph TD
    A[Check Binding Type] --> B{Jina?}
    B -->|Yes| C[Force send_dimensions=True]
    B -->|No| D[Check EMBEDDING_SEND_DIM]
    D --> E{Environment Var?}
    E -->|true| F[Check Function Signature]
    E -->|false| G[send_dimensions=False]
    F --> H{Has embedding_dim param?}
    H -->|Yes| I[send_dimensions=True]
    H -->|No| G
    C --> J[Auto-inject dimension on each call]
    I --> J
    G --> K[Use function without dimension]

Loading

Validation Logic

When send_dimensions=True:

1. Check if user manually provided embedding_dim parameter

2. If provided value differs from class attribute, log warning:

   WARNING: Ignoring user-provided embedding_dim=X, using declared embedding_dim=Y
   

3. Always inject the decorator-declared embedding_dim value

🎁 Benefits

1. API Compliance: Jina embedding requests always include required dimension parameter
2. Flexibility: Optional dimension control for OpenAI without breaking existing code
3. Safety: Automatic validation prevents dimension mismatches
4. Compatibility: Seamless operation with backends that don't support dimensions
5. Debugging: Comprehensive logging for troubleshooting configuration issues
6. Optimization: Enables API-level dimension reduction for supported providers

⚙️ Configuration Guide

For Jina (Automatic - No Configuration Needed)

EMBEDDING_BINDING=jina
EMBEDDING_DIM=1024
# Dimension parameter ALWAYS sent to Jina API (required)
# EMBEDDING_SEND_DIM setting is ignored

For OpenAI (Optional Control)

EMBEDDING_BINDING=openai
EMBEDDING_DIM=1536

# Enable dimension parameter sending (optional)
EMBEDDING_SEND_DIM=true

# OR disable (default behavior)
# EMBEDDING_SEND_DIM=false

For Other Backends (No Impact)

EMBEDDING_BINDING=ollama
EMBEDDING_DIM=1024
# EMBEDDING_SEND_DIM automatically ignored (not supported)

🧪 Backward Compatibility

✅ Fully Backward Compatible:

• Existing code works without changes
• Default behavior unchanged (EMBEDDING_SEND_DIM=false)
• Non-supporting backends automatically skip dimension parameter
• Manual dimension passing still works with validation warnings
• Jina users experience no breaking changes (dimension always sent as required)

🔍 Testing Recommendations

Required Tests

1. ✅ Jina with EMBEDDING_SEND_DIM=false - Dimension still sent (forced)
2. ✅ Jina with EMBEDDING_SEND_DIM=true - Dimension sent (forced)
3. ✅ OpenAI with EMBEDDING_SEND_DIM=true - Dimension sent
4. ✅ OpenAI with EMBEDDING_SEND_DIM=false - Dimension not sent (default)
5. ✅ Ollama with any config - Dimension parameter ignored
6. ✅ Validation warnings - Test manual embedding_dim parameter conflict

Test Scenarios

# Test 1: Jina always sends dimension
EMBEDDING_BINDING=jina
EMBEDDING_DIM=1024
EMBEDDING_SEND_DIM=false  # Should be ignored - dimension still sent

# Test 2: OpenAI with dimension enabled
EMBEDDING_BINDING=openai
EMBEDDING_DIM=1536
EMBEDDING_SEND_DIM=true   # Dimension sent to API

# Test 3: OpenAI default behavior
EMBEDDING_BINDING=openai
EMBEDDING_DIM=1536
# EMBEDDING_SEND_DIM not set (defaults to false) - dimension not sent

📊 Logging Examples

Jina Binding (Forced Dimension)

INFO: Embedding configuration: send_dimensions=True (forced (Jina API requirement), has_param=True, binding=jina)

OpenAI with EMBEDDING_SEND_DIM=true

INFO: Embedding configuration: send_dimensions=True (env_var=True, has_param=True, binding=openai)

OpenAI with EMBEDDING_SEND_DIM=false (default)

INFO: Embedding configuration: send_dimensions=False (env_var=False, has_param=True, binding=openai)

🚀 Migration Path

For Existing Deployments

No action required - the changes are backward compatible:

• Default behavior remains unchanged
• Existing configurations continue to work
• Jina users automatically get compliant behavior

For New Deployments

OpenAI users can optionally enable dimension reduction:

# Add to .env file
EMBEDDING_SEND_DIM=true

Jina users require no configuration - dimension parameter automatically sent.

🔗 Related Information

API Requirements

• Jina Embeddings API: Requires dimensions parameter for proper operation
• OpenAI Embeddings API: Supports optional dimensions parameter for dimension reduction
• Other APIs: May not support dimension parameter (automatically skipped)

Performance Impact

• Minimal - parameter injection happens once during initialization
• No overhead for backends that don't use the parameter
• Potential optimization for OpenAI (API-level dimension reduction)

✅ Checklist

• Code implements Jina mandatory dimension requirement
• Code maintains backward compatibility
• Documentation updated (env.example)
• Logging provides clear configuration status
• Validation prevents dimension mismatches
• Default behavior unchanged
• All binding types handled correctly

[danielaskdd]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ 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]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ 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".