fix(ci): Reformat code to comply with ruff standards

Author: [email protected]

examples/mem_scheduler/api_w_scheduler.py

@@ -27,7 +27,9 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
     print(f"My test handler received {len(messages)} messages:")
     for msg in messages:
         print(f" my_test_handler - {msg.item_id}: {msg.content}")
-        user_status_running = handle_scheduler_status(user_id=msg.user_id, status_tracker=status_tracker)
+        user_status_running = handle_scheduler_status(
+            user_id=msg.user_id, status_tracker=status_tracker
+        )
         print("[Monitor] Status after submit:", user_status_running)
 
 
@@ -37,9 +39,7 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
 mem_scheduler.register_handlers({TEST_HANDLER_LABEL: my_test_handler})
 
 # 2.1 Monitor global scheduler status before submitting tasks
-global_status_before = handle_scheduler_status(
-    user_id=TEST_USER_ID, status_tracker=status_tracker
-)
+global_status_before = handle_scheduler_status(user_id=TEST_USER_ID, status_tracker=status_tracker)
 print("[Monitor] Global status before submit:", global_status_before)
 
 # 3. Create messages
@@ -78,9 +78,7 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
 print(f"[Monitor] Wait result for {USER_MEM_CUBE}:", wait_result)
 
 # 6.2 Monitor global scheduler status after processing
-global_status_after = handle_scheduler_status(
-    user_id=TEST_USER_ID, status_tracker=status_tracker
-)
+global_status_after = handle_scheduler_status(user_id=TEST_USER_ID, status_tracker=status_tracker)
 print("[Monitor] Global status after processing:", global_status_after)
 
 # 7. Stop the scheduler

src/memos/api/handlers/component_init.py

@@ -131,14 +131,17 @@ def init_server() -> dict[str, Any]:
     # Initialize Redis client first as it is a core dependency for features like scheduler status tracking
     try:
         from memos.mem_scheduler.orm_modules.api_redis_model import APIRedisDBManager
+
         redis_client = APIRedisDBManager.load_redis_engine_from_env()
         if redis_client:
             logger.info("Redis client initialized successfully.")
         else:
-            logger.error("Failed to initialize Redis client. Check REDIS_HOST etc. in environment variables.")
+            logger.error(
+                "Failed to initialize Redis client. Check REDIS_HOST etc. in environment variables."
+            )
     except Exception as e:
         logger.error(f"Failed to initialize Redis client: {e}", exc_info=True)
-        redis_client = None # Ensure redis_client exists even on failure
+        redis_client = None  # Ensure redis_client exists even on failure
 
     # Get default cube configuration
     default_cube_config = APIConfig.get_default_cube_config()

src/memos/api/routers/server_router.py

@@ -112,12 +112,12 @@ def add_memories(add_req: APIADDRequest):
 # =============================================================================
 
 
-@router.get( # Changed from post to get
+@router.get(  # Changed from post to get
     "/scheduler/status", summary="Get scheduler running status", response_model=StatusResponse
 )
 def scheduler_status(
     user_id: str = Query(..., description="User ID"),
-    task_id: str | None = Query(None, description="Optional Task ID to query a specific task")
+    task_id: str | None = Query(None, description="Optional Task ID to query a specific task"),
 ):
     """Get scheduler running status."""
     return handlers.scheduler_handler.handle_scheduler_status(

src/memos/mem_scheduler/base_scheduler.py

@@ -731,14 +731,14 @@ def _monitor_loop(self):
                     if len(parts) >= 3:
                         user_id = parts[2]
                         self.metrics.update_queue_length(queue_length, user_id)
-                    elif not self.use_redis_queue: # local queue
-                         user_id = stream_key
-                         self.metrics.update_queue_length(queue_length, user_id)
+                    elif not self.use_redis_queue:  # local queue
+                        user_id = stream_key
+                        self.metrics.update_queue_length(queue_length, user_id)
 
             except Exception as e:
                 logger.error(f"Error in metrics monitor loop: {e}", exc_info=True)
 
-            time.sleep(15) # 每 15 秒采样一次
+            time.sleep(15)  # 每 15 秒采样一次
 
     def start(self) -> None:
         """
@@ -760,7 +760,9 @@ def start(self) -> None:
     def start_background_monitor(self):
         if self._monitor_thread and self._monitor_thread.is_alive():
             return
-        self._monitor_thread = ContextThread(target=self._monitor_loop, daemon=True, name="SchedulerMetricsMonitor")
+        self._monitor_thread = ContextThread(
+            target=self._monitor_loop, daemon=True, name="SchedulerMetricsMonitor"
+        )
         self._monitor_thread.start()
         logger.info("Scheduler metrics monitor thread started.")
 

src/memos/mem_scheduler/schemas/message_schemas.py

@@ -141,7 +141,9 @@ class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
     )
     memcube_name: str | None = Field(default=None, description="Display name for memcube")
     memory_len: int | None = Field(default=None, description="Count of items involved in the event")
-    status: str | None = Field(default=None, description="Completion status of the task (e.g., 'completed', 'failed')")
+    status: str | None = Field(
+        default=None, description="Completion status of the task (e.g., 'completed', 'failed')"
+    )
 
     def debug_info(self) -> dict[str, Any]:
         """Return structured debug information for logging purposes."""

src/memos/mem_scheduler/task_schedule_modules/dispatcher.py

@@ -46,7 +46,7 @@ def __init__(
         config=None,
         status_tracker: TaskStatusTracker | None = None,
         metrics: Any | None = None,
-        submit_web_logs: Callable | None = None, # ADDED
+        submit_web_logs: Callable | None = None,  # ADDED
     ):
         super().__init__()
         self.config = config
@@ -98,7 +98,7 @@ def __init__(
 
         self.metrics = metrics
         self.status_tracker = status_tracker
-        self.submit_web_logs = submit_web_logs # ADDED
+        self.submit_web_logs = submit_web_logs  # ADDED
 
     def on_messages_enqueued(self, msgs: list[ScheduleMessageItem]) -> None:
         if not msgs:
@@ -120,11 +120,13 @@ def _create_task_wrapper(self, handler: Callable, task_item: RunningTaskItem):
         def wrapped_handler(messages: list[ScheduleMessageItem]):
             start_time = time.time()
             if self.status_tracker:
-                self.status_tracker.task_started(task_id=task_item.item_id, user_id=task_item.user_id)
+                self.status_tracker.task_started(
+                    task_id=task_item.item_id, user_id=task_item.user_id
+                )
             try:
                 # --- mark start: record queuing time(now - enqueue_ts)---
                 now = time.time()
-                m = messages[0] # All messages in this batch have same user and type
+                m = messages[0]  # All messages in this batch have same user and type
                 enq_ts = getattr(m, "timestamp", None)
 
                 # Path 1: epoch seconds (preferred)
@@ -145,26 +147,29 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                 wait_sec = max(0.0, now - enq_epoch)
                 self.metrics.observe_task_wait_duration(wait_sec, m.user_id, m.label)
 
-
                 # Execute the original handler
                 result = handler(messages)
 
                 # --- mark done ---
                 duration = time.time() - start_time
                 self.metrics.observe_task_duration(duration, m.user_id, m.label)
                 if self.status_tracker:
-                    self.status_tracker.task_completed(task_id=task_item.item_id, user_id=task_item.user_id)
+                    self.status_tracker.task_completed(
+                        task_id=task_item.item_id, user_id=task_item.user_id
+                    )
                 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'
+                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"
+                        status="completed",
                     )
                     self.submit_web_logs([status_log])
 
@@ -194,7 +199,9 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                 m = messages[0]
                 self.metrics.task_failed(m.user_id, m.label, type(e).__name__)
                 if self.status_tracker:
-                    self.status_tracker.task_failed(task_id=task_item.item_id, user_id=task_item.user_id, error_message=str(e))
+                    self.status_tracker.task_failed(
+                        task_id=task_item.item_id, user_id=task_item.user_id, error_message=str(e)
+                    )
                 # Mark task as failed and remove from tracking
                 with self._task_lock:
                     if task_item.item_id in self._running_tasks:
@@ -204,7 +211,9 @@ 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'
+                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,
@@ -213,7 +222,7 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                         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)
+                        exception=str(e),
                     )
                     self.submit_web_logs([status_log])
                 raise

src/memos/mem_scheduler/utils/metrics.py

@@ -11,31 +11,31 @@
 TASKS_ENQUEUED_TOTAL = Counter(
     "memos_scheduler_tasks_enqueued_total",
     "Total number of tasks enqueued",
-    ["user_id", "task_type"]
+    ["user_id", "task_type"],
 )
 
 TASKS_DEQUEUED_TOTAL = Counter(
     "memos_scheduler_tasks_dequeued_total",
     "Total number of tasks dequeued",
-    ["user_id", "task_type"]
+    ["user_id", "task_type"],
 )
 
 TASK_DURATION_SECONDS = Summary(
     "memos_scheduler_task_duration_seconds",
     "Task processing duration in seconds",
-    ["user_id", "task_type"]
+    ["user_id", "task_type"],
 )
 
 TASK_WAIT_DURATION_SECONDS = Summary(
     "memos_scheduler_task_wait_duration_seconds",
     "Task waiting duration in seconds",
-    ["user_id", "task_type"]
+    ["user_id", "task_type"],
 )
 
 TASKS_FAILED_TOTAL = Counter(
     "memos_scheduler_tasks_failed_total",
     "Total number of failed tasks",
-    ["user_id", "task_type", "error_type"]
+    ["user_id", "task_type", "error_type"],
 )
 
 TASKS_COMPLETED_TOTAL = Counter(
@@ -45,47 +45,56 @@
 )
 
 QUEUE_LENGTH = Gauge(
-    "memos_scheduler_queue_length",
-    "Current length of the task queue",
-    ["user_id"]
+    "memos_scheduler_queue_length", "Current length of the task queue", ["user_id"]
 )
 
 INTERNAL_SPAN_DURATION = Histogram(
-    'memos_scheduler_internal_span_duration_seconds',
-    'Duration of internal operations',
-    ['span_name', 'user_id', 'task_id']
+    "memos_scheduler_internal_span_duration_seconds",
+    "Duration of internal operations",
+    ["span_name", "user_id", "task_id"],
 )
 
 
 # --- Instrumentation Functions ---
 
+
 def task_enqueued(user_id: str, task_type: str, count: int = 1):
     TASKS_ENQUEUED_TOTAL.labels(user_id=user_id, task_type=task_type).inc(count)
 
+
 def task_dequeued(user_id: str, task_type: str, count: int = 1):
     TASKS_DEQUEUED_TOTAL.labels(user_id=user_id, task_type=task_type).inc(count)
 
+
 def observe_task_duration(duration: float, user_id: str, task_type: str):
     TASK_DURATION_SECONDS.labels(user_id=user_id, task_type=task_type).observe(duration)
 
+
 def observe_task_wait_duration(duration: float, user_id: str, task_type: str):
     TASK_WAIT_DURATION_SECONDS.labels(user_id=user_id, task_type=task_type).observe(duration)
 
+
 def task_failed(user_id: str, task_type: str, error_type: str):
     TASKS_FAILED_TOTAL.labels(user_id=user_id, task_type=task_type, error_type=error_type).inc()
 
+
 def task_completed(user_id: str, task_type: str, count: int = 1):
     TASKS_COMPLETED_TOTAL.labels(user_id=user_id, task_type=task_type).inc(count)
 
+
 def update_queue_length(length: int, user_id: str):
     QUEUE_LENGTH.labels(user_id=user_id).set(length)
 
+
 def observe_internal_span(duration: float, span_name: str, user_id: str, task_id: str):
-    INTERNAL_SPAN_DURATION.labels(span_name=span_name, user_id=user_id, task_id=task_id).observe(duration)
+    INTERNAL_SPAN_DURATION.labels(span_name=span_name, user_id=user_id, task_id=task_id).observe(
+        duration
+    )
 
 
 # --- TimingSpan Context Manager ---
 
+
 class TimingSpan(ContextDecorator):
     """
     A context manager/decorator to measure the duration of a code block and record it
@@ -100,6 +109,7 @@ def my_function():
     with TimingSpan("another_op", user_id="user456", task_id="t1"):
         ...
     """
+
     def __init__(self, span_name: str, user_id: str = "unknown", task_id: str = "unknown"):
         self.span_name = span_name
         self.user_id = user_id
@@ -112,4 +122,4 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         duration = time.perf_counter() - self.start_time
-        observe_internal_span(duration, self.span_name, self.user_id, self.task_id)
+        observe_internal_span(duration, self.span_name, self.user_id, self.task_id)

tests/mem_scheduler/test_dispatcher.py

@@ -156,7 +156,9 @@ def test_register_handlers(self):
     def test_dispatch_serial(self):
         """Test dispatching messages in serial mode."""
         # Create a new dispatcher with parallel dispatch disabled
-        serial_dispatcher = SchedulerDispatcher(max_workers=2, enable_parallel_dispatch=False, metrics=MagicMock())
+        serial_dispatcher = SchedulerDispatcher(
+            max_workers=2, enable_parallel_dispatch=False, metrics=MagicMock()
+        )
 
         # Create fresh mock handlers for this test
         mock_handler1 = MagicMock()