feat: Add task-aware embedding support

docs/DockerDeployment.md

@@ -48,6 +48,8 @@ LightRAG can be configured using environment variables in the `.env` file:
 - `EMBEDDING_BINDING`: Embedding backend (lollms/ollama/openai)
 - `EMBEDDING_BINDING_HOST`: Embedding server host URL
 - `EMBEDDING_MODEL`: Embedding model name
+- `EMBEDDING_DOCUMENT_PREFIX`: Optional prefix for document embeddings (e.g., "search_document: ")
+- `EMBEDDING_QUERY_PREFIX`: Optional prefix for query embeddings (e.g., "search_query: ")
 
 **RAG Configuration**
 

env.example

@@ -408,6 +408,15 @@ OLLAMA_EMBEDDING_NUM_CTX=8192
 # # EMBEDDING_DIM=2048
 # # EMBEDDING_BINDING_API_KEY=your_api_key
 
+### Optional: asymmetric embeddings (query/document behavior split)
+### For Jina embeddings set the flag, default is false for backward compatibility
+# EMBEDDING_ASYMMETRIC=false
+### For OpenAI-compatible APIs, you can also set prefixes to enable asymmetric
+### behavior without explicitly setting the toggle.
+### Wrap with quotes if there are trailing spaces
+# EMBEDDING_DOCUMENT_PREFIX="search_document: "
+# EMBEDDING_QUERY_PREFIX="search_query: "
+
 ####################################################################
 ### WORKSPACE sets workspace name for all storage types
 ### for the purpose of isolating data from LightRAG instances.

examples/unofficial-sample/lightrag_embedding_prefixes.py

