How to Implement Custom Embedding Functions in LightRAG Using wrap_embedding_func_with_attrs

Use the @wrap_embedding_func_with_attrs decorator to convert any async embedding coroutine into an EmbeddingFunc instance that automatically injects embedding_dim and max_token_size attributes for seamless integration with LightRAG's indexing and retrieval pipelines.

HKUDS/LightRAG abstracts embedding providers behind async-compatible functions to keep the framework model-agnostic. When you need to integrate a custom embedding model—whether it is a local transformer, a third-party API, or a specialized encoder—you must wrap it using the wrap_embedding_func_with_attrs utility found in lightrag/utils.py. This decorator ensures your function carries the metadata required for dimension validation and token management while maintaining compatibility with the rest of the pipeline.

Understanding the Decorator and EmbeddingFunc

The wrap_embedding_func_with_attrs decorator serves as the bridge between raw async embedding logic and LightRAG’s internal pipelines. When applied to an async function, it returns an EmbeddingFunc instance (defined in lightrag/utils.py lines 411–447) that stores the original coroutine in the .func attribute and overloads __call__ to automatically inject configuration parameters.

The decorator attaches two critical attributes to the resulting callable:

• embedding_dim: The expected vector size, used for runtime dimension validation.
• max_token_size: An optional token limit that enables automatic truncation before sending text to the model.
• model_name: An optional identifier that enables workspace isolation when multiple models share the same vector store.

LightRAG’s internal components—such as vector-store loaders and workspace updaters—invoke the EmbeddingFunc instance directly. This means the injected embedding_dim and max_token_size values are supplied to your function automatically without requiring manual parameter plumbing in every call.

Avoiding Double-Wrapping with Partial Functions

When binding extra arguments using functools.partial, you must reference the underlying raw function via the .func attribute to prevent double-wrapping. The EmbeddingFunc.__post_init__ logic (lines 50–58 in lightrag/utils.py) automatically unwraps nested EmbeddingFunc objects, but calling the wrapper a second time would still cause the outer decorator to re-inject parameters and produce warnings.

Always access .func when using partial to pre-configure parameters like model name or host URL. This ensures you are extending the original coroutine rather than wrapping the already-wrapped EmbeddingFunc instance.

Internal Validation and Pipeline Integration

The EmbeddingFunc.__call__ method performs several runtime checks to maintain pipeline integrity. After the underlying provider returns a NumPy array, the wrapper validates that every embedding vector matches the declared embedding_dim. A mismatch raises a ValueError, guaranteeing consistency across the vector store.

If max_token_size is set, the wrapper automatically truncates any input text that exceeds the limit before forwarding it to the provider. This logic follows the pattern found in lightrag/llm/openai.py (lines 75–81), preventing API-level "token limit exceeded" errors without requiring manual preprocessing in your embedding function.

The optional model_name attribute enables LightRAG to keep vectors from different models separate when they share the same vector store, a feature utilized by the workspace manager in lightrag/api/config.py.

Implementation Examples

Basic Local Model Wrapper

Define a minimal custom embedder that returns vectors compatible with LightRAG’s expected signature. The decorator injects embedding_dim=768 automatically.

from lightrag.utils import wrap_embedding_func_with_attrs
import numpy as np

@wrap_embedding_func_with_attrs(embedding_dim=768, max_token_size=0)
async def local_embed(texts: list[str], embedding_dim: int | None = None) -> np.ndarray:
    """
    Dummy example that returns a zero-vector of the requested dimension.
    Replace the body with a real model inference call.
    """
    # Real model inference would go here.

    return np.zeros((len(texts), embedding_dim or 768), dtype=np.float32)

Safe Partial Configuration

When using functools.partial to bind parameters like embed_model or host, reference the raw coroutine via .func to avoid double-decoration.

from functools import partial
from lightrag.llm.ollama import ollama_embed  # already wrapped

# ❌ Wrong – would wrap the already-wrapped function again

# embed_fn = partial(ollama_embed, embed_model="bge-m3", host="http://localhost:11434")

# ✅ Correct – access the raw coroutine via .func

embed_fn = partial(
    ollama_embed.func,
    embed_model="bge-m3",
    host="http://localhost:11434",
)

