MemoryOS 2.0.3__py3-none-any.whl

memos/api/handlers/search_handler.py

@@ -0,0 +1,222 @@
+"""
+Search handler for memory search functionality (Class-based version).
+
+This module provides a class-based implementation of search handlers,
+using dependency injection for better modularity and testability.
+"""
+
+import time
+
+from typing import Any
+
+from memos.api.handlers.base_handler import BaseHandler, HandlerDependencies
+from memos.api.handlers.formatters_handler import rerank_knowledge_mem
+from memos.api.product_models import APISearchRequest, SearchResponse
+from memos.log import get_logger
+from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import (
+    cosine_similarity_matrix,
+)
+from memos.multi_mem_cube.composite_cube import CompositeCubeView
+from memos.multi_mem_cube.single_cube import SingleCubeView
+from memos.multi_mem_cube.views import MemCubeView
+
+
+logger = get_logger(__name__)
+
+
+class SearchHandler(BaseHandler):
+    """
+    Handler for memory search operations.
+
+    Provides fast, fine-grained, and mixture-based search modes.
+    """
+
+    def __init__(self, dependencies: HandlerDependencies):
+        """
+        Initialize search handler.
+
+        Args:
+            dependencies: HandlerDependencies instance
+        """
+        super().__init__(dependencies)
+        self._validate_dependencies(
+            "naive_mem_cube", "mem_scheduler", "searcher", "deepsearch_agent"
+        )
+
+    def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse:
+        """
+        Main handler for search memories endpoint.
+
+        Orchestrates the search process based on the requested search mode,
+        supporting both text and preference memory searches.
+
+        Args:
+            search_req: Search request containing query and parameters
+
+        Returns:
+            SearchResponse with formatted results
+        """
+        self.logger.info(f"[SearchHandler] Search Req is: {search_req}")
+
+        # Increase recall pool if deduplication is enabled to ensure diversity
+        original_top_k = search_req.top_k
+        if search_req.dedup == "sim":
+            search_req.top_k = original_top_k * 5
+
+        cube_view = self._build_cube_view(search_req)
+
+        results = cube_view.search_memories(search_req)
+        if search_req.dedup == "sim":
+            results = self._dedup_text_memories(results, original_top_k)
+            self._strip_embeddings(results)
+            # Restore original top_k for downstream logic or response metadata
+            search_req.top_k = original_top_k
+
+        start_time = time.time()
+        text_mem = results["text_mem"]
+        results["text_mem"] = rerank_knowledge_mem(
+            self.reranker,
+            query=search_req.query,
+            text_mem=text_mem,
+            top_k=original_top_k,
+            file_mem_proportion=0.5,
+        )
+        rerank_time = time.time() - start_time
+
+        self.logger.info(f"[Knowledge_replace_memory_time] Rerank time: {rerank_time} seconds")
+        self.logger.info(
+            f"[SearchHandler] Final search results: count={len(results)} results={results}"
+        )
+
+        return SearchResponse(
+            message="Search completed successfully",
+            data=results,
+        )
+
+    def _dedup_text_memories(self, results: dict[str, Any], target_top_k: int) -> dict[str, Any]:
+        buckets = results.get("text_mem", [])
+        if not buckets:
+            return results
+
+        flat: list[tuple[int, dict[str, Any], float]] = []
+        for bucket_idx, bucket in enumerate(buckets):
+            for mem in bucket.get("memories", []):
+                score = mem.get("metadata", {}).get("relativity", 0.0)
+                flat.append((bucket_idx, mem, score))
+
+        if len(flat) <= 1:
+            return results
+
+        embeddings = self._extract_embeddings([mem for _, mem, _ in flat])
+        if embeddings is None:
+            documents = [mem.get("memory", "") for _, mem, _ in flat]
+            embeddings = self.searcher.embedder.embed(documents)
+
+        similarity_matrix = cosine_similarity_matrix(embeddings)
+
+        indices_by_bucket: dict[int, list[int]] = {i: [] for i in range(len(buckets))}
+        for flat_index, (bucket_idx, _, _) in enumerate(flat):
+            indices_by_bucket[bucket_idx].append(flat_index)
+
+        selected_global: list[int] = []
+        selected_by_bucket: dict[int, list[int]] = {i: [] for i in range(len(buckets))}
+
+        ordered_indices = sorted(range(len(flat)), key=lambda idx: flat[idx][2], reverse=True)
+        for idx in ordered_indices:
+            bucket_idx = flat[idx][0]
+            if len(selected_by_bucket[bucket_idx]) >= target_top_k:
+                continue
+            # Use 0.92 threshold strictly
+            if self._is_unrelated(idx, selected_global, similarity_matrix, 0.92):
+                selected_by_bucket[bucket_idx].append(idx)
+                selected_global.append(idx)
+
+        # Removed the 'filling' logic that was pulling back similar items.
+        # Now it will only return items that truly pass the 0.92 threshold,
+        # up to target_top_k.
+
+        for bucket_idx, bucket in enumerate(buckets):
+            selected_indices = selected_by_bucket.get(bucket_idx, [])
+            bucket["memories"] = [flat[i][1] for i in selected_indices]
+        return results
+
+    @staticmethod
+    def _is_unrelated(
+        index: int,
+        selected_indices: list[int],
+        similarity_matrix: list[list[float]],
+        similarity_threshold: float,
+    ) -> bool:
+        return all(similarity_matrix[index][j] <= similarity_threshold for j in selected_indices)
+
+    @staticmethod
+    def _max_similarity(
+        index: int, selected_indices: list[int], similarity_matrix: list[list[float]]
+    ) -> float:
+        if not selected_indices:
+            return 0.0
+        return max(similarity_matrix[index][j] for j in selected_indices)
+
+    @staticmethod
+    def _extract_embeddings(memories: list[dict[str, Any]]) -> list[list[float]] | None:
+        embeddings: list[list[float]] = []
+        for mem in memories:
+            embedding = mem.get("metadata", {}).get("embedding")
+            if not embedding:
+                return None
+            embeddings.append(embedding)
+        return embeddings
+
+    @staticmethod
+    def _strip_embeddings(results: dict[str, Any]) -> None:
+        for bucket in results.get("text_mem", []):
+            for mem in bucket.get("memories", []):
+                metadata = mem.get("metadata", {})
+                if "embedding" in metadata:
+                    metadata["embedding"] = []
+        for bucket in results.get("tool_mem", []):
+            for mem in bucket.get("memories", []):
+                metadata = mem.get("metadata", {})
+                if "embedding" in metadata:
+                    metadata["embedding"] = []
+
+    def _resolve_cube_ids(self, search_req: APISearchRequest) -> list[str]:
+        """
+        Normalize target cube ids from search_req.
+        Priority:
+        1) readable_cube_ids (deprecated mem_cube_id is converted to this in model validator)
+        2) fallback to user_id
+        """
+        if search_req.readable_cube_ids:
+            return list(dict.fromkeys(search_req.readable_cube_ids))
+
+        return [search_req.user_id]
+
+    def _build_cube_view(self, search_req: APISearchRequest) -> MemCubeView:
+        cube_ids = self._resolve_cube_ids(search_req)
+
+        if len(cube_ids) == 1:
+            cube_id = cube_ids[0]
+            return SingleCubeView(
+                cube_id=cube_id,
+                naive_mem_cube=self.naive_mem_cube,
+                mem_reader=self.mem_reader,
+                mem_scheduler=self.mem_scheduler,
+                logger=self.logger,
+                searcher=self.searcher,
+                deepsearch_agent=self.deepsearch_agent,
+            )
+        else:
+            single_views = [
+                SingleCubeView(
+                    cube_id=cube_id,
+                    naive_mem_cube=self.naive_mem_cube,
+                    mem_reader=self.mem_reader,
+                    mem_scheduler=self.mem_scheduler,
+                    logger=self.logger,
+                    searcher=self.searcher,
+                    deepsearch_agent=self.deepsearch_agent,
+                )
+                for cube_id in cube_ids
+            ]
+            return CompositeCubeView(cube_views=single_views, logger=self.logger)

