MemoryOS 0.2.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

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

@@ -4,6 +4,8 @@ from typing import Any, ClassVar
 
 from memos.configs.internet_retriever import InternetRetrieverConfigFactory
 from memos.embedders.base import BaseEmbedder
+from memos.mem_reader.factory import MemReaderFactory
+from memos.memories.textual.tree_text_memory.retrieve.bochasearch import BochaAISearchRetriever
 from memos.memories.textual.tree_text_memory.retrieve.internet_retriever import (
     InternetGoogleRetriever,
 )
@@ -17,6 +19,7 @@ class InternetRetrieverFactory:
         "google": InternetGoogleRetriever,
         "bing": InternetGoogleRetriever,  # TODO: Implement BingRetriever
         "xinyu": XinyuSearchRetriever,
+        "bocha": BochaAISearchRetriever,
     }
 
     @classmethod
@@ -66,6 +69,14 @@ class InternetRetrieverFactory:
                 access_key=config.api_key,  # Use api_key as access_key for xinyu
                 search_engine_id=config.search_engine_id,
                 embedder=embedder,
+                reader=MemReaderFactory.from_config(config.reader),
+                max_results=config.max_results,
+            )
+        elif backend == "bocha":
+            return retriever_class(
+                access_key=config.api_key,  # Use api_key as access_key for xinyu
+                embedder=embedder,
+                reader=MemReaderFactory.from_config(config.reader),
                 max_results=config.max_results,
             )
         else:

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

@@ -44,16 +44,23 @@ class GraphMemoryRetriever:
 
         if memory_scope == "WorkingMemory":
             # For working memory, retrieve all entries (no filtering)
-            working_memories = self.graph_store.get_all_memory_items(scope="WorkingMemory")
+            working_memories = self.graph_store.get_all_memory_items(
+                scope="WorkingMemory", include_embedding=True
+            )
             return [TextualMemoryItem.from_dict(record) for record in working_memories]
 
-        # Step 1: Structured graph-based retrieval
-        graph_results = self._graph_recall(parsed_goal, memory_scope)
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+            # Structured graph-based retrieval
+            future_graph = executor.submit(self._graph_recall, parsed_goal, memory_scope)
+            # Vector similarity search
+            future_vector = executor.submit(
+                self._vector_recall, query_embedding, memory_scope, top_k
+            )
 
-        # Step 2: Vector similarity search
-        vector_results = self._vector_recall(query_embedding, memory_scope, top_k)
+            graph_results = future_graph.result()
+            vector_results = future_vector.result()
 
-        # Step 3: Merge and deduplicate results
+        # Merge and deduplicate by ID
         combined = {item.id: item for item in graph_results + vector_results}
 
         graph_ids = {item.id for item in graph_results}
@@ -101,7 +108,7 @@ class GraphMemoryRetriever:
             return []
 
         # Load nodes and post-filter
-        node_dicts = self.graph_store.get_nodes(list(candidate_ids))
+        node_dicts = self.graph_store.get_nodes(list(candidate_ids), include_embedding=True)
 
         final_nodes = []
         for node in node_dicts:
@@ -152,6 +159,6 @@ class GraphMemoryRetriever:
 
         # Step 3: Extract matched IDs and retrieve full nodes
         unique_ids = set({r["id"] for r in all_matches})
-        node_dicts = self.graph_store.get_nodes(list(unique_ids))
+        node_dicts = self.graph_store.get_nodes(list(unique_ids), include_embedding=True)
 
         return [TextualMemoryItem.from_dict(record) for record in node_dicts]

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

@@ -78,7 +78,7 @@ class MemoryReranker:
         embeddings = [item.metadata.embedding for item in items_with_embeddings]
 
         if not embeddings:
-            return graph_results[:top_k]
+            return [(item, 0.5) for item in graph_results[:top_k]]
 
         # Step 2: Compute cosine similarities
         similarity_scores = batch_cosine_similarity(query_embedding, embeddings)

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

@@ -10,4 +10,6 @@ class ParsedTaskGoal:
     memories: list[str] = field(default_factory=list)
     keys: list[str] = field(default_factory=list)
     tags: list[str] = field(default_factory=list)
+    rephrased_query: str | None = None
+    internet_search: bool = False
     goal_type: str | None = None  # e.g., 'default', 'explanation', etc.

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

@@ -6,7 +6,9 @@ from datetime import datetime
 from memos.embedders.factory import OllamaEmbedder
 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.utils import timed
 
 from .internet_retriever_factory import InternetRetrieverFactory
 from .reasoner import MemoryReasoner
@@ -15,6 +17,9 @@ from .reranker import MemoryReranker
 from .task_goal_parser import TaskGoalParser
 
 
+logger = get_logger(__name__)
+
+
 class Searcher:
     def __init__(
         self,
@@ -34,8 +39,9 @@ class Searcher:
         # Create internet retriever from config if provided
         self.internet_retriever = internet_retriever
 
+    @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.
@@ -53,152 +59,222 @@ 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": ""}
+        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.")
+        return final_results
 
-        # Step 1: Parse task structure into topic, concept, and fact levels
+    @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
-        parsed_goal = self.task_goal_parser.parse(query, "\n".join(context))
+        # parse goal using LLM
+        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 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 []
             )
 
-            # Rerank combined results
-            ranked_memories = self.reranker.rerank(
+            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,
+        )
+
+    # --- 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:
-                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
-            )
-
-            # 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=top_k * 2,
                 parsed_goal=parsed_goal,
+                query_embedding=query_embedding,
+                top_k=top_k * 2,
+                memory_scope="UserMemory",
             )