@@ -0,0 +1,234 @@
+import os
+import asyncio
+import inspect
+import logging
+import logging.config
+from functools import partial
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import openai_complete_if_cache
+from lightrag.llm.ollama import ollama_embed
+from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
+
+from dotenv import load_dotenv
+
+load_dotenv(dotenv_path=".env", override=False)
+
+WORKING_DIR = "./dickens"
+
+
+def configure_logging():
+    """Configure logging for the application"""
+
+    # Reset any existing handlers to ensure clean configuration
+    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
+        logger_instance = logging.getLogger(logger_name)
+        logger_instance.handlers = []
+        logger_instance.filters = []
+
+    # Get log directory path from environment variable or use current directory
+    log_dir = os.getenv("LOG_DIR", os.getcwd())
+    log_file_path = os.path.abspath(
+        os.path.join(log_dir, "lightrag_compatible_demo.log")
+    )
+
+    print(f"\nLightRAG compatible demo log file: {log_file_path}\n")
+    os.makedirs(os.path.dirname(log_dir), exist_ok=True)
+
+    # Get log file max size and backup count from environment variables
+    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))  # Default 10MB
+    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))  # Default 5 backups
+
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "default": {
+                    "format": "%(levelname)s: %(message)s",
+                },
+                "detailed": {
+                    "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+                },
+            },
+            "handlers": {
+                "console": {
+                    "formatter": "default",
+                    "class": "logging.StreamHandler",
+                    "stream": "ext://sys.stderr",
+                },
+                "file": {
+                    "formatter": "detailed",
+                    "class": "logging.handlers.RotatingFileHandler",
+                    "filename": log_file_path,
+                    "maxBytes": log_max_bytes,
+                    "backupCount": log_backup_count,
+                    "encoding": "utf-8",
+                },
+            },
+            "loggers": {
+                "lightrag": {
+                    "handlers": ["console", "file"],
+                    "level": "INFO",
+                    "propagate": False,
+                },
+            },
+        }
+    )
+
+    # Set the logger level to INFO
+    logger.setLevel(logging.INFO)
+    # Enable verbose debug if needed
+    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
+
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+
+async def llm_model_func(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        os.getenv("LLM_MODEL", "deepseek-chat"),
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=os.getenv("LLM_BINDING_API_KEY") or os.getenv("OPENAI_API_KEY"),
+        base_url=os.getenv("LLM_BINDING_HOST", "https://api.deepseek.com"),
+        **kwargs,
+    )
+
+
+async def print_stream(stream):
+    async for chunk in stream:
+        if chunk:
+            print(chunk, end="", flush=True)
+
+
+async def initialize_rag():
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        # Note: ollama_embed is decorated with @wrap_embedding_func_with_attrs,
+        # which wraps it in an EmbeddingFunc. Using .func accesses the original
+        # unwrapped function to avoid double wrapping when we create our own
+        # EmbeddingFunc with custom configuration (embedding_dim, max_token_size and prefixes).
+        embedding_func=EmbeddingFunc(
+            embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
+            max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "8192")),
+            supports_asymmetric=True,
+            func=partial(
+                ollama_embed.func,  # Access the unwrapped function to avoid double EmbeddingFunc wrapping
+                embed_model=os.getenv("EMBEDDING_MODEL", "FRIDA:latest"),
+                host=os.getenv("EMBEDDING_BINDING_HOST", "http://localhost:11434"),
+                query_prefix=os.getenv("EMBEDDING_QUERY_PREFIX", "search_query: "),
+                document_prefix=os.getenv(
+                    "EMBEDDING_DOCUMENT_PREFIX", "search_document: "
+                ),
+            ),
+        ),
+    )
+
+    await rag.initialize_storages()  # Auto-initializes pipeline_status
+    return rag
+
+
+async def main():
+    try:
+        # Clear old data files
+        files_to_delete = [
+            "graph_chunk_entity_relation.graphml",
+            "kv_store_doc_status.json",
+            "kv_store_full_docs.json",
+            "kv_store_text_chunks.json",
+            "vdb_chunks.json",
+            "vdb_entities.json",
+            "vdb_relationships.json",
+        ]
+
+        for file in files_to_delete:
+            file_path = os.path.join(WORKING_DIR, file)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                print(f"Deleting old file:: {file_path}")
+
+        # Initialize RAG instance
+        rag = await initialize_rag()
+
+        # Test embedding function
+        test_text = ["This is a test string for embedding."]
+        embedding = await rag.embedding_func(test_text)
+        embedding_dim = embedding.shape[1]
+        print("\n=======================")
+        print("Test embedding function")
+        print("========================")
+        print(f"Test dict: {test_text}")
+        print(f"Detected embedding dimension: {embedding_dim}\n\n")
+
+        with open("./book.txt", "r", encoding="utf-8") as f:
+            await rag.ainsert(f.read())
+
+        # Perform naive search
+        print("\n=====================")
+        print("Query mode: naive")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="naive", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform local search
+        print("\n=====================")
+        print("Query mode: local")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="local", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform global search
+        print("\n=====================")
+        print("Query mode: global")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="global", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+        # Perform hybrid search
+        print("\n=====================")
+        print("Query mode: hybrid")
+        print("=====================")
+        resp = await rag.aquery(
+            "What are the top themes in this story?",
+            param=QueryParam(mode="hybrid", stream=True),
+        )
+        if inspect.isasyncgen(resp):
+            await print_stream(resp)
+        else:
+            print(resp)
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    finally:
+        if rag:
+            await rag.finalize_storages()
+
+
+if __name__ == "__main__":
+    # Configure logging before running the main function
+    configure_logging()
+    asyncio.run(main())
+    print("\nDone!")

lightrag/api/config.py

@@ -459,6 +459,16 @@ def parse_args() -> argparse.Namespace:
         "MAX_UPLOAD_SIZE", 104857600, int, special_none=True
     )
 
+    # Embedding prefix configuration for context-aware embeddings
+    args.embedding_document_prefix = get_env_value(
+        "EMBEDDING_DOCUMENT_PREFIX", None, special_none=True
+    )
+    args.embedding_query_prefix = get_env_value(
+        "EMBEDDING_QUERY_PREFIX", None, special_none=True
+    )
+    # Asymmetric embedding behavior toggle
+    args.embedding_asymmetric = get_env_value("EMBEDDING_ASYMMETRIC", False, bool)
+
     ollama_server_infos.LIGHTRAG_NAME = args.simulated_model_name
     ollama_server_infos.LIGHTRAG_TAG = args.simulated_model_tag