git-sync/openwebui-memory-system
1
0
Fork 0
mirror of https://github.com/mtayfur/openwebui-memory-system.git synced 2026-01-22 06:51:01 +01:00
Code Issues Packages Projects Releases Wiki Activity

refactor(memory_system): remove excessive try/except and input validation, streamline async operations, and add skip state cache

Browse Source 
Removes redundant try/except blocks and input validation in several
methods to simplify logic and improve readability. Moves error handling
to higher levels where appropriate. Adds a skip state cache to track
when memory operations should be skipped, improving efficiency by
avoiding repeated skip checks. Cleans up batch operation execution and
cache clearing to include the new skip state. These changes reduce
unnecessary code complexity and improve maintainability, while also
optimizing memory operation flow and cache management.
This commit is contained in:
mtayfur
2025-10-28 17:51:09 +03:00
parent 8ced9aace5
commit 8997f71f05
1 changed files with 118 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
memory_system.py
View File

@@ -254,6 +254,7 @@ class UnifiedCacheManager:
self.EMBEDDING_CACHE = "embedding"
self.RETRIEVAL_CACHE = "retrieval"
self.MEMORY_CACHE = "memory"
self.SKIP_STATE_CACHE = "skip"
async def get(self, user_id: str, cache_type: str, key: str) -> Optional[Any]:
"""Get value from cache with LRU updates."""
@@ -747,7 +748,6 @@ class LLMConsolidationService:
if not existing_memories:
return None
try:
content_embedding = await self.memory_system._generate_embeddings(content, user_id)
for memory in existing_memories:
@@ -755,7 +755,6 @@ class LLMConsolidationService:
continue
memory_embedding = await self.memory_system._generate_embeddings(memory.content, user_id)
similarity = float(np.dot(content_embedding, memory_embedding))
if similarity >= Constants.DEDUPLICATION_SIMILARITY_THRESHOLD:
@@ -763,9 +762,6 @@ class LLMConsolidationService:
return str(memory.id)
return None
except Exception as e:
logger.warning(f"⚠️ Semantic duplicate check failed: {str(e)}")
return None
def _filter_consolidation_candidates(self, similarities: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], str]:
"""Filter consolidation candidates by threshold and return candidates with threshold info."""
@@ -945,21 +941,13 @@ class LLMConsolidationService:
async def execute_memory_operations(self, operations: List[Dict[str, Any]], user_id: str, emitter: Optional[Callable] = None) -> Tuple[int, int, int, int]:
"""Execute consolidation operations with simplified tracking."""
if not operations or not user_id:
if not operations:
return 0, 0, 0, 0
try:
user = await asyncio.wait_for(
asyncio.to_thread(Users.get_user_by_id, user_id),
timeout=Constants.DATABASE_OPERATION_TIMEOUT_SEC,
)
except asyncio.TimeoutError:
raise TimeoutError(f"⏱️ User lookup timed out after {Constants.DATABASE_OPERATION_TIMEOUT_SEC}s")
except Exception as e:
raise RuntimeError(f"👤 User lookup failed: {str(e)}")
if not user:
raise ValueError(f"👤 User not found for consolidation: {user_id}")
created_count = updated_count = deleted_count = failed_count = 0
@@ -982,14 +970,10 @@ class LLMConsolidationService:
memory_contents_for_deletion = {}
if operations_by_type["DELETE"]:
try:
user_memories = await self.memory_system._get_user_memories(user_id)
memory_contents_for_deletion = {str(mem.id): mem.content for mem in user_memories}
except Exception as e:
logger.warning(f"⚠️ Failed to fetch memories for DELETE preview: {str(e)}")
if operations_by_type["CREATE"] or operations_by_type["UPDATE"]:
try:
current_memories = await self.memory_system._get_user_memories(user_id)
if operations_by_type["CREATE"]:
@@ -1002,9 +986,6 @@ class LLMConsolidationService:
operations_by_type["UPDATE"], current_memories, user_id, operation_type="UPDATE", delete_operations=operations_by_type["DELETE"]
)
except Exception as e:
logger.warning(f"⚠️ Semantic deduplication check failed, proceeding with original operations: {str(e)}")
for operation_type, ops in operations_by_type.items():
if not ops:
continue
@@ -1014,7 +995,6 @@ class LLMConsolidationService:
task = self.memory_system._execute_single_operation(operation, user)
batch_tasks.append(task)
try:
results = await asyncio.gather(*batch_tasks, return_exceptions=True)
for idx, result in enumerate(results):
operation = ops[idx]
@@ -1042,10 +1022,6 @@ class LLMConsolidationService:
]:
failed_count += 1
await self.memory_system._emit_status(emitter, f"❌ Failed {operation_type}", done=False)
except Exception as e:
failed_count += len(ops)
logger.error(f"❌ Batch {operation_type} operations failed during memory consolidation: {str(e)}")
await self.memory_system._emit_status(emitter, f"❌ Batch {operation_type} Failed", done=False)
total_executed = created_count + updated_count + deleted_count
logger.info(
@@ -1296,17 +1272,9 @@ class Filter:
async def _generate_embeddings(self, texts: Union[str, List[str]], user_id: str) -> Union[np.ndarray, List[np.ndarray]]:
"""Unified embedding generation for single text or batch with optimized caching using OpenWebUI's embedding function."""
if self._embedding_function is None:
raise RuntimeError("🤖 Embedding function not initialized. Ensure pipeline context is set.")
is_single = isinstance(texts, str)
text_list = [texts] if is_single else texts
if not text_list:
if is_single:
raise ValueError("📏 Empty text provided for embedding generation")
return []
result_embeddings = []
uncached_texts = []
uncached_indices = []
@@ -1360,9 +1328,6 @@ class Filter:
return result_embeddings
def _should_skip_memory_operations(self, user_message: str) -> Tuple[bool, str]:
if self._skip_detector is None:
raise RuntimeError("🤖 Skip detector not initialized")
skip_reason = self._skip_detector.detect_skip_reason(user_message, self.valves.max_message_chars, memory_system=self)
if skip_reason:
status_key = SkipDetector.SkipReason(skip_reason)
@@ -1371,13 +1336,6 @@ class Filter:
def _process_user_message(self, body: Dict[str, Any]) -> Tuple[Optional[str], bool, str]:
"""Extract user message and determine if memory operations should be skipped."""
if not body or "messages" not in body or not isinstance(body["messages"], list):
return (
None,
True,
SkipDetector.STATUS_MESSAGES[SkipDetector.SkipReason.SKIP_SIZE],
)
messages = body["messages"]
user_message = None
@@ -1421,10 +1379,6 @@ class Filter:
return
scores = [memory["relevance"] for memory in memories]
if not scores:
return
top_score = max(scores)
lowest_score = min(scores)
median_score = statistics.median(scores)
@@ -1485,13 +1439,9 @@ class Filter:
return
payload = {"type": "status", "data": {"description": description, "done": done}}
try:
result = emitter(payload)
if asyncio.iscoroutine(result):
await result
except Exception:
pass
async def _retrieve_relevant_memories(
self,
@@ -1549,10 +1499,6 @@ class Filter:
emitter: Optional[Callable] = None,
) -> None:
"""Add memory context to request body with simplified logic."""
if not body or "messages" not in body or not body["messages"]:
logger.warning("⚠️ Invalid request body or no messages found")
return
content_parts = [f"Current Date/Time: {self.format_current_datetime()}"]
memory_count = 0
@@ -1610,26 +1556,16 @@ class Filter:
memory_contents = [memory.content for memory in user_memories]
memory_embeddings = await self._generate_embeddings(memory_contents, user_id)
if len(memory_embeddings) != len(user_memories):
logger.error(f"🔢 Embedding generation failed: generated {len(memory_embeddings)} embeddings but expected {len(user_memories)} for user memories")
return [], self.valves.semantic_retrieval_threshold, []
similarity_scores = []
memory_data = []
for memory_index, memory in enumerate(user_memories):
memory_embedding = memory_embeddings[memory_index]
if memory_embedding is None:
continue
similarity = float(np.dot(query_embedding, memory_embedding))
similarity_scores.append(similarity)
memory_dict = self._build_memory_dict(memory, similarity)
memory_data.append(memory_dict)
if not similarity_scores:
return [], self.valves.semantic_retrieval_threshold, []
memory_data.sort(key=lambda x: x["relevance"], reverse=True)
threshold = self.valves.semantic_retrieval_threshold
@@ -1657,10 +1593,19 @@ class Filter:
return body
user_message, should_skip, skip_reason = self._process_user_message(body)
if not user_message or should_skip:
if __event_emitter__ and skip_reason:
await self._emit_status(__event_emitter__, skip_reason, done=True)
await self._add_memory_context(body, [], user_id, __event_emitter__)
skip_cache_key = self._cache_key(self._cache_manager.SKIP_STATE_CACHE, user_id, user_message or "")
await self._cache_manager.put(
user_id,
self._cache_manager.SKIP_STATE_CACHE,
skip_cache_key,
True,
)
return body
try:
memory_cache_key = self._cache_key(self._cache_manager.MEMORY_CACHE, user_id)
@@ -1709,11 +1654,20 @@ class Filter:
user_id = __user__.get("id") if body and __user__ else None
if not user_id:
return body
user_message, should_skip, skip_reason = self._process_user_message(body)
if not user_message or should_skip:
user_message, _, _ = self._process_user_message(body)
if not user_message:
return body
cache_key = self._cache_key(self._cache_manager.RETRIEVAL_CACHE, user_id, user_message)
cached_similarities = await self._cache_manager.get(user_id, self._cache_manager.RETRIEVAL_CACHE, cache_key)
skip_cache_key = self._cache_key(self._cache_manager.SKIP_STATE_CACHE, user_id, user_message)
should_skip = await self._cache_manager.get(user_id, self._cache_manager.SKIP_STATE_CACHE, skip_cache_key)
if should_skip:
logger.info("⏭️ Skipping outlet consolidation: inlet already detected skip condition")
return body
retrieval_cache_key = self._cache_key(self._cache_manager.RETRIEVAL_CACHE, user_id, user_message)
cached_similarities = await self._cache_manager.get(user_id, self._cache_manager.RETRIEVAL_CACHE, retrieval_cache_key)
task = asyncio.create_task(self._llm_consolidation_service.run_consolidation_pipeline(user_message, user_id, __event_emitter__, cached_similarities))
self._background_tasks.add(task)
@@ -1745,7 +1699,10 @@ class Filter:
try:
retrieval_cleared = await self._cache_manager.clear_user_cache(user_id, self._cache_manager.RETRIEVAL_CACHE)
embedding_cleared = await self._cache_manager.clear_user_cache(user_id, self._cache_manager.EMBEDDING_CACHE)
logger.info(f"🔄 Cleared {retrieval_cleared} retrieval + {embedding_cleared} embedding cache entries for user {user_id}")
skip_state_cleared = await self._cache_manager.clear_user_cache(user_id, self._cache_manager.SKIP_STATE_CACHE)
logger.info(
f"🔄 Cleared {retrieval_cleared} retrieval + {embedding_cleared} embedding + {skip_state_cleared} skip state cache entries for user {user_id}"
)
user_memories = await self._get_user_memories(user_id)
memory_cache_key = self._cache_key(self._cache_manager.MEMORY_CACHE, user_id)
@@ -1774,11 +1731,9 @@ class Filter:
async def _execute_single_operation(self, operation: Models.MemoryOperation, user: Any) -> str:
"""Execute a single memory operation."""
try:
if operation.operation == Models.MemoryOperationType.CREATE:
content_stripped = operation.content.strip()
if not content_stripped:
logger.warning(f"⚠️ Skipping CREATE operation: empty content")
return Models.OperationResult.SKIPPED_EMPTY_CONTENT.value
await asyncio.wait_for(
@@ -1790,12 +1745,10 @@ class Filter:
elif operation.operation == Models.MemoryOperationType.UPDATE:
id_stripped = operation.id.strip()
if not id_stripped:
logger.warning(f"⚠️ Skipping UPDATE operation: empty ID")
return Models.OperationResult.SKIPPED_EMPTY_ID.value
content_stripped = operation.content.strip()
if not content_stripped:
logger.warning(f"⚠️ Skipping UPDATE operation for {id_stripped}: empty content")
return Models.OperationResult.SKIPPED_EMPTY_CONTENT.value
await asyncio.wait_for(
@@ -1812,7 +1765,6 @@ class Filter:
elif operation.operation == Models.MemoryOperationType.DELETE:
id_stripped = operation.id.strip()
if not id_stripped:
logger.warning(f"⚠️ Skipping DELETE operation: empty ID")
return Models.OperationResult.SKIPPED_EMPTY_ID.value
await asyncio.wait_for(
@@ -1820,13 +1772,8 @@ class Filter:
timeout=Constants.DATABASE_OPERATION_TIMEOUT_SEC,
)
return Models.MemoryOperationType.DELETE.value
else:
logger.error(f"❓ Unsupported operation: {operation}")
return Models.OperationResult.UNSUPPORTED.value
except Exception as e:
logger.error(f"💾 Database operation failed for {operation.operation.value}: {str(e)}")
return Models.OperationResult.FAILED.value
return Models.OperationResult.UNSUPPORTED.value
def _remove_refs_from_schema(self, schema: Dict[str, Any], schema_defs: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
"""Remove $ref references and ensure required fields for Azure OpenAI."""