-            return ranked_memories
+        return self.reranker.rerank(
+            query=query,
+            query_embedding=query_embedding[0],
+            graph_results=results,
+            top_k=top_k * 2,
+            parsed_goal=parsed_goal,
+        )
 
-        # Step 3: Parallel execution of all paths
-        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
-
-        # 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
-        ]:
+    # --- 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)
-            )
-
-        # Step 4: Reasoning over all retrieved and ranked memory
-        if mode == "fine":
-            searched_res = self.reasoner.reason(
-                query=query,
-                ranked_memories=searched_res,
-                parsed_goal=parsed_goal,
+            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()
-        usage_record = json.dumps(
-            {"time": now_time, "info": info}
-        )  # `info` should be a serializable dict or string
-
-        for item in searched_res:
+        info.pop("chat_history", None)
+        # `info` should be a serializable dict or string
+        usage_record = json.dumps({"time": now_time, "info": info})
+        for item in items:
             if (
                 hasattr(item, "id")
                 and hasattr(item, "metadata")
@@ -206,4 +282,3 @@ class Searcher:
             ):
                 item.metadata.usage.append(usage_record)
                 self.graph_store.update_node(item.id, {"usage": item.metadata.usage})
-        return searched_res

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

@@ -1,12 +1,16 @@
-import json
+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:
@@ -14,11 +18,16 @@ class TaskGoalParser:
     - mode == 'fine': use LLM to parse structured topic/keys/tags
     """
 
-    def __init__(self, llm=BaseLLM, mode: str = "fast"):
+    def __init__(self, llm=BaseLLM):
         self.llm = llm
-        self.mode = mode
 
-    def parse(self, task_description: str, context: str = "") -> ParsedTaskGoal:
+    def parse(
+        self,
+        task_description: str,
+        context: str = "",
+        conversation: list[dict] | None = None,
+        mode: str = "fast",
+    ) -> ParsedTaskGoal:
         """
         Parse user input into structured semantic layers.
         Returns:
@@ -26,42 +35,65 @@ class TaskGoalParser:
         - mode == 'fast': use jieba to split words only
         - mode == 'fine': use LLM to parse structured topic/keys/tags
         """
-        if self.mode == "fast":
+        if mode == "fast":
             return self._parse_fast(task_description)
-        elif self.mode == "fine":
+        elif mode == "fine":
             if not self.llm:
                 raise ValueError("LLM not provided for slow mode.")
-            return self._parse_fine(task_description, context)
+            return self._parse_fine(task_description, context, conversation)
         else:
-            raise ValueError(f"Unknown mode: {self.mode}")
+            raise ValueError(f"Unknown mode: {mode}")
 
     def _parse_fast(self, task_description: str, limit_num: int = 5) -> ParsedTaskGoal:
         """
         Fast mode: simple jieba word split.
         """
         return ParsedTaskGoal(
-            memories=[task_description], keys=[task_description], tags=[], goal_type="default"
+            memories=[task_description],
+            keys=[task_description],
+            tags=[],
+            goal_type="default",
+            rephrased_query=task_description,
+            internet_search=False,
         )
 
-    def _parse_fine(self, query: str, context: str = "") -> ParsedTaskGoal:
+    def _parse_fine(
+        self, query: str, context: str = "", conversation: list[dict] | None = None
+    ) -> ParsedTaskGoal:
         """
         Slow mode: LLM structured parse.
         """
-        prompt = Template(TASK_PARSE_PROMPT).substitute(task=query.strip(), context=context)
-        response = self.llm.generate(messages=[{"role": "user", "content": prompt}])
-        return self._parse_response(response)
+        try:
+            if conversation:
+                conversation_prompt = "\n".join(
+                    [f"{each['role']}: {each['content']}" for each in conversation]
+                )
+            else:
+                conversation_prompt = ""
+            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:
+            logger.warning(f"Fail to fine-parse query {query}: {traceback.format_exc()}")
+            return self._parse_fast(query)
 
     def _parse_response(self, response: str) -> ParsedTaskGoal:
         """
         Parse LLM JSON output safely.
         """
         try:
-            response = response.replace("```", "").replace("json", "")
-            response_json = json.loads(response.strip())
+            response = response.replace("```", "").replace("json", "").strip()
+            response_json = eval(response)
             return ParsedTaskGoal(
                 memories=response_json.get("memories", []),
                 keys=response_json.get("keys", []),
                 tags=response_json.get("tags", []),
+                rephrased_query=response_json.get("rephrased_instruction", None),
+                internet_search=response_json.get("internet_search", False),
                 goal_type=response_json.get("goal_type", "default"),
             )
         except Exception as e:

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

@@ -1,19 +1,21 @@
 # Prompt for task parsing
 TASK_PARSE_PROMPT = """
-You are a task parsing expert. Given a user's task instruction, extract the following structured information:
-
-Given a user task instruction and optional related memory context,
-extract the following structured information:
+You are a task parsing expert. Given a user task instruction, optional former conversation and optional related memory context,extract the following structured information:
 1. Keys: the high-level keywords directly relevant to the user’s task.
 2. Tags: thematic tags to help categorize and retrieve related memories.
 3. Goal Type: retrieval | qa | generation
-4. Memories: Provide 2–5 short semantic expansions or rephrasings of the task instruction.
-   These are used for improved embedding search coverage.
-   Each should be clear, concise, and meaningful for retrieval.
+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 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
+\"\"\"
+
 Context (if any):
 \"\"\"$context\"\"\"
 
@@ -22,6 +24,8 @@ Return strictly in this JSON format:
   "keys": [...],
   "tags": [...],
   "goal_type": "retrieval | qa | generation",
+  "rephrased_instruction": "...", # return an empty string if the original instruction is easy enough to understand
+  "internet_search": True/False,
   "memories": ["...", "...", ...]
 }
 """