:Merge branch 'test' into dev_mcp

Author: WenjinW

src/memos/api/routers/server_router.py

@@ -243,6 +243,57 @@ def search_memories(search_req: APISearchRequest):
     )
 
 
+@router.post("/search_ws", summary="Search memories with scheduler", response_model=SearchResponse)
+def search_memories_ws(search_req: APISearchRequest):
+    """Search memories for a specific user."""
+    # Create UserContext object - how to assign values
+    user_context = UserContext(
+        user_id=search_req.user_id,
+        mem_cube_id=search_req.mem_cube_id,
+        session_id=search_req.session_id or "default_session",
+    )
+    logger.info(f"Search user_id is: {user_context.mem_cube_id}")
+    memories_result: MOSSearchResult = {
+        "text_mem": [],
+        "act_mem": [],
+        "para_mem": [],
+    }
+    target_session_id = search_req.session_id
+    if not target_session_id:
+        target_session_id = "default_session"
+    search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
+
+    # Create MemCube and perform search
+    naive_mem_cube = _create_naive_mem_cube()
+    search_results = naive_mem_cube.text_mem.search(
+        query=search_req.query,
+        user_name=user_context.mem_cube_id,
+        top_k=search_req.top_k,
+        mode=search_req.mode,
+        manual_close_internet=not search_req.internet_search,
+        moscube=search_req.moscube,
+        search_filter=search_filter,
+        info={
+            "user_id": search_req.user_id,
+            "session_id": target_session_id,
+            "chat_history": search_req.chat_history,
+        },
+    )
+    formatted_memories = [_format_memory_item(data) for data in search_results]
+
+    memories_result["text_mem"].append(
+        {
+            "cube_id": search_req.mem_cube_id,
+            "memories": formatted_memories,
+        }
+    )
+
+    return SearchResponse(
+        message="Search completed successfully",
+        data=memories_result,
+    )
+
+
 @router.post("/add", summary="Add memories", response_model=MemoryResponse)
 def add_memories(add_req: APIADDRequest):
     """Add memories for a specific user."""

src/memos/configs/mem_scheduler.py

@@ -28,13 +28,11 @@ class BaseSchedulerConfig(BaseConfig):
     thread_pool_max_workers: int = Field(
         default=DEFAULT_THREAD_POOL_MAX_WORKERS,
         gt=1,
-        lt=20,
         description=f"Maximum worker threads in pool (default: {DEFAULT_THREAD_POOL_MAX_WORKERS})",
     )
     consume_interval_seconds: float = Field(
         default=DEFAULT_CONSUME_INTERVAL_SECONDS,
         gt=0,
-        le=60,
         description=f"Interval for consuming messages from queue in seconds (default: {DEFAULT_CONSUME_INTERVAL_SECONDS})",
     )
     auth_config_path: str | None = Field(

src/memos/llms/hf.py

@@ -379,10 +379,52 @@ def build_kv_cache(self, messages) -> DynamicCache:
             raise ValueError(
                 "Prompt after chat template is empty, cannot build KV cache. Check your messages input."
             )
-        kv = DynamicCache()
+        # Create cache and perform forward pass without pre-existing cache
         with torch.no_grad():
-            self.model(**inputs, use_cache=True, past_key_values=kv)
-        for i, (k, v) in enumerate(zip(kv.key_cache, kv.value_cache, strict=False)):
-            kv.key_cache[i] = k[:, :, :seq_len, :]
-            kv.value_cache[i] = v[:, :, :seq_len, :]
-        return kv
+            outputs = self.model(**inputs, use_cache=True)
+
+        # Get the cache from model outputs
+        if hasattr(outputs, "past_key_values") and outputs.past_key_values is not None:
+            kv = outputs.past_key_values
+
+            # Convert from legacy tuple format to DynamicCache if needed
+            if isinstance(kv, tuple):
+                kv = DynamicCache.from_legacy_cache(kv)
+
+            # Handle compatibility between old and new transformers versions
+            # In newer versions, DynamicCache uses 'layers' attribute
+            # In older versions, it uses 'key_cache' and 'value_cache' attributes
+            if hasattr(kv, "layers"):
+                # New version: trim cache using layers attribute
+                for layer in kv.layers:
+                    if hasattr(layer, "key_cache") and hasattr(layer, "value_cache"):
+                        # Trim each layer's cache to the sequence length
+                        if layer.key_cache is not None:
+                            layer.key_cache = layer.key_cache[:, :, :seq_len, :]
+                        if layer.value_cache is not None:
+                            layer.value_cache = layer.value_cache[:, :, :seq_len, :]
+                    elif hasattr(layer, "keys") and hasattr(layer, "values"):
+                        # Alternative attribute names in some versions
+                        if layer.keys is not None:
+                            layer.keys = layer.keys[:, :, :seq_len, :]
+                        if layer.values is not None:
+                            layer.values = layer.values[:, :, :seq_len, :]
+            elif hasattr(kv, "key_cache") and hasattr(kv, "value_cache"):
+                # Old version: trim cache using key_cache and value_cache attributes
+                for i in range(len(kv.key_cache)):
+                    if kv.key_cache[i] is not None:
+                        kv.key_cache[i] = kv.key_cache[i][:, :, :seq_len, :]
+                    if kv.value_cache[i] is not None:
+                        kv.value_cache[i] = kv.value_cache[i][:, :, :seq_len, :]
+            else:
+                # Fallback: log warning but continue without trimming
+                logger.warning(
+                    f"DynamicCache object of type {type(kv)} has unexpected structure. "
+                    f"Cache trimming skipped. Available attributes: {dir(kv)}"
+                )
+
+            return kv
+        else:
+            raise RuntimeError(
+                "Failed to build KV cache: no cache data available from model outputs"
+            )

src/memos/mem_os/core.py

@@ -308,18 +308,20 @@ def chat(self, query: str, user_id: str | None = None, base_prompt: str | None =
         past_key_values = None
 
         if self.config.enable_activation_memory:
-            assert self.config.chat_model.backend == "huggingface", (
-                "Activation memory only used for huggingface backend."
-            )
-            # TODO this only one cubes
-            for mem_cube_id, mem_cube in self.mem_cubes.items():
-                if mem_cube_id not in user_cube_ids:
-                    continue
-                if mem_cube.act_mem:
-                    kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
-                    past_key_values = (
-                        kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
-                    )
+            if self.config.chat_model.backend != "huggingface":
+                logger.error(
+                    "Activation memory only used for huggingface backend. Skipping activation memory."
+                )
+            else:
+                # TODO this only one cubes
+                for mem_cube_id, mem_cube in self.mem_cubes.items():
+                    if mem_cube_id not in user_cube_ids:
+                        continue
+                    if mem_cube.act_mem:
+                        kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
+                        past_key_values = (
+                            kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
+                        )
                     break
             # Generate response
             response = self.chat_llm.generate(current_messages, past_key_values=past_key_values)

src/memos/mem_os/main.py

@@ -312,23 +312,25 @@ def _generate_enhanced_response_with_context(
         # Handle activation memory if enabled (same as core method)
         past_key_values = None
         if self.config.enable_activation_memory:
-            assert self.config.chat_model.backend == "huggingface", (
-                "Activation memory only used for huggingface backend."
-            )
-            # Get accessible cubes for the user
-            target_user_id = user_id if user_id is not None else self.user_id
-            accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
-            user_cube_ids = [cube.cube_id for cube in accessible_cubes]
-
-            for mem_cube_id, mem_cube in self.mem_cubes.items():
-                if mem_cube_id not in user_cube_ids:
-                    continue
-                if mem_cube.act_mem:
-                    kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
-                    past_key_values = (
-                        kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
-                    )
-                    break
+            if self.config.chat_model.backend != "huggingface":
+                logger.error(
+                    "Activation memory only used for huggingface backend. Skipping activation memory."
+                )
+            else:
+                # Get accessible cubes for the user
+                target_user_id = user_id if user_id is not None else self.user_id
+                accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
+                user_cube_ids = [cube.cube_id for cube in accessible_cubes]
+
+                for mem_cube_id, mem_cube in self.mem_cubes.items():
+                    if mem_cube_id not in user_cube_ids:
+                        continue
+                    if mem_cube.act_mem:
+                        kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
+                        past_key_values = (
+                            kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
+                        )
+                        break
 
         try:
             # Generate the enhanced response using the chat LLM with same parameters as core