merge dev

Author: yuan.wang

src/memos/graph_dbs/neo4j.py

@@ -367,7 +367,8 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                             if kb_log_content:
                                 event = self.create_event_log(
                                     label="knowledgeBaseUpdate",
-                                    log_content=f"Knowledge Base Memory Update: {len(kb_log_content)} changes.",
+                                    from_memory_type=USER_INPUT_TYPE,
+                                    to_memory_type=LONG_TERM_MEMORY_TYPE,
                                     user_id=msg.user_id,
                                     mem_cube_id=msg.mem_cube_id,
                                     mem_cube=self.current_mem_cube,
@@ -376,6 +377,9 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                                     memory_len=len(kb_log_content),
                                     memcube_name=self._map_memcube_name(msg.mem_cube_id),
                                 )
+                                event.log_content = (
+                                    f"Knowledge Base Memory Update: {len(kb_log_content)} changes."
+                                )
                                 event.task_id = msg.task_id
                                 self._submit_web_logs([event])
                         else:
@@ -504,6 +508,8 @@ def process_message(message: ScheduleMessageItem):
                     text_mem=text_mem,
                     user_name=user_name,
                     custom_tags=info.get("custom_tags", None),
+                    task_id=message.task_id,
+                    info=info,
                 )
 
                 logger.info(
@@ -529,6 +535,8 @@ def _process_memories_with_reader(
         text_mem: TreeTextMemory,
         user_name: str,
         custom_tags: list[str] | None = None,
+        task_id: str | None = None,
+        info: dict | None = None,
     ) -> None:
         """
         Process memories using mem_reader for enhanced memory processing.
@@ -540,6 +548,7 @@ def _process_memories_with_reader(
             text_mem: Text memory instance
             custom_tags: Optional list of custom tags for memory processing
         """
+        kb_log_content: list[dict] = []
         try:
             # Get the mem_reader from the parent MOSCore
             if not hasattr(self, "mem_reader") or self.mem_reader is None:
@@ -602,6 +611,90 @@ def _process_memories_with_reader(
                     logger.info(
                         f"Added {len(enhanced_mem_ids)} enhanced memories: {enhanced_mem_ids}"
                     )
+
+                    # LOGGING BLOCK START
+                    # This block is replicated from _add_message_consumer to ensure consistent logging
+                    is_cloud_env = (
+                        os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME") == "memos-memory-change"
+                    )
+                    if is_cloud_env:
+                        # New: Knowledge Base Logging (Cloud Service)
+                        kb_log_content = []
+                        for item in flattened_memories:
+                            kb_log_content.append(
+                                {
+                                    "log_source": "KNOWLEDGE_BASE_LOG",
+                                    "trigger_source": info.get("trigger_source", "Messages")
+                                    if info
+                                    else "Messages",
+                                    "operation": "ADD",
+                                    "memory_id": item.id,
+                                    "content": item.memory,
+                                    "original_content": None,
+                                    "source_doc_id": getattr(item.metadata, "source_doc_id", None),
+                                }
+                            )
+                        if kb_log_content:
+                            event = self.create_event_log(
+                                label="knowledgeBaseUpdate",
+                                from_memory_type=USER_INPUT_TYPE,
+                                to_memory_type=LONG_TERM_MEMORY_TYPE,
+                                user_id=user_id,
+                                mem_cube_id=mem_cube_id,
+                                mem_cube=self.current_mem_cube,
+                                memcube_log_content=kb_log_content,
+                                metadata=None,
+                                memory_len=len(kb_log_content),
+                                memcube_name=self._map_memcube_name(mem_cube_id),
+                            )
+                            event.log_content = (
+                                f"Knowledge Base Memory Update: {len(kb_log_content)} changes."
+                            )
+                            event.task_id = task_id
+                            self._submit_web_logs([event])
+                    else:
+                        # Existing: Playground/Default Logging
+                        add_content_legacy: list[dict] = []
+                        add_meta_legacy: list[dict] = []
+                        for item_id, item in zip(
+                            enhanced_mem_ids, flattened_memories, strict=False
+                        ):
+                            key = getattr(item.metadata, "key", None) or transform_name_to_key(
+                                name=item.memory
+                            )
+                            add_content_legacy.append(
+                                {"content": f"{key}: {item.memory}", "ref_id": item_id}
+                            )
+                            add_meta_legacy.append(
+                                {
+                                    "ref_id": item_id,
+                                    "id": item_id,
+                                    "key": item.metadata.key,
+                                    "memory": item.memory,
+                                    "memory_type": item.metadata.memory_type,
+                                    "status": item.metadata.status,
+                                    "confidence": item.metadata.confidence,
+                                    "tags": item.metadata.tags,
+                                    "updated_at": getattr(item.metadata, "updated_at", None)
+                                    or getattr(item.metadata, "update_at", None),
+                                }
+                            )
+                        if add_content_legacy:
+                            event = self.create_event_log(
+                                label="addMemory",
+                                from_memory_type=USER_INPUT_TYPE,
+                                to_memory_type=LONG_TERM_MEMORY_TYPE,
+                                user_id=user_id,
+                                mem_cube_id=mem_cube_id,
+                                mem_cube=self.current_mem_cube,
+                                memcube_log_content=add_content_legacy,
+                                metadata=add_meta_legacy,
+                                memory_len=len(add_content_legacy),
+                                memcube_name=self._map_memcube_name(mem_cube_id),
+                            )
+                            event.task_id = task_id
+                            self._submit_web_logs([event])
+                    # LOGGING BLOCK END
                 else:
                     logger.info("No enhanced memories generated by mem_reader")
             else:
@@ -630,10 +723,47 @@ def _process_memories_with_reader(
             logger.info("Remove and Refresh Memories")
             logger.debug(f"Finished add {user_id} memory: {mem_ids}")
 
-        except Exception:
+        except Exception as exc:
             logger.error(
                 f"Error in _process_memories_with_reader: {traceback.format_exc()}", exc_info=True
             )
+            with contextlib.suppress(Exception):
+                is_cloud_env = (
+                    os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME") == "memos-memory-change"
+                )
+                if is_cloud_env:
+                    if not kb_log_content:
+                        trigger_source = (
+                            info.get("trigger_source", "Messages") if info else "Messages"
+                        )
+                        kb_log_content = [
+                            {
+                                "log_source": "KNOWLEDGE_BASE_LOG",
+                                "trigger_source": trigger_source,
+                                "operation": "ADD",
+                                "memory_id": mem_id,
+                                "content": None,
+                                "original_content": None,
+                                "source_doc_id": None,
+                            }
+                            for mem_id in mem_ids
+                        ]
+                    event = self.create_event_log(
+                        label="knowledgeBaseUpdate",
+                        from_memory_type=USER_INPUT_TYPE,
+                        to_memory_type=LONG_TERM_MEMORY_TYPE,
+                        user_id=user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=self.current_mem_cube,
+                        memcube_log_content=kb_log_content,
+                        metadata=None,
+                        memory_len=len(kb_log_content),
+                        memcube_name=self._map_memcube_name(mem_cube_id),
+                    )
+                    event.log_content = f"Knowledge Base Memory Update failed: {exc!s}"
+                    event.task_id = task_id
+                    event.status = "failed"
+                    self._submit_web_logs([event])
 
     def _mem_reorganize_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         logger.info(f"Messages {messages} assigned to {MEM_ORGANIZE_LABEL} handler.")

src/memos/graph_dbs/neo4j_community.py

@@ -1,5 +1,4 @@
 import concurrent
-import os
 import threading
 import time
 
@@ -15,7 +14,7 @@
 from memos.mem_scheduler.schemas.general_schemas import (
     DEFAULT_STOP_WAIT,
 )
-from memos.mem_scheduler.schemas.message_schemas import ScheduleLogForWebItem, ScheduleMessageItem
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 from memos.mem_scheduler.schemas.task_schemas import RunningTaskItem
 from memos.mem_scheduler.utils.misc_utils import group_messages_by_user_and_mem_cube
 from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
@@ -159,20 +158,6 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                     )
                 self.metrics.task_completed(user_id=m.user_id, task_type=m.label)
 
-                is_cloud_env = (
-                    os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME") == "memos-memory-change"
-                )
-                if self.submit_web_logs and is_cloud_env:
-                    status_log = ScheduleLogForWebItem(
-                        user_id=task_item.user_id,
-                        mem_cube_id=task_item.mem_cube_id,
-                        item_id=task_item.item_id,
-                        label=m.label,
-                        log_content=f"Task {task_item.item_id} completed successfully for user {task_item.user_id}.",
-                        status="completed",
-                    )
-                    self.submit_web_logs([status_log])
-
                 # acknowledge redis messages
                 if self.use_redis_queue and self.memos_message_queue is not None:
                     for msg in messages:
@@ -211,20 +196,6 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                             self._completed_tasks.pop(0)
                 logger.error(f"Task failed: {task_item.get_execution_info()}, Error: {e}")
 
-                is_cloud_env = (
-                    os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME") == "memos-memory-change"
-                )
-                if self.submit_web_logs and is_cloud_env:
-                    status_log = ScheduleLogForWebItem(
-                        user_id=task_item.user_id,
-                        mem_cube_id=task_item.mem_cube_id,
-                        item_id=task_item.item_id,
-                        label=m.label,
-                        log_content=f"Task {task_item.item_id} failed for user {task_item.user_id} with error: {e!s}.",
-                        status="failed",
-                        exception=str(e),
-                    )
-                    self.submit_web_logs([status_log])
                 raise
 
         return wrapped_handler

src/memos/graph_dbs/polardb.py

@@ -203,6 +203,7 @@ class PreferenceTextualMemoryMetadata(TextualMemoryMetadata):
     preference: str | None = Field(default=None, description="Preference.")
     created_at: str | None = Field(default=None, description="Timestamp of the dialog.")
     mem_cube_id: str | None = Field(default=None, description="ID of the MemCube.")
+    score: float | None = Field(default=None, description="Score of the retrieval result.")
 
 
 class TextualMemoryItem(BaseModel):

src/memos/mem_scheduler/general_scheduler.py

@@ -90,7 +90,8 @@ def extract_implicit_preference(self, qa_pair: MessageList | str) -> dict[str, A
             response = self.llm_provider.generate([{"role": "user", "content": prompt}])
             response = response.strip().replace("```json", "").replace("```", "").strip()
             result = json.loads(response)
-            result["preference"] = result.pop("implicit_preference")
+            for d in result:
+                d["preference"] = d.pop("implicit_preference")
             return result
         except Exception as e:
             logger.error(f"Error extracting implicit preferences: {e}, return None")
@@ -137,20 +138,24 @@ def _process_single_chunk_implicit(
         if not implicit_pref:
             return None
 
-        vector_info = {
-            "embedding": self.embedder.embed([implicit_pref["context_summary"]])[0],
-        }
+        memories = []
+        for pref in implicit_pref:
+            vector_info = {
+                "embedding": self.embedder.embed([pref["context_summary"]])[0],
+            }
 
-        extract_info = {**basic_info, **implicit_pref, **vector_info, **info}
+            extract_info = {**basic_info, **pref, **vector_info, **info}
 
-        metadata = PreferenceTextualMemoryMetadata(
-            type=msg_type, preference_type="implicit_preference", **extract_info
-        )
-        memory = TextualMemoryItem(
-            id=extract_info["dialog_id"], memory=implicit_pref["context_summary"], metadata=metadata
-        )
+            metadata = PreferenceTextualMemoryMetadata(
+                type=msg_type, preference_type="implicit_preference", **extract_info
+            )
+            memory = TextualMemoryItem(
+                id=str(uuid.uuid4()), memory=pref["context_summary"], metadata=metadata
+            )
 
-        return memory
+            memories.append(memory)
+
+        return memories
 
     def extract(
         self,

src/memos/mem_scheduler/task_schedule_modules/dispatcher.py

@@ -1,3 +1,5 @@
+import os
+
 from abc import ABC, abstractmethod
 from typing import Any
 
@@ -34,9 +36,12 @@ def _naive_reranker(
         self, query: str, prefs_mem: list[TextualMemoryItem], top_k: int, **kwargs: Any
     ) -> list[TextualMemoryItem]:
         if self.reranker:
-            prefs_mem = self.reranker.rerank(query, prefs_mem, top_k)
-            return [item for item, _ in prefs_mem]
-        return prefs_mem
+            prefs_mem_reranked = []
+            prefs_mem_tuple = self.reranker.rerank(query, prefs_mem, top_k)
+            for item, score in prefs_mem_tuple:
+                item.metadata.score = score
+                prefs_mem_reranked.append(item)
+        return prefs_mem_reranked
 
     def _original_text_reranker(
         self,
@@ -52,11 +57,22 @@ def _original_text_reranker(
             prefs_mem_for_reranker = deepcopy(prefs_mem)
             for pref_mem, pref in zip(prefs_mem_for_reranker, prefs, strict=False):
                 pref_mem.memory = pref_mem.memory + "\n" + pref.original_text
-            prefs_mem_for_reranker = self.reranker.rerank(query, prefs_mem_for_reranker, top_k)
-            prefs_mem_for_reranker = [item for item, _ in prefs_mem_for_reranker]
+            reranked_results = self.reranker.rerank(query, prefs_mem_for_reranker, top_k)
+            prefs_mem_for_reranker = [item for item, _ in reranked_results]
             prefs_ids = [item.id for item in prefs_mem_for_reranker]
             prefs_dict = {item.id: item for item in prefs_mem}
-            return [prefs_dict[item_id] for item_id in prefs_ids if item_id in prefs_dict]
+
+            # Create mapping from id to score from reranked results
+            reranked_scores = {item.id: score for item, score in reranked_results}
+
+            # Assign scores to the original items
+            result_items = []
+            for item_id in prefs_ids:
+                if item_id in prefs_dict:
+                    original_item = prefs_dict[item_id]
+                    original_item.metadata.score = reranked_scores.get(item_id)
+                    result_items.append(original_item)
+            return result_items
         return prefs_mem
 
     def retrieve(
@@ -119,24 +135,34 @@ def retrieve(
             if pref.payload.get("preference", None)
         ]
 
-        # store explicit id and score, use it after reranker
-        explicit_id_scores = {item.id: item.score for item in explicit_prefs}
-
         reranker_map = {
             "naive": self._naive_reranker,
             "original_text": self._original_text_reranker,
         }
         reranker_func = reranker_map["naive"]
-        explicit_prefs_mem = reranker_func(
-            query=query, prefs_mem=explicit_prefs_mem, prefs=explicit_prefs, top_k=top_k
+        prefs_mem_explicit = reranker_func(
+            query=query,
+            prefs_mem=explicit_prefs_mem,
+            prefs=explicit_prefs,
+            top_k=top_k,
         )
-        implicit_prefs_mem = reranker_func(
-            query=query, prefs_mem=implicit_prefs_mem, prefs=implicit_prefs, top_k=top_k
+        prefs_mem_implicit = reranker_func(
+            query=query,
+            prefs_mem=implicit_prefs_mem,
+            prefs=implicit_prefs,
+            top_k=top_k,
         )
 
         # filter explicit mem by score bigger than threshold
-        explicit_prefs_mem = [
-            item for item in explicit_prefs_mem if explicit_id_scores.get(item.id, 0) >= 0.0
+        prefs_mem_explicit = [
+            item
+            for item in prefs_mem_explicit
+            if item.metadata.score >= float(os.getenv("PREFERENCE_SEARCH_THRESHOLD", 0.0))
+        ]
+        prefs_mem_implicit = [
+            item
+            for item in prefs_mem_implicit
+            if item.metadata.score >= float(os.getenv("PREFERENCE_SEARCH_THRESHOLD", 0.0))
         ]
 
-        return explicit_prefs_mem + implicit_prefs_mem
+        return prefs_mem_explicit + prefs_mem_implicit