feat: add more log in tree text mem retriever

Author: CaralHsi

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

@@ -1,5 +1,6 @@
 import concurrent.futures
 import json
+import time
 
 from datetime import datetime
 
@@ -57,67 +58,110 @@ def search(
         Returns:
             list[TextualMemoryItem]: List of matching memories.
         """
+        overall_start = time.perf_counter()
+        logger.info(
+            f"[SEARCH] 🚀 Starting search for 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}")
+
+        # ===== Step 1: Parse task structure =====
+        step_start = time.perf_counter()
         context = []
         if mode == "fine":
+            logger.info("[SEARCH] Fine mode enabled, performing initial embedding search...")
+            embed_start = time.perf_counter()
             query_embedding = self.embedder.embed([query])[0]
+            logger.debug(f"[SEARCH] Query embedding vector length: {len(query_embedding)}")
+            logger.info(
+                f"[TIMER] Embedding query took {(time.perf_counter() - embed_start) * 1000:.2f} ms"
+            )
+
+            search_start = time.perf_counter()
             related_node_ids = self.graph_store.search_by_embedding(query_embedding, top_k=top_k)
             related_nodes = [
                 self.graph_store.get_node(related_node["id"]) for related_node in related_node_ids
             ]
-
             context = [related_node["memory"] for related_node in related_nodes]
             context = list(set(context))
+            logger.info(f"[SEARCH] Found {len(related_nodes)} related nodes from graph_store.")
+            logger.info(
+                f"[TIMER] Graph embedding search took {(time.perf_counter() - search_start) * 1000:.2f} ms"
+            )
 
-        # Step 1a: Parse task structure into topic, concept, and fact levels
+        parse_start = time.perf_counter()
         parsed_goal = self.task_goal_parser.parse(
             task_description=query,
             context="\n".join(context),
             conversation=info.get("chat_history", []),
             mode=mode,
         )
-
-        query = (
-            parsed_goal.rephrased_query
-            if parsed_goal.rephrased_query and len(parsed_goal.rephrased_query) > 0
-            else query
+        logger.info(
+            f"[TIMER] TaskGoalParser took {(time.perf_counter() - parse_start) * 1000:.2f} ms"
         )
+        logger.info(f"TaskGoalParser result is {parsed_goal}")
 
+        query = parsed_goal.rephrased_query or query
         if parsed_goal.memories:
+            embed_extra_start = time.perf_counter()
             query_embedding = self.embedder.embed(list({query, *parsed_goal.memories}))
+            logger.info(
+                f"[TIMER] Embedding parsed_goal memories took {(time.perf_counter() - embed_extra_start) * 1000:.2f} ms"
+            )
+        step_end = time.perf_counter()
+        logger.info(f"[TIMER] Step 1 (Parsing & Embedding) took {(step_end - step_start):.2f} s")
+
+        # ===== Step 2: Define retrieval paths =====
+        def timed(func):
+            """Decorator to measure and log time of retrieval steps."""
+
+            def wrapper(*args, **kwargs):
+                start = time.perf_counter()
+                result = func(*args, **kwargs)
+                elapsed = time.perf_counter() - start
+                logger.info(f"[TIMER] {func.__name__} took {elapsed:.2f} s")
+                return result
 
-        # Step 2a: Working memory retrieval (Path A)
+            return wrapper
+
+        @timed
         def retrieve_from_working_memory():
             """
             Direct structure-based retrieval from working memory.
             """
+            logger.info("[PATH-A] Retrieving from WorkingMemory...")
             if memory_type not in ["All", "WorkingMemory"]:
+                logger.info("[PATH-A] Skipped (memory_type does not match)")
                 return []
-
             working_memory = self.graph_retriever.retrieve(
                 query=query, parsed_goal=parsed_goal, top_k=top_k, memory_scope="WorkingMemory"
             )
+
+            logger.debug(f"[PATH-A] Retrieved {len(working_memory)} items.")
             # Rerank working_memory results
+            rerank_start = time.perf_counter()
             ranked_memories = self.reranker.rerank(
                 query=query,
                 query_embedding=query_embedding[0],
                 graph_results=working_memory,
                 top_k=top_k,
                 parsed_goal=parsed_goal,
             )
+            logger.info(
+                f"[TIMER] PATH-A rerank took {(time.perf_counter() - rerank_start) * 1000:.2f} ms"
+            )
             return ranked_memories
 
-        # Step 2b: Parallel long-term and user memory retrieval (Path B)
+        @timed
         def retrieve_ranked_long_term_and_user():
-            """
-            Retrieve from both long-term and user memory, then rank and merge results.
-            """
+            logger.info("[PATH-B] Retrieving from LongTermMemory & UserMemory...")
             long_term_items = (
                 self.graph_retriever.retrieve(
                     query=query,
@@ -140,7 +184,10 @@ def retrieve_ranked_long_term_and_user():
                 if memory_type in ["All", "UserMemory"]
                 else []
             )
-
+            logger.debug(
+                f"[PATH-B] Retrieved {len(long_term_items)} LongTerm + {len(user_items)} UserMemory items."
+            )
+            rerank_start = time.perf_counter()
             # Rerank combined results
             ranked_memories = self.reranker.rerank(
                 query=query,
@@ -149,21 +196,30 @@ def retrieve_ranked_long_term_and_user():
                 top_k=top_k * 2,
                 parsed_goal=parsed_goal,
             )
+            logger.info(
+                f"[TIMER] PATH-B rerank took {(time.perf_counter() - rerank_start) * 1000:.2f} ms"
+            )
             return ranked_memories
 
-        # Step 2c: Internet retrieval (Path C)
+        @timed
         def retrieve_from_internet():
             """
             Retrieve information from the internet using Google Custom Search API.
             """
+            logger.info("[PATH-C] Retrieving from Internet...")
             if not self.internet_retriever or mode == "fast" or not parsed_goal.internet_search:
+                logger.info(
+                    "[PATH-C] Skipped (no retriever, fast mode, or no internet_search flag)"
+                )
                 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
             )
 
+            logger.debug(f"[PATH-C] Retrieved {len(internet_items)} internet items.")
+            rerank_start = time.perf_counter()
             # Convert to the format expected by reranker
             ranked_memories = self.reranker.rerank(
                 query=query,
@@ -172,9 +228,13 @@ def retrieve_from_internet():
                 top_k=min(top_k, 5),
                 parsed_goal=parsed_goal,
             )
+            logger.info(
+                f"[TIMER] PATH-C rerank took {(time.perf_counter() - rerank_start) * 1000:.2f} ms"
+            )
             return ranked_memories
 
-        # Step 3: Parallel execution of all paths (enable internet search accoeding to parameter in the parsed goal)
+        # ===== Step 3: Run retrieval in parallel =====
+        path_start = time.perf_counter()
         if parsed_goal.internet_search:
             with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
                 future_working = executor.submit(retrieve_from_working_memory)
@@ -193,14 +253,24 @@ def retrieve_from_internet():
                 working_results = future_working.result()
                 hybrid_results = future_hybrid.result()
                 searched_res = working_results + hybrid_results
+        logger.info(
+            f"[TIMER] Step 3 (Retrieval paths) took {(time.perf_counter() - path_start):.2f} s"
+        )
+        logger.info(f"[SEARCH] Total results before deduplication: {len(searched_res)}")
 
-        # Deduplicate by item.memory, keep higher score
+        # ===== Step 4: Deduplication =====
+        dedup_start = time.perf_counter()
         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)
+        logger.info(
+            f"[TIMER] Deduplication took {(time.perf_counter() - dedup_start) * 1000:.2f} ms"
+        )
 
+        # ===== Step 5: Sorting & trimming =====
+        sort_start = time.perf_counter()
         searched_res = []
         for item, score in sorted(deduped_result.values(), key=lambda pair: pair[1], reverse=True)[
             :top_k
@@ -212,15 +282,18 @@ def retrieve_from_internet():
             searched_res.append(
                 TextualMemoryItem(id=item.id, memory=item.memory, metadata=new_meta)
             )
+        logger.info(
+            f"[TIMER] Sorting & trimming took {(time.perf_counter() - sort_start) * 1000:.2f} ms"
+        )
 
-        # Step 5: Update usage history with current timestamp
+        # ===== Step 6: Update usage history =====
+        usage_start = time.perf_counter()
         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")
@@ -229,4 +302,11 @@ def retrieve_from_internet():
             ):
                 item.metadata.usage.append(usage_record)
                 self.graph_store.update_node(item.id, {"usage": item.metadata.usage})
+        logger.info(
+            f"[TIMER] Usage history update took {(time.perf_counter() - usage_start) * 1000:.2f} ms"
+        )
+
+        # ===== Finish =====
+        logger.info(f"[SEARCH] ✅ Final top_k results: {len(searched_res)}")
+        logger.info(f"[SEARCH] 🔚 Total search took {(time.perf_counter() - overall_start):.2f} s")
         return searched_res

src/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 @@ def _parse_fine(
             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: