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:
parent
bf176b38ee
commit
519f7f61c4
2 changed files with 14 additions and 3 deletions
|
|
@ -226,7 +226,18 @@ class BaseVectorStorage(StorageNameSpace, ABC):
|
||||||
Returns:
|
Returns:
|
||||||
str: Suffix string, e.g. "text_embedding_3_large_3072d"
|
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:
|
def _get_legacy_collection_name(self) -> str:
|
||||||
"""Get legacy collection/table name (without suffix).
|
"""Get legacy collection/table name (without suffix).
|
||||||
|
|
|
||||||
|
|
@ -207,13 +207,13 @@ class UnifiedLock(Generic[T]):
|
||||||
self._lock.release()
|
self._lock.release()
|
||||||
else:
|
else:
|
||||||
self._lock.release()
|
self._lock.release()
|
||||||
|
|
||||||
direct_log(
|
direct_log(
|
||||||
f"== Lock == Process {self._pid}: Released lock {self._name} (async={self._is_async})",
|
f"== Lock == Process {self._pid}: Released lock {self._name} (async={self._is_async})",
|
||||||
level="INFO",
|
level="INFO",
|
||||||
enable_output=self._enable_logging,
|
enable_output=self._enable_logging,
|
||||||
)
|
)
|
||||||
main_lock_released = True
|
main_lock_released = True
|
||||||
|
|
||||||
# Then release async lock if in multiprocess mode
|
# Then release async lock if in multiprocess mode
|
||||||
if not self._is_async and self._async_lock is not None:
|
if not self._is_async and self._async_lock is not None:
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue