# How to Implement Custom Embedding Functions in LightRAG Using wrap_embedding_func_with_attrs

> Learn to implement custom embedding functions in LightRAG using wrap_embedding_func_with_attrs. Seamlessly integrate your models into indexing and retrieval pipelines with this powerful decorator.

- Repository: [✨Data Intelligence Lab@HKU✨/LightRAG](https://github.com/HKUDS/LightRAG)
- Tags: how-to-guide
- Published: 2026-03-23

---

**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`](https://github.com/HKUDS/LightRAG/blob/main/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`](https://github.com/HKUDS/LightRAG/blob/main/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`](https://github.com/HKUDS/LightRAG/blob/main/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`](https://github.com/HKUDS/LightRAG/blob/main/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`](https://github.com/HKUDS/LightRAG/blob/main/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.

```python
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.

```python
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.

```python
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.

```python
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`](https://github.com/HKUDS/LightRAG/blob/main/lightrag/utils.py)**: Defines `EmbeddingFunc` and the `wrap_embedding_func_with_attrs` decorator (core of custom embedding support).
- **[`lightrag/llm/openai.py`](https://github.com/HKUDS/LightRAG/blob/main/lightrag/llm/openai.py)**: Example provider-specific embed function wrapped with the decorator, demonstrating token-limit handling (lines 75–81).
- **[`lightrag/llm/ollama.py`](https://github.com/HKUDS/LightRAG/blob/main/lightrag/llm/ollama.py)**: Another provider example that demonstrates proper use of `.func` when partially binding arguments.
- **[`lightrag/api/config.py`](https://github.com/HKUDS/LightRAG/blob/main/lightrag/api/config.py)**: Contains workspace management logic that uses the `model_name` attribute for vector isolation.
- **[`examples/lightrag_ollama_demo.py`](https://github.com/HKUDS/LightRAG/blob/main/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`](https://github.com/HKUDS/LightRAG/blob/main/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`](https://github.com/HKUDS/LightRAG/blob/main/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.