Merge branch 'dev' into feat/deep-search

Author: fridayL

examples/mem_reader/multimodel_struct_reader.py

@@ -1,7 +1,12 @@
 import argparse
 import json
+import os
 import time
 
+from typing import Any
+
+from dotenv import load_dotenv
+
 from memos.configs.mem_reader import SimpleStructMemReaderConfig
 from memos.mem_reader.simple_struct import SimpleStructMemReader
 from memos.memories.textual.item import (
@@ -11,6 +16,10 @@
 )
 
 
+# Load environment variables from .env file
+load_dotenv()
+
+
 def print_textual_memory_item(
     item: TextualMemoryItem, max_memory_length: int = 200, indent: int = 0
 ):
@@ -98,6 +107,104 @@ def print_textual_memory_item_json(item: TextualMemoryItem, indent: int = 2):
     print(json.dumps(data, indent=indent, ensure_ascii=False))
 
 
+def get_reader_config() -> dict[str, Any]:
+    """
+    Get reader configuration from environment variables.
+
+    Returns a dictionary that can be used to create SimpleStructMemReaderConfig.
+    Similar to APIConfig.get_reader_config() in server_router_api.py.
+
+    Returns:
+        Configuration dictionary for SimpleStructMemReaderConfig
+    """
+    openai_api_key = os.getenv("OPENAI_API_KEY")
+    openai_base_url = os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1")
+    ollama_api_base = os.getenv("OLLAMA_API_BASE", "http://localhost:11434")
+
+    # Get LLM backend and config
+    llm_backend = os.getenv("MEM_READER_LLM_BACKEND", "openai")
+    if llm_backend == "ollama":
+        llm_config = {
+            "backend": "ollama",
+            "config": {
+                "model_name_or_path": os.getenv("MEM_READER_LLM_MODEL", "qwen3:0.6b"),
+                "api_base": ollama_api_base,
+                "temperature": float(os.getenv("MEM_READER_LLM_TEMPERATURE", "0.0")),
+                "remove_think_prefix": os.getenv(
+                    "MEM_READER_LLM_REMOVE_THINK_PREFIX", "true"
+                ).lower()
+                == "true",
+                "max_tokens": int(os.getenv("MEM_READER_LLM_MAX_TOKENS", "8192")),
+            },
+        }
+    else:  # openai
+        llm_config = {
+            "backend": "openai",
+            "config": {
+                "model_name_or_path": os.getenv("MEM_READER_LLM_MODEL", "gpt-4o-mini"),
+                "api_key": openai_api_key or os.getenv("MEMRADER_API_KEY", "EMPTY"),
+                "api_base": openai_base_url,
+                "temperature": float(os.getenv("MEM_READER_LLM_TEMPERATURE", "0.5")),
+                "remove_think_prefix": os.getenv(
+                    "MEM_READER_LLM_REMOVE_THINK_PREFIX", "true"
+                ).lower()
+                == "true",
+                "max_tokens": int(os.getenv("MEM_READER_LLM_MAX_TOKENS", "8192")),
+            },
+        }
+
+    # Get embedder backend and config
+    embedder_backend = os.getenv(
+        "MEM_READER_EMBEDDER_BACKEND", os.getenv("MOS_EMBEDDER_BACKEND", "ollama")
+    )
+    if embedder_backend == "universal_api":
+        embedder_config = {
+            "backend": "universal_api",
+            "config": {
+                "provider": os.getenv(
+                    "MEM_READER_EMBEDDER_PROVIDER", os.getenv("MOS_EMBEDDER_PROVIDER", "openai")
+                ),
+                "api_key": os.getenv(
+                    "MEM_READER_EMBEDDER_API_KEY",
+                    os.getenv("MOS_EMBEDDER_API_KEY", openai_api_key or "sk-xxxx"),
+                ),
+                "model_name_or_path": os.getenv(
+                    "MEM_READER_EMBEDDER_MODEL",
+                    os.getenv("MOS_EMBEDDER_MODEL", "text-embedding-3-large"),
+                ),
+                "base_url": os.getenv(
+                    "MEM_READER_EMBEDDER_API_BASE",
+                    os.getenv("MOS_EMBEDDER_API_BASE", openai_base_url),
+                ),
+            },
+        }
+    else:  # ollama
+        embedder_config = {
+            "backend": "ollama",
+            "config": {
+                "model_name_or_path": os.getenv(
+                    "MEM_READER_EMBEDDER_MODEL",
+                    os.getenv("MOS_EMBEDDER_MODEL", "nomic-embed-text:latest"),
+                ),
+                "api_base": ollama_api_base,
+            },
+        }
+
+    return {
+        "llm": llm_config,
+        "embedder": embedder_config,
+        "chunker": {
+            "backend": "sentence",
+            "config": {
+                "tokenizer_or_token_counter": "gpt2",
+                "chunk_size": 512,
+                "chunk_overlap": 128,
+                "min_sentences_per_chunk": 1,
+            },
+        },
+    }
+
+
 def main():
     # Parse command line arguments
     parser = argparse.ArgumentParser(description="Test Mem-Reader with structured output")
