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)
This commit is contained in:
BukeLy 2025-11-20 00:51:47 +08:00
parent bf176b38ee
commit 519f7f61c4
2 changed files with 14 additions and 3 deletions

View file

@ -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).

View file

@ -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: