Feat/redis scheduler

docker/requirements.txt

@@ -158,3 +158,4 @@ watchfiles==1.1.0
 websockets==15.0.1
 xlrd==2.0.2
 xlsxwriter==3.2.5
+prometheus-client==0.23.1

evaluation/scripts/locomo/locomo_eval.py

@@ -3,6 +3,7 @@
 import json
 import logging
 import os
+import re
 import time
 
 import nltk
@@ -47,6 +48,29 @@ class LLMGrade(BaseModel):
     llm_reasoning: str = Field(description="Explain why the answer is correct or incorrect.")
 
 
+def extract_label_json(text: str) -> str | None:
+    """
+    Extracts a JSON object of the form {"label": "VALUE"} from a given text string.
+    This function is designed to handle cases where the LLM response contains
+    natural language alongside a final JSON snippet, ensuring robust parsing.
+
+    Supports both single and double quotes around the label value.
+    Ignores surrounding whitespace and formatting.
+
+    Returns:
+        The full matching JSON string (e.g., '{"label": "CORRECT"}') if found.
+        None if no valid label JSON is found.
+    """
+    # Regex pattern to match: { "label": "value" } with optional whitespace
+    # Matches both single and double quotes, allows spaces around keys and values
+    pattern = r'\{\s*"label"\s*:\s*["\']([^"\']*)["\']\s*\}'
+    match = re.search(pattern, text)
+    if match:
+        # Return the complete matched JSON string for safe json.loads()
+        return match.group(0)
+    return None
+
+
 async def locomo_grader(llm_client, question: str, gold_answer: str, response: str) -> bool:
     system_prompt = """
         You are an expert grader that determines if answers to questions match a gold standard answer
@@ -77,20 +101,23 @@ async def locomo_grader(llm_client, question: str, gold_answer: str, response: s
 
     Just return the label CORRECT or WRONG in a json format with the key as "label".
     """
-
-    response = await llm_client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=[
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": accuracy_prompt},
-        ],
-        temperature=0,
-    )
-    message_content = response.choices[0].message.content
-    label = json.loads(message_content)["label"]
-    parsed = LLMGrade(llm_judgment=label, llm_reasoning="")
-
-    return parsed.llm_judgment.strip().lower() == "correct"
+    try:
+        response = await llm_client.chat.completions.create(
+            model=os.getenv("EVAL_MODEL", "gpt-4o-mini"),
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": accuracy_prompt},
+            ],
+            temperature=0,
+        )
+        message_content = response.choices[0].message.content
+        message_content = extract_label_json(text=message_content)
+        label = json.loads(message_content)["label"]
+        parsed = LLMGrade(llm_judgment=label, llm_reasoning="")
+        return parsed.llm_judgment.strip().lower() == "correct"
+    except Exception as e:
+        print(f"======== {e}, {response} ===========")
+        exit()
 
 
 def calculate_rouge_scores(gold_answer, response):