@@ -115,10 +222,18 @@ def main():
     )
     args = parser.parse_args()
 
-    # 1. Create Configuration
-    reader_config = SimpleStructMemReaderConfig.from_json_file(
-        "examples/data/config/simple_struct_reader_config.json"
-    )
+    # 1. Create Configuration from environment variables or JSON file
+    # Try to get config from environment variables first
+    openai_api_key = os.getenv("OPENAI_API_KEY")
+    if openai_api_key:
+        # Use environment variables (similar to server_router_api.py)
+        config_dict = get_reader_config()
+        reader_config = SimpleStructMemReaderConfig.model_validate(config_dict)
+    else:
+        # Fall back to JSON file
+        reader_config = SimpleStructMemReaderConfig.from_json_file(
+            "examples/data/config/simple_struct_reader_config.json"
+        )
     reader = SimpleStructMemReader(reader_config)
 
     # 2. Define scene data

examples/mem_reader/reader.py …mples/mem_reader/simple_struct_reader.pyexamples/mem_reader/reader.py renamed to examples/mem_reader/simple_struct_reader.py examples/mem_reader/reader.py renamed to examples/mem_reader/simple_struct_reader.py

@@ -894,7 +894,7 @@ def _send_message_to_scheduler(
                 content=query,
                 timestamp=datetime.utcnow(),
             )
-            self.mem_scheduler.memos_message_queue.submit_messages(messages=[message_item])
+            self.mem_scheduler.submit_messages(messages=[message_item])
             self.logger.info(f"Sent message to scheduler with label: {label}")
         except Exception as e:
             self.logger.error(f"Failed to send message to scheduler: {e}", exc_info=True)

src/memos/api/handlers/chat_handler.py

@@ -15,6 +15,7 @@
     MemoryResponse,
 )
 from memos.log import get_logger
