feat: add image parser in file

Author: CaralHsi

examples/mem_scheduler/try_schedule_modules.py

@@ -41,14 +41,26 @@ def init_task():
             "role": "user",
             "content": "Max needs joint supplements. Actually, we're moving to Chicago next month.",
         },
+        {
+            "role": "assistant",
+            "content": "Got it — Max is on joint supplements, and you’re relocating to Chicago soon. That’s a big move! Have you looked into how the change in climate or vet access might affect his needs?",
+        },
         {
             "role": "user",
             "content": "Correction: Bella is 6, not 5. And she's allergic to chicken.",
         },
+        {
+            "role": "assistant",
+            "content": "Thanks for the update! So Bella is 6 years old and has a chicken allergy — good to know. You’ll want to double-check her food and treats, especially during the move. Has she had any reactions recently?",
+        },
         {
             "role": "user",
             "content": "My partner's cat Whiskers visits weekends. Bella chases her sometimes.",
         },
+        {
+            "role": "assistant",
+            "content": "Ah, the classic dog-and-cat dynamic! Since Bella chases Whiskers, it might help to give them gradual supervised interactions or create safe zones for the cat—especially important as you settle into a new home in Chicago. Keeping Bella’s routine stable during the move could also reduce her urge to chase. How do they usually get along when Whiskers visits?",
+        },
     ]
 
     questions = [
@@ -145,18 +157,25 @@ def start_conversation(self, user_id="test_user", mem_cube_id="test_cube", sessi
         print(f"   User ID: {self.current_user_id}")
         print(f"   Mem Cube ID: {self.current_mem_cube_id}")
 
-    def add_msgs(self, messages: list[dict]):
+    def add_msgs(
+        self,
+        messages: list[dict],
+        extract_mode: str = "fine",
+        async_mode: str = "sync",
+    ):
         # Create add request
         add_req = self.create_test_add_request(
             user_id=self.current_user_id,
             mem_cube_id=self.current_mem_cube_id,
             messages=messages,
             session_id=self.current_session_id,
+            extract_mode=extract_mode,
+            async_mode=async_mode,
         )
 
         # Add to memory
         result = self.add_memories(add_req)
-        print(f"   ✅ Added to memory successfully: \n{messages}")
+        print(f"   ✅ Added to memory successfully: \n{result}")
 
         return result
 

src/memos/mem_reader/read_multi_modal/file_content_parser.py

@@ -259,6 +259,7 @@ def create_source(
         chunk_index: int | None = None,
         chunk_total: int | None = None,
         chunk_content: str | None = None,
+        file_url_flag: bool = False,
     ) -> SourceMessage:
         """Create SourceMessage from file content part."""
         if isinstance(message, dict):
@@ -267,6 +268,7 @@ def create_source(
                 "type": "file",
                 "doc_path": file_info.get("filename") or file_info.get("file_id", ""),
                 "content": chunk_content if chunk_content else file_info.get("file_data", ""),
+                "file_info": file_info if file_url_flag else {},
             }
             # Add chunk ordering information if provided
             if chunk_index is not None:
@@ -291,10 +293,7 @@ def rebuild_from_source(
         # Rebuild from source fields
         return {
             "type": "file",
-            "file": {
-                "filename": source.doc_path or "",
-                "file_data": source.content or "",
-            },
+            "file": source.file_info,
         }
 
     def _parse_file(self, file_info: dict[str, Any]) -> str:
@@ -367,7 +366,7 @@ def parse_fast(
         file_data = file_info.get("file_data", "")
         file_id = file_info.get("file_id", "")
         filename = file_info.get("filename", "")
-
+        file_url_flag = False
         # Build content string based on available information
         content_parts = []
 
@@ -386,6 +385,7 @@ def parse_fast(
                     content_parts.append(f"[File Data (base64/encoded): {len(file_data)} chars]")
                 # Check if it looks like a URL
                 elif file_data.startswith(("http://", "https://", "file://")):
+                    file_url_flag = True
                     content_parts.append(f"[File URL: {file_data}]")
                 else:
                     # TODO: split into multiple memory items
@@ -437,6 +437,7 @@ def parse_fast(
                 chunk_index=chunk_idx,
                 chunk_total=total_chunks,
                 chunk_content=chunk_text,
+                file_url_flag=file_url_flag,
             )
 
             memory_item = TextualMemoryItem(
@@ -473,6 +474,7 @@ def parse_fast(
                 chunk_index=None,
                 chunk_total=0,
                 chunk_content=content,
+                file_url_flag=file_url_flag,
             )
             memory_item = TextualMemoryItem(
                 memory=content,

src/memos/mem_reader/read_multi_modal/user_parser.py

@@ -80,6 +80,7 @@ def create_source(
                                 message_id=message_id,
                                 doc_path=file_info.get("filename") or file_info.get("file_id", ""),
                                 content=file_info.get("file_data", ""),
+                                file_info=file_info,
                             )
                         )
                     elif part_type == "image_url":

src/memos/mem_reader/simple_struct.py

@@ -1,6 +1,7 @@
 import concurrent.futures
 import copy
 import json
+import os
 import re
 import traceback
 
@@ -25,6 +26,7 @@
 from memos.templates.mem_reader_prompts import (
     CUSTOM_TAGS_INSTRUCTION,
     CUSTOM_TAGS_INSTRUCTION_ZH,
+    PROMPT_MAPPING,
     SIMPLE_STRUCT_DOC_READER_PROMPT,
     SIMPLE_STRUCT_DOC_READER_PROMPT_ZH,
     SIMPLE_STRUCT_MEM_READER_EXAMPLE,
@@ -80,6 +82,7 @@ def from_config(_config):
     "custom_tags": {"en": CUSTOM_TAGS_INSTRUCTION, "zh": CUSTOM_TAGS_INSTRUCTION_ZH},
 }
 
+
 try:
     import tiktoken
 
@@ -448,6 +451,81 @@ def get_memory(
         standard_scene_data = coerce_scene_data(scene_data, type)
         return self._read_memory(standard_scene_data, type, info, mode)
 
+    @staticmethod
+    def _parse_hallucination_filter_response(text: str) -> tuple[bool, dict[int, dict]]:
+        """Parse index-keyed JSON from hallucination filter response.
+        Expected shape: { "0": {"if_delete": bool, "rewritten memory content": str}, ... }
+        Returns (success, parsed_dict) with int keys.
+        """
+        try:
+            data = json.loads(text)
+        except Exception:
+            return False, {}
+
+        if not isinstance(data, dict):
+            return False, {}
+
+        result: dict[int, dict] = {}
+        for k, v in data.items():
+            try:
+                idx = int(k)
+            except Exception:
+                # allow integer keys as-is
+                if isinstance(k, int):
+                    idx = k
+                else:
+                    continue
+            if not isinstance(v, dict):
+                continue
+            delete_flag = v.get("delete_flag")
+            rewritten = v.get("rewritten memory content", "")
+            if isinstance(delete_flag, bool) and isinstance(rewritten, str):
+                result[idx] = {"delete_flag": delete_flag, "rewritten memory content": rewritten}
+
+        return (len(result) > 0), result
+
+    def filter_hallucination_in_memories(
+        self, user_messages: list[str], memory_list: list[list[TextualMemoryItem]]
+    ):
+        filtered_memory_list = []
+        for group in memory_list:
+            try:
+                flat_memories = [one.memory for one in group]
+                template = PROMPT_MAPPING["hallucination_filter"]
+                prompt_args = {
+                    "user_messages_inline": "\n".join(user_messages),
+                    "memories_inline": json.dumps(flat_memories, ensure_ascii=False, indent=2),
+                }
+                prompt = template.format(**prompt_args)
+
+                # Optionally run filter and parse the output
+                try:
+                    raw = self.llm.generate(prompt)
+                    success, parsed = self._parse_hallucination_filter_response(raw)
+                    logger.info(f"Hallucination filter parsed successfully: {success}")
+                    new_mem_list = []
+                    if success:
+                        logger.info(f"Hallucination filter result: {parsed}")
+                        for mem_idx, (delete_flag, rewritten_mem_content) in parsed.items():
+                            if not delete_flag:
+                                group[mem_idx].memory = rewritten_mem_content
+                                new_mem_list.append(group[mem_idx])
+                        filtered_memory_list.append(new_mem_list)
+                        logger.info(
+                            f"Successfully transform origianl memories from {group} to {new_mem_list}."
+                        )
+                    else:
+                        logger.warning(
+                            "Hallucination filter parsing failed or returned empty result."
+                        )
+                except Exception as e:
+                    logger.error(f"Hallucination filter execution error: {e}", stack_info=True)
+                    filtered_memory_list.append(group)
+            except Exception:
+                logger.error("Fail to filter memories", stack_info=True)
+                filtered_memory_list.append(group)
+        return filtered_memory_list
+
     def _read_memory(
         self, messages: list[MessagesType], type: str, info: dict[str, Any], mode: str = "fine"
     ) -> list[list[TextualMemoryItem]]:
@@ -492,6 +570,14 @@ def _read_memory(
                 except Exception as e:
                     logger.error(f"Task failed with exception: {e}")
                     logger.error(traceback.format_exc())
+
+        if os.getenv("SIMPLE_STRUCT_ADD_FILTER", "false") == "true":
+            # Build inputs
+            user_messages = [msg.content for msg in messages if msg.role == "user"]
+            memory_list = self.filter_hallucination_in_memories(
+                user_messages=user_messages, memory_list=memory_list
+            )
+
         return memory_list
 
     def fine_transfer_simple_mem(

src/memos/mem_scheduler/analyzer/api_analyzer.py

@@ -599,6 +599,8 @@ def create_test_add_request(
         messages=None,
         memory_content=None,
         session_id=None,
+        extract_mode=None,
+        async_mode="sync",
     ):
         """
         Create a test APIADDRequest object with the given parameters.
@@ -637,6 +639,8 @@ def create_test_add_request(
             source="api_analyzer_test",
             chat_history=None,
             operation=None,
+            mode=extract_mode,
+            async_mode=async_mode,
         )
 
     def run_all_tests(self, mode=SearchMode.MIXTURE):

src/memos/mem_scheduler/base_scheduler.py

@@ -140,12 +140,7 @@ def __init__(self, config: BaseSchedulerConfig):
             "max_internal_message_queue_size", DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE
         )
         self.orchestrator = SchedulerOrchestrator()
-        self.memos_message_queue = ScheduleTaskQueue(
-            use_redis_queue=self.use_redis_queue,
-            maxsize=self.max_internal_message_queue_size,
-            disabled_handlers=self.disabled_handlers,
-            orchestrator=self.orchestrator,
-        )
+
         self.searcher: Searcher | None = None
         self.retriever: SchedulerRetriever | None = None
         self.db_engine: Engine | None = None
@@ -155,6 +150,13 @@ def __init__(self, config: BaseSchedulerConfig):
         self.status_tracker: TaskStatusTracker | None = None
         self.metrics = metrics
         self._monitor_thread = None
+        self.memos_message_queue = ScheduleTaskQueue(
+            use_redis_queue=self.use_redis_queue,
+            maxsize=self.max_internal_message_queue_size,
+            disabled_handlers=self.disabled_handlers,
+            orchestrator=self.orchestrator,
+            status_tracker=self.status_tracker,
+        )
         self.dispatcher = SchedulerDispatcher(
             config=self.config,
             memos_message_queue=self.memos_message_queue,
@@ -228,6 +230,8 @@ def initialize_modules(
                 self.status_tracker = TaskStatusTracker(redis_client)
                 if self.dispatcher:
                     self.dispatcher.status_tracker = self.status_tracker
+                if self.memos_message_queue:
+                    self.memos_message_queue.status_tracker = self.status_tracker
             # initialize submodules
             self.chat_llm = chat_llm
             self.process_llm = process_llm
@@ -712,7 +716,13 @@ def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageIt
             # emit enqueue events for consistency
             for m in immediate_msgs:
                 emit_monitor_event(
-                    "enqueue", m, {"enqueue_ts": to_iso(getattr(m, "timestamp", None))}
+                    "enqueue",
+                    m,
+                    {
+                        "enqueue_ts": to_iso(getattr(m, "timestamp", None)),
+                        "event_duration_ms": 0,
+                        "total_duration_ms": 0,
+                    },
                 )
 
             # simulate dequeue for immediately dispatched messages so monitor logs stay complete
@@ -741,6 +751,8 @@ def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageIt
                             "enqueue_ts": to_iso(enqueue_ts_obj),
                             "dequeue_ts": datetime.fromtimestamp(now, tz=timezone.utc).isoformat(),
                             "queue_wait_ms": queue_wait_ms,
+                            "event_duration_ms": queue_wait_ms,
+                            "total_duration_ms": queue_wait_ms,
                         },
                     )
                     self.metrics.task_dequeued(user_id=m.user_id, task_type=m.label)
@@ -919,6 +931,8 @@ def _message_consumer(self) -> None:
                                         now, tz=timezone.utc
                                     ).isoformat(),
                                     "queue_wait_ms": queue_wait_ms,
+                                    "event_duration_ms": queue_wait_ms,
+                                    "total_duration_ms": queue_wait_ms,
                                 },
                             )
                             self.metrics.task_dequeued(user_id=msg.user_id, task_type=msg.label)

src/memos/mem_scheduler/general_modules/misc.py

@@ -233,7 +233,7 @@ def put(self, item: T, block: bool = False, timeout: float | None = None) -> Non
 
     def get(
         self, block: bool = True, timeout: float | None = None, batch_size: int | None = None
-    ) -> list[T] | T:
+    ) -> list[T]:
         """Get items from the queue.
 
         Args:

src/memos/mem_scheduler/schemas/task_schemas.py

@@ -62,10 +62,9 @@ class TaskPriorityLevel(Enum):
 
 
 # task queue
-DEFAULT_STREAM_KEY_PREFIX = "scheduler:messages:stream:v1.7"
-exchange_name = os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME", None)
-if exchange_name is not None:
-    DEFAULT_STREAM_KEY_PREFIX += f":{exchange_name}"
+DEFAULT_STREAM_KEY_PREFIX = os.getenv(
+    "MEMSCHEDULER_STREAM_KEY_PREFIX", "scheduler:messages:stream:v2.0"
+)
 
 
 # ============== Running Tasks ==============

src/memos/mem_scheduler/task_schedule_modules/dispatcher.py

@@ -185,6 +185,8 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                             if isinstance(dequeue_ts, int | float)
                             else None
                         ),
+                        "event_duration_ms": start_delay_ms,
+                        "total_duration_ms": self._calc_total_duration_ms(start_time, enq_ts),
                     },
                 )
 
@@ -210,6 +212,7 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                             finish_time, tz=timezone.utc
                         ).isoformat(),
                         "exec_duration_ms": duration * 1000,
+                        "event_duration_ms": duration * 1000,
                         "total_duration_ms": self._calc_total_duration_ms(
                             finish_time, getattr(first_msg, "timestamp", None)
                         ),
@@ -244,6 +247,7 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                             finish_time, tz=timezone.utc
                         ).isoformat(),
                         "exec_duration_ms": (finish_time - start_time) * 1000,
+                        "event_duration_ms": (finish_time - start_time) * 1000,
                         "error_type": type(e).__name__,
                         "error_msg": str(e),
                         "total_duration_ms": self._calc_total_duration_ms(
@@ -273,6 +277,7 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                                 mem_cube_id=msg.mem_cube_id,
                                 task_label=msg.label,
                                 redis_message_id=redis_message_id,
+                                message=msg,
                             )
                     except Exception as ack_err:
                         logger.warning(f"Ack in finally failed: {ack_err}")