Advanced Token Truncation

Implement preprocessing logic that respects the max_token_size injected by the decorator, as demonstrated in the OpenAI-compatible provider implementation.

from lightrag.utils import wrap_embedding_func_with_attrs
import tiktoken   # any tokenizer you prefer

import numpy as np

@wrap_embedding_func_with_attrs(
    embedding_dim=1024,
    max_token_size=2048,
    model_name="my_local_encoder"
)
async def my_encoder_embed(
    texts: list[str],
    embedding_dim: int | None = None,
    max_token_size: int | None = None,
) -> np.ndarray:
    # Truncate texts that exceed the token budget

    if max_token_size:
        enc = tiktoken.get_encoding("cl100k_base")
        texts = [
            t if len(enc.encode(t)) <= max_token_size else
            enc.decode(enc.encode(t)[:max_token_size])
            for t in texts
        ]

    # Replace with actual encoder inference

    # Example: vectors = my_encoder.encode(texts)

    vectors = np.random.rand(len(texts), embedding_dim or 1024).astype(np.float32)
    return vectors

Integration with LightRAG

Pass the decorated EmbeddingFunc instance directly to the LightRAG constructor. The framework invokes your custom embedder automatically during document ingestion and query processing.

from lightrag import LightRAG
from my_custom_embed import my_encoder_embed   # the EmbeddingFunc instance

rag = LightRAG(
    embed_func=my_encoder_embed,   # accepts the EmbeddingFunc directly

    # ... other config ...

)

# Now any call that needs embeddings (e.g., rag.add_documents) will use the custom embedder.

Key Source Files

• lightrag/utils.py: Defines EmbeddingFunc and the wrap_embedding_func_with_attrs decorator (core of custom embedding support).
• lightrag/llm/openai.py: Example provider-specific embed function wrapped with the decorator, demonstrating token-limit handling (lines 75–81).
• lightrag/llm/ollama.py: Another provider example that demonstrates proper use of .func when partially binding arguments.
• lightrag/api/config.py: Contains workspace management logic that uses the model_name attribute for vector isolation.
• examples/lightrag_ollama_demo.py: User-level demo showing how to bind parameters with partial and access .func to avoid double wrapping.

Summary

• Decorate async functions with @wrap_embedding_func_with_attrs(embedding_dim=..., max_token_size=...) to create an EmbeddingFunc instance compatible with LightRAG.
• Access .func when using functools.partial to prevent double-wrapping and parameter re-injection.
• Rely on automatic injection of embedding_dim and max_token_size into your function signature for validation and truncation logic.
• Validate dimensions automatically—the wrapper raises ValueError if returned vectors do not match the declared size.
• Support workspace isolation by setting the optional model_name attribute, allowing multiple models to share vector stores without collision.

Frequently Asked Questions

What is the purpose of wrap_embedding_func_with_attrs in LightRAG?

The decorator converts a raw async embedding function into an EmbeddingFunc object that carries metadata required by the framework, including vector dimensions and token limits. It ensures that custom embedding providers integrate seamlessly with LightRAG’s validation and storage pipelines without requiring changes to internal logic.

How do I avoid double-wrapping when using functools.partial?

Always reference the underlying coroutine via the .func attribute of the EmbeddingFunc instance when creating a partial. For example, use partial(my_embedder.func, ...) rather than partial(my_embedder, ...). This prevents the decorator from wrapping an already-wrapped function, which would cause parameter injection conflicts.

Can I implement custom token truncation logic in my embedding function?

Yes. When you declare max_token_size in the decorator and accept it as a parameter in your function signature, LightRAG automatically passes the value during invocation. You can use this to truncate text with a tokenizer (such as tiktoken) before sending it to your model, following the pattern used in lightrag/llm/openai.py.

Where is the EmbeddingFunc class defined?

The EmbeddingFunc dataclass and the wrap_embedding_func_with_attrs decorator are defined in lightrag/utils.py (lines 411–447). This file also contains the __post_init__ logic (lines 50–58) that handles unwrapping of nested embedding functions and the __call__ method that performs runtime dimension validation.

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