memos/api/handlers/suggestion_handler.py

@@ -0,0 +1,117 @@
+"""
+Suggestion handler for generating suggestion queries.
+
+This module handles suggestion query generation based on user's recent memories
+or further suggestions from chat history.
+"""
+
+import json
+
+from typing import Any
+
+from memos.api.product_models import SuggestionResponse
+from memos.log import get_logger
+from memos.mem_os.utils.format_utils import clean_json_response
+from memos.templates.mos_prompts import (
+    FURTHER_SUGGESTION_PROMPT,
+    SUGGESTION_QUERY_PROMPT_EN,
+    SUGGESTION_QUERY_PROMPT_ZH,
+)
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+def _get_further_suggestion(
+    llm: Any,
+    message: MessageList,
+) -> list[str]:
+    """
+    Get further suggestion based on recent dialogue.
+
+    Args:
+        llm: LLM instance for generating suggestions
+        message: Recent chat messages
+
+    Returns:
+        List of suggestion queries
+    """
+    try:
+        dialogue_info = "\n".join([f"{msg['role']}: {msg['content']}" for msg in message[-2:]])
+        further_suggestion_prompt = FURTHER_SUGGESTION_PROMPT.format(dialogue=dialogue_info)
+        message_list = [{"role": "system", "content": further_suggestion_prompt}]
+        response = llm.generate(message_list)
+        clean_response = clean_json_response(response)
+        response_json = json.loads(clean_response)
+        return response_json["query"]
+    except Exception as e:
+        logger.error(f"Error getting further suggestion: {e}", exc_info=True)
+        return []
+
+
+def handle_get_suggestion_queries(
+    user_id: str,
+    language: str,
+    message: MessageList | None,
+    llm: Any,
+    naive_mem_cube: Any,
+) -> SuggestionResponse:
+    """
+    Main handler for suggestion queries endpoint.
+
+    Generates suggestion queries based on user's recent memories or chat history.
+
+    Args:
+        user_id: User ID
+        language: Language preference ("zh" or "en")
+        message: Optional chat message list for further suggestions
+        llm: LLM instance
+        naive_mem_cube: Memory cube instance
+
+    Returns:
+        SuggestionResponse with generated queries
+    """
+    try:
+        # If message is provided, get further suggestions based on dialogue
+        if message:
+            suggestions = _get_further_suggestion(llm, message)
+            return SuggestionResponse(
+                message="Suggestions retrieved successfully",
+                data={"query": suggestions},
+            )
+
+        # Otherwise, generate suggestions based on recent memories
+        if language == "zh":
+            suggestion_prompt = SUGGESTION_QUERY_PROMPT_ZH
+        else:  # English
+            suggestion_prompt = SUGGESTION_QUERY_PROMPT_EN
+
+        # Search for recent memories
+        text_mem_results = naive_mem_cube.text_mem.search(
+            query="my recently memories",
+            user_name=user_id,
+            top_k=3,
+            mode="fast",
+            info={"user_id": user_id},
+        )
+
+        # Extract memory content
+        memories = ""
+        if text_mem_results:
+            memories = "\n".join([m.memory[:200] for m in text_mem_results])
+
+        # Generate suggestions using LLM
+        message_list = [{"role": "system", "content": suggestion_prompt.format(memories=memories)}]
+        response = llm.generate(message_list)
+        clean_response = clean_json_response(response)
+        response_json = json.loads(clean_response)
+
+        return SuggestionResponse(
+            message="Suggestions retrieved successfully",
+            data={"query": response_json["query"]},
+        )
+
+    except Exception as e:
+        logger.error(f"Failed to get suggestions: {e}", exc_info=True)
+        raise