From 519f7f61c4f4c273b7017250e1453c953fd343bc Mon Sep 17 00:00:00 2001 From: BukeLy Date: Thu, 20 Nov 2025 00:51:47 +0800 Subject: [PATCH] fix: handle wrapped embedding_func and lock flag logic Why these changes are needed: 1. LightRAG wraps embedding_func with priority_limit_async_func_call decorator, causing loss of get_model_identifier method 2. UnifiedLock.__aexit__ set main_lock_released flag incorrectly How it solves them: 1. _generate_collection_suffix now tries multiple approaches: - First check if embedding_func has get_model_identifier - Fallback to original EmbeddingFunc in global_config - Return empty string for backward compatibility 2. Move main_lock_released = True inside the if block so flag is only set when lock actually exists and is released Impact: - Fixes E2E tests that initialize complete LightRAG instances - Fixes incorrect async lock cleanup in exception scenarios - Maintains backward compatibility Testing: All unit tests pass (test_qdrant_migration.py, test_postgres_migration.py) --- lightrag/base.py | 13 ++++++++++++- lightrag/kg/shared_storage.py | 4 ++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/lightrag/base.py b/lightrag/base.py index 1997c231..9671f1b7 100644 --- a/lightrag/base.py +++ b/lightrag/base.py @@ -226,7 +226,18 @@ class BaseVectorStorage(StorageNameSpace, ABC): Returns: str: Suffix string, e.g. "text_embedding_3_large_3072d" """ - return self.embedding_func.get_model_identifier() + # Try to get model identifier from the embedding function + # If it's a wrapped function (doesn't have get_model_identifier), + # fallback to the original embedding_func from global_config + if hasattr(self.embedding_func, 'get_model_identifier'): + return self.embedding_func.get_model_identifier() + elif 'embedding_func' in self.global_config: + original_embedding_func = self.global_config['embedding_func'] + if hasattr(original_embedding_func, 'get_model_identifier'): + return original_embedding_func.get_model_identifier() + + # Fallback: no model identifier available + return "" def _get_legacy_collection_name(self) -> str: """Get legacy collection/table name (without suffix). diff --git a/lightrag/kg/shared_storage.py b/lightrag/kg/shared_storage.py index 6a3fa8c7..d56e0be0 100644 --- a/lightrag/kg/shared_storage.py +++ b/lightrag/kg/shared_storage.py @@ -207,13 +207,13 @@ class UnifiedLock(Generic[T]): self._lock.release() else: self._lock.release() - + direct_log( f"== Lock == Process {self._pid}: Released lock {self._name} (async={self._is_async})", level="INFO", enable_output=self._enable_logging, ) - main_lock_released = True + main_lock_released = True # Then release async lock if in multiprocess mode if not self._is_async and self._async_lock is not None: