MemoryOS 0.2.2__py3-none-any.whl → 1.0.1__py3-none-any.whl

memos/memories/textual/tree_text_memory/retrieve/searcher.py

@@ -8,11 +8,12 @@ from memos.graph_dbs.factory import Neo4jGraphDB
 from memos.llms.factory import AzureLLM, OllamaLLM, OpenAILLM
 from memos.log import get_logger
 from memos.memories.textual.item import SearchedTreeNodeTextualMemoryMetadata, TextualMemoryItem
+from memos.reranker.base import BaseReranker
+from memos.utils import timed
 
 from .internet_retriever_factory import InternetRetrieverFactory
 from .reasoner import MemoryReasoner
 from .recall import GraphMemoryRetriever
-from .reranker import MemoryReranker
 from .task_goal_parser import TaskGoalParser
 
 
@@ -25,21 +26,29 @@ class Searcher:
         dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM,
         graph_store: Neo4jGraphDB,
         embedder: OllamaEmbedder,
+        reranker: BaseReranker,
         internet_retriever: InternetRetrieverFactory | None = None,
+        moscube: bool = False,
     ):
         self.graph_store = graph_store
         self.embedder = embedder
 
         self.task_goal_parser = TaskGoalParser(dispatcher_llm)
         self.graph_retriever = GraphMemoryRetriever(self.graph_store, self.embedder)
-        self.reranker = MemoryReranker(dispatcher_llm, self.embedder)
+        self.reranker = reranker
         self.reasoner = MemoryReasoner(dispatcher_llm)
 
         # Create internet retriever from config if provided
         self.internet_retriever = internet_retriever
+        self.moscube = moscube
 
+        self._usage_executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=4, thread_name_prefix="usage"
+        )
+
+    @timed
     def search(
-        self, query: str, top_k: int, info=None, mode: str = "fast", memory_type: str = "All"
+        self, query: str, top_k: int, info=None, mode="fast", memory_type="All"
     ) -> list[TextualMemoryItem]:
         """
         Search for memories based on a query.
@@ -57,25 +66,60 @@ class Searcher:
         Returns:
             list[TextualMemoryItem]: List of matching memories.
         """
+        logger.info(
+            f"[SEARCH] Start query='{query}', top_k={top_k}, mode={mode}, memory_type={memory_type}"
+        )
         if not info:
             logger.warning(
                 "Please input 'info' when use tree.search so that "
                 "the database would store the consume history."
             )
             info = {"user_id": "", "session_id": ""}
-        # Step 1: Parse task structure into topic, concept, and fact levels
+        else:
+            logger.debug(f"[SEARCH] Received info dict: {info}")
+
+        parsed_goal, query_embedding, context, query = self._parse_task(query, info, mode)
+        results = self._retrieve_paths(
+            query, parsed_goal, query_embedding, info, top_k, mode, memory_type
+        )
+        deduped = self._deduplicate_results(results)
+        final_results = self._sort_and_trim(deduped, top_k)
+        self._update_usage_history(final_results, info)
+
+        logger.info(f"[SEARCH] Done. Total {len(final_results)} results.")
+        res_results = ""
+        for _num_i, result in enumerate(final_results):
+            res_results += "\n" + (
+                result.id + "|" + result.metadata.memory_type + "|" + result.memory
+            )
+        logger.info(f"[SEARCH] Results. {res_results}")
+        return final_results
+
+    @timed
+    def _parse_task(self, query, info, mode, top_k=5):
+        """Parse user query, do embedding search and create context"""
         context = []
+        query_embedding = None
+
+        # fine mode will trigger initial embedding search
         if mode == "fine":
+            logger.info("[SEARCH] Fine mode: embedding search")
             query_embedding = self.embedder.embed([query])[0]
