diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 8a638759..b6bd7f4a 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -5,6 +5,7 @@ import asyncio import configparser import inspect import os +import json import time import warnings from dataclasses import asdict, dataclass, field @@ -117,6 +118,7 @@ from dotenv import load_dotenv # use the .env that is inside the current folder # allows to use different .env file for each lightrag instance + # the OS environment variables take precedence over the .env file load_dotenv(dotenv_path=".env", override=False) @@ -525,7 +527,7 @@ class LightRAG: logger.debug(f"LightRAG init with param:\n {_print_config}\n") # Init Embedding - # Step 1: Capture max_token_size before applying decorator (decorator strips dataclass attributes) + # Step 1: Capture max_token_size and embedding_dim before applying decorator embedding_max_token_size = None if self.embedding_func and hasattr(self.embedding_func, "max_token_size"): embedding_max_token_size = self.embedding_func.max_token_size @@ -534,6 +536,17 @@ class LightRAG: ) self.embedding_token_limit = embedding_max_token_size + # Capture embedding model name before decoration so we don't lose it to wrappers + self.embedding_model_name = ( + self.embedding_func.__class__.__name__ if self.embedding_func else "unknown" + ) + + # --- CAPTURE EMBEDDING DIMENSION (NEW) --- + self.embedding_dim = None + if self.embedding_func and hasattr(self.embedding_func, "embedding_dim"): + self.embedding_dim = self.embedding_func.embedding_dim + # ----------------------------------------- + # Step 2: Apply priority wrapper decorator self.embedding_func = priority_limit_async_func_call( self.embedding_func_max_async, @@ -658,8 +671,56 @@ class LightRAG: self._storages_status = StoragesStatus.CREATED + def _check_embedding_config(self): + """ + Validates that the current embedding dimension matches the existing data. + Prevents data corruption when switching models without clearing storage. + """ + if self.embedding_dim is None: + # If we couldn't capture dimensions, skip the check to avoid blocking valid custom models + return + + meta_file = os.path.join(self.working_dir, "lightrag_meta.json") + + if os.path.exists(meta_file): + with open(meta_file, "r", encoding="utf-8") as f: + try: + meta_data = json.load(f) + saved_dim = meta_data.get("embedding_dim") + saved_model = meta_data.get("embedding_model_func", "unknown") + + if saved_dim and saved_dim != self.embedding_dim: + raise ValueError( + f"Embedding dimension mismatch! " + f"Existing data uses dimension {saved_dim} (Model: {saved_model}), " + f"but current configuration uses {self.embedding_dim}. " + f"Please clear the '{self.working_dir}' directory or switch back to the original model." + ) + except json.JSONDecodeError: + logger.warning( + f"Could not parse {meta_file}. Skipping dimension check." + ) + else: + # First run: Save the metadata + meta_data = { + "embedding_dim": self.embedding_dim, + "embedding_model_func": self.embedding_model_name, + "created_at": datetime.now(timezone.utc).isoformat(), + } + + # Ensure directory exists + if not os.path.exists(self.working_dir): + os.makedirs(self.working_dir) + + with open(meta_file, "w", encoding="utf-8") as f: + json.dump(meta_data, f, indent=4) + async def initialize_storages(self): """Storage initialization must be called one by one to prevent deadlock""" + # --- NEW SAFETY CHECK CALL --- + self._check_embedding_config() + # ----------------------------- + if self._storages_status == StoragesStatus.CREATED: # Set the first initialized workspace will set the default workspace # Allows namespace operation without specifying workspace for backward compatibility