+from memos.mem_cube.navie import NaiveMemCube
 from memos.mem_os.utils.format_utils import (
     convert_graph_to_tree_forworkmem,
     ensure_unique_tree_ids,
@@ -162,11 +163,13 @@ def handle_get_subgraph(
         raise
 
 
-def handle_get_memories(get_mem_req: GetMemoryRequest, naive_mem_cube: Any) -> GetMemoryResponse:
+def handle_get_memories(
+    get_mem_req: GetMemoryRequest, naive_mem_cube: NaiveMemCube
+) -> GetMemoryResponse:
     # TODO: Implement get memory with filter
     memories = naive_mem_cube.text_mem.get_all(user_name=get_mem_req.mem_cube_id)["nodes"]
     preferences: list[TextualMemoryItem] = []
-    if get_mem_req.include_preference:
+    if get_mem_req.include_preference and naive_mem_cube.pref_mem is not None:
         filter_params: dict[str, Any] = {}
         if get_mem_req.user_id is not None:
             filter_params["user_id"] = get_mem_req.user_id
@@ -183,10 +186,11 @@ def handle_get_memories(get_mem_req: GetMemoryRequest, naive_mem_cube: Any) -> G
     )
 
 
-def handle_delete_memories(delete_mem_req: DeleteMemoryRequest, naive_mem_cube: Any):
+def handle_delete_memories(delete_mem_req: DeleteMemoryRequest, naive_mem_cube: NaiveMemCube):
     try:
         naive_mem_cube.text_mem.delete(delete_mem_req.memory_ids)
-        naive_mem_cube.pref_mem.delete(delete_mem_req.memory_ids)
+        if naive_mem_cube.pref_mem is not None:
+            naive_mem_cube.pref_mem.delete(delete_mem_req.memory_ids)
     except Exception as e:
         logger.error(f"Failed to delete memories: {e}", exc_info=True)
         return DeleteMemoryResponse(

src/memos/api/handlers/memory_handler.py

@@ -34,7 +34,9 @@ def handle_scheduler_status(
     Args:
         user_id: User ID to query for.
         status_tracker: The TaskStatusTracker instance.
-        task_id: Optional Task ID to query a specific task.
+        task_id: Optional Task ID to query. Can be either:
+                 - business_task_id (will aggregate all related item statuses)
+                 - item_id (will return single item status)
 
     Returns:
         StatusResponse with a list of task statuses.
@@ -46,12 +48,22 @@ def handle_scheduler_status(
 
     try:
         if task_id:
-            task_data = status_tracker.get_task_status(task_id, user_id)
-            if not task_data:
-                raise HTTPException(
-                    status_code=404, detail=f"Task {task_id} not found for user {user_id}"
+            # First try as business_task_id (aggregated query)
+            business_task_data = status_tracker.get_task_status_by_business_id(task_id, user_id)
+            if business_task_data:
+                response_data.append(
+                    StatusResponseItem(task_id=task_id, status=business_task_data["status"])
+                )
+            else:
+                # Fallback: try as item_id (single item query)
+                item_task_data = status_tracker.get_task_status(task_id, user_id)
+                if not item_task_data:
+                    raise HTTPException(
+                        status_code=404, detail=f"Task {task_id} not found for user {user_id}"
+                    )
+                response_data.append(
+                    StatusResponseItem(task_id=task_id, status=item_task_data["status"])
                 )
-            response_data.append(StatusResponseItem(task_id=task_id, status=task_data["status"]))
         else:
             all_tasks = status_tracker.get_all_tasks_for_user(user_id)
             # The plan returns an empty list, which is good.

src/memos/api/handlers/scheduler_handler.py

@@ -258,6 +258,7 @@ class MemoryCreateRequest(BaseRequest):
     source: str | None = Field(None, description="Source of the memory")
     user_profile: bool = Field(False, description="User profile memory")
     session_id: str | None = Field(None, description="Session id")
+    task_id: str | None = Field(None, description="Task ID for monitoring async tasks")
 
 
 class SearchRequest(BaseRequest):

src/memos/api/product_models.py

@@ -188,9 +188,43 @@ def get_all_memories(memory_req: GetMemoryPlaygroundRequest):
 @router.post("/add", summary="add a new memory", response_model=SimpleResponse)
 def create_memory(memory_req: MemoryCreateRequest):
     """Create a new memory for a specific user."""
+    # Initialize status_tracker outside try block to avoid NameError in except blocks
+    status_tracker = None
+
     try:
         time_start_add = time.time()
         mos_product = get_mos_product_instance()
+
+        # Track task if task_id is provided
+        item_id: str | None = None
+        if (
+            memory_req.task_id
+            and hasattr(mos_product, "mem_scheduler")
+            and mos_product.mem_scheduler
+        ):
+            from uuid import uuid4
+
+            from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
+
+            item_id = str(uuid4())  # Generate a unique item_id for this submission
+
+            # Get Redis client from scheduler
+            if (
+                hasattr(mos_product.mem_scheduler, "redis_client")
+                and mos_product.mem_scheduler.redis_client
+            ):
+                status_tracker = TaskStatusTracker(mos_product.mem_scheduler.redis_client)
+                # Submit task with "product_add" type
+                status_tracker.task_submitted(
+                    task_id=item_id,  # Use generated item_id for internal tracking
+                    user_id=memory_req.user_id,
+                    task_type="product_add",
+                    mem_cube_id=memory_req.mem_cube_id or memory_req.user_id,
+                    business_task_id=memory_req.task_id,  # Use memory_req.task_id as business_task_id
+                )
+                status_tracker.task_started(item_id, memory_req.user_id)  # Use item_id here
+
+        # Execute the add operation
         mos_product.add(
             user_id=memory_req.user_id,
             memory_content=memory_req.memory_content,
@@ -200,15 +234,27 @@ def create_memory(memory_req: MemoryCreateRequest):
             source=memory_req.source,
             user_profile=memory_req.user_profile,
             session_id=memory_req.session_id,
+            task_id=memory_req.task_id,
         )
+
+        # Mark task as completed
+        if status_tracker and item_id:
+            status_tracker.task_completed(item_id, memory_req.user_id)
+
         logger.info(
             f"time add api : add time user_id: {memory_req.user_id} time is: {time.time() - time_start_add}"
         )
         return SimpleResponse(message="Memory created successfully")
 
     except ValueError as err:
+        # Mark task as failed if tracking
+        if status_tracker and item_id:
+            status_tracker.task_failed(item_id, memory_req.user_id, str(err))
         raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err
     except Exception as err:
+        # Mark task as failed if tracking
+        if status_tracker and item_id:
+            status_tracker.task_failed(item_id, memory_req.user_id, str(err))
         logger.error(f"Failed to create memory: {traceback.format_exc()}")
         raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err