-            related_node_ids = self.graph_store.search_by_embedding(query_embedding, top_k=top_k)
+
+            # retrieve related nodes by embedding
             related_nodes = [
-                self.graph_store.get_node(related_node["id"]) for related_node in related_node_ids
+                self.graph_store.get_node(n["id"])
+                for n in self.graph_store.search_by_embedding(query_embedding, top_k=top_k)
             ]
+            context = list({node["memory"] for node in related_nodes})
 
-            context = [related_node["memory"] for related_node in related_nodes]
-            context = list(set(context))
+            # optional: supplement context with internet knowledge
+            """if self.internet_retriever:
+                extra = self.internet_retriever.retrieve_from_internet(query=query, top_k=3)
+                context.extend(item.memory.partition("\nContent: ")[-1] for item in extra)
+            """
 
-        # Step 1a: Parse task structure into topic, concept, and fact levels
+        # parse goal using LLM
         parsed_goal = self.task_goal_parser.parse(
             task_description=query,
             context="\n".join(context),
@@ -83,150 +127,218 @@ class Searcher:
             mode=mode,
         )
 
-        query = (
-            parsed_goal.rephrased_query
-            if parsed_goal.rephrased_query and len(parsed_goal.rephrased_query) > 0
-            else query
-        )
-
+        query = parsed_goal.rephrased_query or query
+        # if goal has extra memories, embed them too
         if parsed_goal.memories:
             query_embedding = self.embedder.embed(list({query, *parsed_goal.memories}))
 
-        # Step 2a: Working memory retrieval (Path A)
-        def retrieve_from_working_memory():
-            """
-            Direct structure-based retrieval from working memory.
-            """
-            if memory_type not in ["All", "WorkingMemory"]:
-                return []
+        return parsed_goal, query_embedding, context, query
 