@@ -284,7 +311,7 @@ async def main(frame, version="default", options=None, num_runs=1, max_workers=4
     with open(response_path) as file:
         locomo_responses = json.load(file)
 
-    num_users = 10
+    num_users = 2
     all_grades = {}
 
     total_responses_count = sum(

evaluation/scripts/utils/client.py

@@ -189,9 +189,7 @@ def search(self, query, user_id, top_k):
         )
         response = requests.request("POST", url, data=payload, headers=self.headers)
         assert response.status_code == 200, response.text
-        assert json.loads(response.text)["message"] == "Search completed successfully", (
-            response.text
-        )
+        assert json.loads(response.text)["message"] == "Memory searched successfully", response.text
         return json.loads(response.text)["data"]
 
 

examples/mem_scheduler/api_w_scheduler.py

@@ -1,8 +1,10 @@
+from time import sleep
+
 from memos.api.handlers.scheduler_handler import (
     handle_scheduler_status,
     handle_scheduler_wait,
 )
-from memos.api.routers.server_router import mem_scheduler
+from memos.api.routers.server_router import mem_scheduler, status_tracker
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
@@ -26,26 +28,25 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
     for msg in messages:
         print(f" my_test_handler - {msg.item_id}: {msg.content}")
         user_status_running = handle_scheduler_status(
-            user_name=USER_MEM_CUBE, mem_scheduler=mem_scheduler, instance_id="api_w_scheduler"
+            user_id=msg.user_id, status_tracker=status_tracker
         )
-        print(f"[Monitor] Status for {USER_MEM_CUBE} after submit:", user_status_running)
+        print("[Monitor] Status after submit:", user_status_running)
 
 
 # 2. Register the handler
 TEST_HANDLER_LABEL = "test_handler"
+TEST_USER_ID = "test_user"
 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_name=None, mem_scheduler=mem_scheduler, instance_id="api_w_scheduler"
-)
+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
 messages_to_send = [
     ScheduleMessageItem(
         item_id=f"test_item_{i}",
-        user_id="test_user",
+        user_id=TEST_USER_ID,
         mem_cube_id="test_mem_cube",
         label=TEST_HANDLER_LABEL,
         content=f"This is test message {i}",
@@ -56,28 +57,28 @@ def my_test_handler(messages: list[ScheduleMessageItem]):
 # 5. Submit messages
 for mes in messages_to_send:
     print(f"Submitting message {mes.item_id} to the scheduler...")
-    mem_scheduler.memos_message_queue.submit_messages([mes])
+    mem_scheduler.submit_messages([mes])
+    sleep(1)
 
 # 5.1 Monitor status for specific mem_cube while running
 USER_MEM_CUBE = "test_mem_cube"
 
 # 6. Wait for messages to be processed (limited to 100 checks)
-print("Waiting for messages to be consumed (max 100 checks)...")
-mem_scheduler.mem_scheduler_wait()
+
+user_status_running = handle_scheduler_status(user_id=TEST_USER_ID, status_tracker=status_tracker)
+print(f"[Monitor] Status for {USER_MEM_CUBE} after submit:", user_status_running)
 
 # 6.1 Wait until idle for specific mem_cube via handler
 wait_result = handle_scheduler_wait(
-    user_name=USER_MEM_CUBE,
+    user_name=TEST_USER_ID,
+    status_tracker=status_tracker,
     timeout_seconds=120.0,
-    poll_interval=0.2,
-    mem_scheduler=mem_scheduler,
+    poll_interval=0.5,
 )
 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_name=None, mem_scheduler=mem_scheduler, instance_id="api_w_scheduler"
-)
+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

pyproject.toml

@@ -46,6 +46,7 @@ dependencies = [
     "scikit-learn (>=1.7.0,<2.0.0)",  # Machine learning
     "fastmcp (>=2.10.5,<3.0.0)",
     "python-dateutil (>=2.9.0.post0,<3.0.0)",
+    "prometheus-client (>=0.23.1,<0.24.0)",
 ]
 
 [project.urls]

src/memos/api/handlers/base_handler.py

@@ -9,7 +9,7 @@
 
 from memos.log import get_logger
 from memos.mem_scheduler.base_scheduler import BaseScheduler
-from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
+from memos.memories.textual.tree_text_memory.retrieve.advanced_searcher import AdvancedSearcher
 
 
 logger = get_logger(__name__)
@@ -132,7 +132,7 @@ def mem_scheduler(self) -> BaseScheduler:
         return self.deps.mem_scheduler
 
     @property
-    def searcher(self) -> Searcher:
+    def searcher(self) -> AdvancedSearcher:
         """Get scheduler instance."""
         return self.deps.searcher
 

src/memos/api/handlers/component_init.py

@@ -129,6 +129,21 @@ def init_server() -> dict[str, Any]:
     """
     logger.info("Initializing MemOS server components...")
 
+    # 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."
+            )
+    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
+
     # Get default cube configuration
     default_cube_config = APIConfig.get_default_cube_config()
 
@@ -272,6 +287,8 @@ def init_server() -> dict[str, Any]:
     tree_mem: TreeTextMemory = naive_mem_cube.text_mem
     searcher: Searcher = tree_mem.get_searcher(
         manual_close_internet=os.getenv("ENABLE_INTERNET", "true").lower() == "false",
+        moscube=False,
+        process_llm=mem_reader.llm,
     )
     logger.debug("Searcher created")
 
@@ -286,6 +303,7 @@ def init_server() -> dict[str, Any]:
         process_llm=mem_reader.llm,
         db_engine=BaseDBManager.create_default_sqlite_engine(),
         mem_reader=mem_reader,
+        redis_client=redis_client,
     )
     mem_scheduler.init_mem_cube(mem_cube=naive_mem_cube, searcher=searcher)
     logger.debug("Scheduler initialized")
@@ -335,5 +353,6 @@ def init_server() -> dict[str, Any]:
         "text_mem": text_mem,
         "pref_mem": pref_mem,
         "online_bot": online_bot,
+        "redis_client": redis_client,
         "deepsearch_agent": deepsearch_agent,
     }