-            working_memory = self.graph_retriever.retrieve(
-                query=query, parsed_goal=parsed_goal, top_k=top_k, memory_scope="WorkingMemory"
-            )
-            # Rerank working_memory results
-            ranked_memories = self.reranker.rerank(
-                query=query,
-                query_embedding=query_embedding[0],
-                graph_results=working_memory,
-                top_k=top_k,
-                parsed_goal=parsed_goal,
+    @timed
+    def _retrieve_paths(self, query, parsed_goal, query_embedding, info, top_k, mode, memory_type):
+        """Run A/B/C retrieval paths in parallel"""
+        tasks = []
+        with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+            tasks.append(
+                executor.submit(
+                    self._retrieve_from_working_memory,
+                    query,
+                    parsed_goal,
+                    query_embedding,
+                    top_k,
+                    memory_type,
+                )
             )
-            return ranked_memories
-
-        # Step 2b: Parallel long-term and user memory retrieval (Path B)
-        def retrieve_ranked_long_term_and_user():
-            """
-            Retrieve from both long-term and user memory, then rank and merge results.
-            """
-            long_term_items = (
-                self.graph_retriever.retrieve(
-                    query=query,
-                    query_embedding=query_embedding,
-                    parsed_goal=parsed_goal,
-                    top_k=top_k * 2,
-                    memory_scope="LongTermMemory",
+            tasks.append(
+                executor.submit(
+                    self._retrieve_from_long_term_and_user,
+                    query,
+                    parsed_goal,
+                    query_embedding,
+                    top_k,
+                    memory_type,
                 )
-                if memory_type in ["All", "LongTermMemory"]
-                else []
             )
-            user_items = (
-                self.graph_retriever.retrieve(
-                    query=query,
-                    query_embedding=query_embedding,
-                    parsed_goal=parsed_goal,
-                    top_k=top_k * 2,
-                    memory_scope="UserMemory",
+            tasks.append(
+                executor.submit(
+                    self._retrieve_from_internet,
+                    query,
+                    parsed_goal,
+                    query_embedding,
+                    top_k,
+                    info,
+                    mode,
+                    memory_type,
                 )
-                if memory_type in ["All", "UserMemory"]
-                else []
             )
+            if self.moscube:
+                tasks.append(
+                    executor.submit(
+                        self._retrieve_from_memcubes,
+                        query,
+                        parsed_goal,
+                        query_embedding,
+                        top_k,
+                        "memos_cube01",
+                    )
+                )
+
+            results = []
+            for t in tasks:
+                results.extend(t.result())
+
+        logger.info(f"[SEARCH] Total raw results: {len(results)}")
+        return results
+
+    # --- Path A
+    @timed
+    def _retrieve_from_working_memory(
+        self, query, parsed_goal, query_embedding, top_k, memory_type
+    ):
+        """Retrieve and rerank from WorkingMemory"""
+        if memory_type not in ["All", "WorkingMemory"]:
+            logger.info(f"[PATH-A] '{query}'Skipped (memory_type does not match)")
+            return []
+        items = self.graph_retriever.retrieve(
+            query=query, parsed_goal=parsed_goal, top_k=top_k, memory_scope="WorkingMemory"
+        )
+        return self.reranker.rerank(
+            query=query,
+            query_embedding=query_embedding[0],
+            graph_results=items,
+            top_k=top_k,
+            parsed_goal=parsed_goal,
+        )
 
-            # Rerank combined results
-            ranked_memories = self.reranker.rerank(
+    # --- Path B
+    @timed
+    def _retrieve_from_long_term_and_user(
+        self, query, parsed_goal, query_embedding, top_k, memory_type
+    ):
+        """Retrieve and rerank from LongTermMemory and UserMemory"""
+        results = []
+        if memory_type in ["All", "LongTermMemory"]:
+            results += self.graph_retriever.retrieve(
                 query=query,
-                query_embedding=query_embedding[0],
-                graph_results=long_term_items + user_items,
-                top_k=top_k * 2,
                 parsed_goal=parsed_goal,
+                query_embedding=query_embedding,
+                top_k=top_k * 2,
+                memory_scope="LongTermMemory",
             )
-            return ranked_memories
-
-        # Step 2c: Internet retrieval (Path C)
-        def retrieve_from_internet():
-            """
-            Retrieve information from the internet using Google Custom Search API.
-            """
-            if not self.internet_retriever or mode == "fast" or not parsed_goal.internet_search:
-                return []
-            if memory_type not in ["All"]:
-                return []
-            internet_items = self.internet_retriever.retrieve_from_internet(
-                query=query, top_k=top_k, parsed_goal=parsed_goal, info=info
-            )
-
-            # Convert to the format expected by reranker
-            ranked_memories = self.reranker.rerank(
+        if memory_type in ["All", "UserMemory"]:
+            results += self.graph_retriever.retrieve(
                 query=query,
-                query_embedding=query_embedding[0],
-                graph_results=internet_items,
-                top_k=min(top_k, 5),
                 parsed_goal=parsed_goal,
+                query_embedding=query_embedding,
+                top_k=top_k * 2,
+                memory_scope="UserMemory",
             )
-            return ranked_memories
-
-        # Step 3: Parallel execution of all paths (enable internet search accoeding to parameter in the parsed goal)
-        if parsed_goal.internet_search:
-            with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
-                future_working = executor.submit(retrieve_from_working_memory)
-                future_hybrid = executor.submit(retrieve_ranked_long_term_and_user)
-                future_internet = executor.submit(retrieve_from_internet)
-
-                working_results = future_working.result()
-                hybrid_results = future_hybrid.result()
-                internet_results = future_internet.result()
-                searched_res = working_results + hybrid_results + internet_results
-        else:
-            with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
-                future_working = executor.submit(retrieve_from_working_memory)
-                future_hybrid = executor.submit(retrieve_ranked_long_term_and_user)
-
-                working_results = future_working.result()
-                hybrid_results = future_hybrid.result()
-                searched_res = working_results + hybrid_results
-
-        # Deduplicate by item.memory, keep higher score
-        deduped_result = {}
-        for item, score in searched_res:
-            mem_key = item.memory
-            if mem_key not in deduped_result or score > deduped_result[mem_key][1]:
-                deduped_result[mem_key] = (item, score)
-
-        searched_res = []
-        for item, score in sorted(deduped_result.values(), key=lambda pair: pair[1], reverse=True)[
-            :top_k
-        ]:
+        return self.reranker.rerank(
+            query=query,
+            query_embedding=query_embedding[0],
+            graph_results=results,
+            top_k=top_k,
+            parsed_goal=parsed_goal,
+        )
+
+    @timed
+    def _retrieve_from_memcubes(
+        self, query, parsed_goal, query_embedding, top_k, cube_name="memos_cube01"
+    ):
+        """Retrieve and rerank from LongTermMemory and UserMemory"""
+        results = self.graph_retriever.retrieve_from_cube(
+            query_embedding=query_embedding,
+            top_k=top_k * 2,
+            memory_scope="LongTermMemory",
+            cube_name=cube_name,
+        )
+        return self.reranker.rerank(
+            query=query,
+            query_embedding=query_embedding[0],
+            graph_results=results,
+            top_k=top_k,
+            parsed_goal=parsed_goal,
+        )
+
+    # --- Path C
+    @timed
+    def _retrieve_from_internet(
+        self, query, parsed_goal, query_embedding, top_k, info, mode, memory_type
+    ):
+        """Retrieve and rerank from Internet source"""
+        if not self.internet_retriever or mode == "fast":
+            logger.info(f"[PATH-C] '{query}' Skipped (no retriever, fast mode)")
+            return []
+        if memory_type not in ["All"]:
+            return []
+        logger.info(f"[PATH-C] '{query}' Retrieving from internet...")
+        items = self.internet_retriever.retrieve_from_internet(
+            query=query, top_k=top_k, parsed_goal=parsed_goal, info=info
+        )
+        logger.info(f"[PATH-C] '{query}' Retrieved from internet {len(items)} items: {items}")
+        return self.reranker.rerank(
+            query=query,
+            query_embedding=query_embedding[0],
+            graph_results=items,
+            top_k=min(top_k, 5),
+            parsed_goal=parsed_goal,
+        )
+
+    @timed
+    def _deduplicate_results(self, results):
+        """Deduplicate results by memory text"""
+        deduped = {}
+        for item, score in results:
+            if item.memory not in deduped or score > deduped[item.memory][1]:
+                deduped[item.memory] = (item, score)
+        return list(deduped.values())
+
+    @timed
+    def _sort_and_trim(self, results, top_k):
+        """Sort results by score and trim to top_k"""
+        sorted_results = sorted(results, key=lambda pair: pair[1], reverse=True)[:top_k]
+        final_items = []
+        for item, score in sorted_results:
             meta_data = item.metadata.model_dump()
             if "relativity" not in meta_data:
                 meta_data["relativity"] = score
-            new_meta = SearchedTreeNodeTextualMemoryMetadata(**meta_data)
-            searched_res.append(
-                TextualMemoryItem(id=item.id, memory=item.memory, metadata=new_meta)
+            final_items.append(
+                TextualMemoryItem(
+                    id=item.id,
+                    memory=item.memory,
+                    metadata=SearchedTreeNodeTextualMemoryMetadata(**meta_data),
+                )
             )
+        return final_items
 
-        # Step 5: Update usage history with current timestamp
+    @timed
+    def _update_usage_history(self, items, info):
+        """Update usage history in graph DB"""
         now_time = datetime.now().isoformat()
-        if "chat_history" in info:
-            info.pop("chat_history")
-        usage_record = json.dumps(
-            {"time": now_time, "info": info}
-        )  # `info` should be a serializable dict or string
-
-        for item in searched_res:
-            if (
-                hasattr(item, "id")
-                and hasattr(item, "metadata")
-                and hasattr(item.metadata, "usage")
-            ):
-                item.metadata.usage.append(usage_record)
-                self.graph_store.update_node(item.id, {"usage": item.metadata.usage})
-        return searched_res
+        info_copy = dict(info or {})
+        info_copy.pop("chat_history", None)
+        usage_record = json.dumps({"time": now_time, "info": info_copy})
+        payload = []
+        for it in items:
+            try:
+                item_id = getattr(it, "id", None)
+                md = getattr(it, "metadata", None)
+                if md is None:
+                    continue
+                if not hasattr(md, "usage") or md.usage is None:
+                    md.usage = []
+                md.usage.append(usage_record)
+                if item_id:
+                    payload.append((item_id, list(md.usage)))
+            except Exception:
+                logger.exception("[USAGE] snapshot item failed")
+
+        if payload:
+            self._usage_executor.submit(self._update_usage_history_worker, payload, usage_record)
+
+    def _update_usage_history_worker(self, payload, usage_record: str):
+        try:
+            for item_id, usage_list in payload:
+                self.graph_store.update_node(item_id, {"usage": usage_list})
+        except Exception:
+            logger.exception("[USAGE] update usage failed")

memos/memories/textual/tree_text_memory/retrieve/task_goal_parser.py

@@ -1,13 +1,16 @@
-import logging
 import traceback
 
 from string import Template
 
 from memos.llms.base import BaseLLM
+from memos.log import get_logger
 from memos.memories.textual.tree_text_memory.retrieve.retrieval_mid_structs import ParsedTaskGoal
 from memos.memories.textual.tree_text_memory.retrieve.utils import TASK_PARSE_PROMPT
 
 
+logger = get_logger(__name__)
+
+
 class TaskGoalParser:
     """
     Unified TaskGoalParser:
@@ -70,10 +73,12 @@ class TaskGoalParser:
             prompt = Template(TASK_PARSE_PROMPT).substitute(
                 task=query.strip(), context=context, conversation=conversation_prompt
             )
+            logger.info(f"Parsing Goal... LLM input is {prompt}")
             response = self.llm.generate(messages=[{"role": "user", "content": prompt}])
+            logger.info(f"Parsing Goal... LLM Response is {response}")
             return self._parse_response(response)
         except Exception:
-            logging.warning(f"Fail to fine-parse query {query}: {traceback.format_exc()}")
+            logger.warning(f"Fail to fine-parse query {query}: {traceback.format_exc()}")
             return self._parse_fast(query)
 
     def _parse_response(self, response: str) -> ParsedTaskGoal:

memos/memories/textual/tree_text_memory/retrieve/utils.py

@@ -5,21 +5,23 @@ You are a task parsing expert. Given a user task instruction, optional former co
 2. Tags: thematic tags to help categorize and retrieve related memories.
 3. Goal Type: retrieval | qa | generation
 4. Rephrased instruction: Give a rephrased task instruction based on the former conversation to make it less confusing to look alone. If you think the task instruction is easy enough to understand, or there is no former conversation, set "rephrased_instruction" to an empty string.
-5. Need for internet search: If you think you need to search the internet to finish the rephrased/original user task instruction, set "internet_search" to True. Otherwise, set it to False.
+5. Need for internet search: If the user's task instruction only involves objective facts or can be completed without introducing external knowledge, set "internet_search" to False. Otherwise, set it to True.
 6. Memories: Provide 2–5 short semantic expansions or rephrasings of the rephrased/original user task instruction. These are used for improved embedding search coverage. Each should be clear, concise, and meaningful for retrieval.
 
-Task description:
-\"\"\"$task\"\"\"
-
 Former conversation (if any):
 \"\"\"
 $conversation
 \"\"\"
 
+Task description(User Question):
+\"\"\"$task\"\"\"
+
 Context (if any):
 \"\"\"$context\"\"\"
 
-Return strictly in this JSON format:
+Return strictly in this JSON format, note that the
+keys/tags/rephrased_instruction/memories should use the same language as the
+input query:
 {
   "keys": [...],
   "tags": [...],

memos/memos_tools/lockfree_dict.py

@@ -0,0 +1,120 @@
+"""
+Lock-free dictionary implementation using copy-on-write strategy.
+This provides better performance but uses more memory.
+"""
+
+import threading
+
+from collections.abc import ItemsView, Iterator, KeysView, ValuesView
+from typing import Generic, TypeVar
+
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class CopyOnWriteDict(Generic[K, V]):
+    """
+    A lock-free dictionary using copy-on-write strategy.
+
+    Reads are completely lock-free and very fast.
+    Writes create a new copy of the dictionary.
+    Uses more memory but provides excellent read performance.
+    """
+
+    def __init__(self, initial_dict: dict[K, V] | None = None):
+        """Initialize with optional initial dictionary."""
+        self._dict = initial_dict.copy() if initial_dict else {}
+        self._write_lock = threading.Lock()  # Only for writes
+
+    def __getitem__(self, key: K) -> V:
+        """Get item by key - completely lock-free."""
+        return self._dict[key]
+
+    def __setitem__(self, key: K, value: V) -> None:
+        """Set item by key - uses copy-on-write."""
+        with self._write_lock:
+            # Create a new dictionary with the update
+            new_dict = self._dict.copy()
+            new_dict[key] = value
+            # Atomic replacement
+            self._dict = new_dict
+
+    def __delitem__(self, key: K) -> None:
+        """Delete item by key - uses copy-on-write."""
+        with self._write_lock:
+            new_dict = self._dict.copy()
+            del new_dict[key]
+            self._dict = new_dict
+
+    def __contains__(self, key: K) -> bool:
+        """Check if key exists - completely lock-free."""
+        return key in self._dict
+
+    def __len__(self) -> int:
+        """Get length - completely lock-free."""
+        return len(self._dict)
+
+    def __bool__(self) -> bool:
+        """Check if not empty - completely lock-free."""
+        return bool(self._dict)
+
+    def __iter__(self) -> Iterator[K]:
+        """Iterate over keys - completely lock-free."""
+        return iter(self._dict.keys())
+
+    def get(self, key: K, default: V | None = None) -> V:
+        """Get with default - completely lock-free."""
+        return self._dict.get(key, default)
+
+    def keys(self) -> KeysView[K]:
+        """Get keys - completely lock-free."""
+        return self._dict.keys()
+
+    def values(self) -> ValuesView[V]:
+        """Get values - completely lock-free."""
+        return self._dict.values()
+
+    def items(self) -> ItemsView[K, V]:
+        """Get items - completely lock-free."""
+        return self._dict.items()
+
+    def copy(self) -> dict[K, V]:
+        """Create a copy - completely lock-free."""
+        return self._dict.copy()
+
+    def update(self, *args, **kwargs) -> None:
+        """Update dictionary - uses copy-on-write."""
+        with self._write_lock:
+            new_dict = self._dict.copy()
+            new_dict.update(*args, **kwargs)
+            self._dict = new_dict
+
+    def clear(self) -> None:
+        """Clear all items."""
+        with self._write_lock:
+            self._dict = {}
+
+    def pop(self, key: K, *args) -> V:
+        """Pop item by key."""
+        with self._write_lock:
+            new_dict = self._dict.copy()
+            result = new_dict.pop(key, *args)
+            self._dict = new_dict
+            return result
+
+    def setdefault(self, key: K, default: V | None = None) -> V:
+        """Set default value for key if not exists."""
+        # Fast path for existing keys
+        if key in self._dict:
+            return self._dict[key]
+
+        with self._write_lock:
+            # Double-check after acquiring lock
+            if key in self._dict:
+                return self._dict[key]
+
+            new_dict = self._dict.copy()
+            result = new_dict.setdefault(key, default)
+            self._dict = new_dict
+            return result