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

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

@@ -5,7 +5,8 @@ from datetime import datetime
 
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.factory import Neo4jGraphDB
-from memos.llms.factory import OllamaLLM, OpenAILLM
+from memos.llms.factory import AzureLLM, OllamaLLM, OpenAILLM
+from memos.log import get_logger
 from memos.memories.textual.item import SearchedTreeNodeTextualMemoryMetadata, TextualMemoryItem
 
 from .internet_retriever_factory import InternetRetrieverFactory
@@ -15,10 +16,13 @@ from .reranker import MemoryReranker
 from .task_goal_parser import TaskGoalParser
 
 
+logger = get_logger(__name__)
+
+
 class Searcher:
     def __init__(
         self,
-        dispatcher_llm: OpenAILLM | OllamaLLM,
+        dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM,
         graph_store: Neo4jGraphDB,
         embedder: OllamaEmbedder,
         internet_retriever: InternetRetrieverFactory | None = None,
@@ -53,7 +57,12 @@ class Searcher:
         Returns:
             list[TextualMemoryItem]: List of matching memories.
         """
-
+        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
         context = []
         if mode == "fine":
@@ -67,7 +76,18 @@ class Searcher:
             context = list(set(context))
 
         # Step 1a: Parse task structure into topic, concept, and fact levels
-        parsed_goal = self.task_goal_parser.parse(query, "\n".join(context))
+        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
+        )
 
         if parsed_goal.memories:
             query_embedding = self.embedder.embed(list({query, *parsed_goal.memories}))
@@ -136,12 +156,12 @@ class Searcher:
             """
             Retrieve information from the internet using Google Custom Search API.
             """
-            if not self.internet_retriever:
+            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
+                query=query, top_k=top_k, parsed_goal=parsed_goal, info=info
             )
 
             # Convert to the format expected by reranker
@@ -149,21 +169,30 @@ class Searcher:
                 query=query,
                 query_embedding=query_embedding[0],
                 graph_results=internet_items,
-                top_k=top_k * 2,
+                top_k=min(top_k, 5),
                 parsed_goal=parsed_goal,
             )
             return ranked_memories
 
-        # 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
+        # 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 = {}
@@ -176,23 +205,18 @@ class Searcher:
         for item, score in sorted(deduped_result.values(), key=lambda pair: pair[1], reverse=True)[
             :top_k
         ]:
-            new_meta = SearchedTreeNodeTextualMemoryMetadata(
-                **item.metadata.model_dump(), relativity=score
-            )
+            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,
-            )
-
         # Step 5: Update usage history with current timestamp
         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

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

@@ -1,4 +1,5 @@
-import json
+import logging
+import traceback
 
 from string import Template
 
@@ -14,11 +15,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 +32,63 @@ 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
+            )
+            response = self.llm.generate(messages=[{"role": "user", "content": prompt}])
+            return self._parse_response(response)
+        except Exception:
+            logging.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 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.
+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": ["...", "...", ...]
 }
 """

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

@@ -3,13 +3,15 @@
 import json
 import uuid
 
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
 
 import requests
 
 from memos.embedders.factory import OllamaEmbedder
 from memos.log import get_logger
-from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
+from memos.mem_reader.base import BaseMemReader
+from memos.memories.textual.item import TextualMemoryItem
 
 
 logger = get_logger(__name__)
@@ -93,8 +95,8 @@ class XinyuSearchAPI:
             "online_search": {
                 "max_entries": max_results,
                 "cache_switch": False,
-                "baidu_field": {"switch": True, "mode": "relevance", "type": "page"},
-                "bing_field": {"switch": False, "mode": "relevance", "type": "page_web"},
+                "baidu_field": {"switch": False, "mode": "relevance", "type": "page"},
+                "bing_field": {"switch": True, "mode": "relevance", "type": "page"},
                 "sogou_field": {"switch": False, "mode": "relevance", "type": "page"},
             },
             "request_id": "memos" + str(uuid.uuid4()),
@@ -112,6 +114,7 @@ class XinyuSearchRetriever:
         access_key: str,
         search_engine_id: str,
         embedder: OllamaEmbedder,
+        reader: BaseMemReader,
         max_results: int = 20,
     ):
         """
@@ -121,12 +124,14 @@ class XinyuSearchRetriever:
             access_key: Xinyu API access key
             embedder: Embedder instance for generating embeddings
             max_results: Maximum number of results to retrieve
+            reader: MemReader Moduel to deal with internet contents
         """
         self.xinyu_api = XinyuSearchAPI(access_key, search_engine_id, max_results=max_results)
         self.embedder = embedder
+        self.reader = reader
 
     def retrieve_from_internet(
-        self, query: str, top_k: int = 10, parsed_goal=None
+        self, query: str, top_k: int = 10, parsed_goal=None, info=None
     ) -> list[TextualMemoryItem]:
         """
         Retrieve information from Xinyu search and convert to TextualMemoryItem format
@@ -135,7 +140,7 @@ class XinyuSearchRetriever:
             query: Search query
             top_k: Number of results to return
             parsed_goal: Parsed task goal (optional)
-
+            info (dict): Leave a record of memory consumption.
         Returns:
             List of TextualMemoryItem
         """
@@ -143,63 +148,25 @@ class XinyuSearchRetriever:
         search_results = self.xinyu_api.search(query, max_results=top_k)
 
         # Convert to TextualMemoryItem format
-        memory_items = []
-
-        for _, result in enumerate(search_results):
-            # Extract basic information from Xinyu response format
-            title = result.get("title", "")
-            content = result.get("content", "")
-            summary = result.get("summary", "")
-            url = result.get("url", "")
-            publish_time = result.get("publish_time", "")
-            if publish_time:
+        memory_items: list[TextualMemoryItem] = []
+
+        with ThreadPoolExecutor(max_workers=8) as executor:
+            futures = [
+                executor.submit(self._process_result, result, query, parsed_goal, info)
+                for result in search_results
+            ]
+            for future in as_completed(futures):
                 try:
-                    publish_time = datetime.strptime(publish_time, "%Y-%m-%d %H:%M:%S").strftime(
-                        "%Y-%m-%d"
-                    )
+                    memory_items.extend(future.result())
                 except Exception as e:
-                    logger.error(f"xinyu search error: {e}")
-                    publish_time = datetime.now().strftime("%Y-%m-%d")
-            else:
-                publish_time = datetime.now().strftime("%Y-%m-%d")
-            source = result.get("source", "")
-            site = result.get("site", "")
-            if site:
-                site = site.split("|")[0]
-
-            # Combine memory content
-            memory_content = (
-                f"Title: {title}\nSummary: {summary}\nContent: {content[:200]}...\nSource: {url}"
-            )
+                    logger.error(f"Error processing search result: {e}")
 
-            # Create metadata
-            metadata = TreeNodeTextualMemoryMetadata(
-                user_id=None,
-                session_id=None,
-                status="activated",
-                type="fact",  # Search results are usually factual information
-                memory_time=publish_time,
-                source="web",
-                confidence=85.0,  # Confidence level for search information
-                entities=self._extract_entities(title, content, summary),
-                tags=self._extract_tags(title, content, summary, parsed_goal),
-                visibility="public",
-                memory_type="LongTermMemory",  # Search results as working memory
-                key=title,
-                sources=[url] if url else [],
-                embedding=self.embedder.embed([memory_content])[0],
-                created_at=datetime.now().isoformat(),
-                usage=[],
-                background=f"Xinyu search result from {site or source}",
-            )
-            # Create TextualMemoryItem
-            memory_item = TextualMemoryItem(
-                id=str(uuid.uuid4()), memory=memory_content, metadata=metadata
-            )
+        unique_memory_items = {}
+        for item in memory_items:
+            if item.memory not in unique_memory_items:
+                unique_memory_items[item.memory] = item
 
-            memory_items.append(memory_item)
-
-        return memory_items
+        return list(unique_memory_items.values())
 
     def _extract_entities(self, title: str, content: str, summary: str) -> list[str]:
         """
@@ -333,3 +300,40 @@ class XinyuSearchRetriever:
             tags.extend(parsed_goal.tags)
 
         return list(set(tags))[:15]  # Limit to 15 tags
+
+    def _process_result(
+        self, result: dict, query: str, parsed_goal: str, info: None
+    ) -> list[TextualMemoryItem]:
+        if not info:
+            info = {"user_id": "", "session_id": ""}
+        title = result.get("title", "")
+        content = result.get("content", "")
+        summary = result.get("summary", "")
+        url = result.get("url", "")
+        publish_time = result.get("publish_time", "")
+        if publish_time:
+            try:
+                publish_time = datetime.strptime(publish_time, "%Y-%m-%d %H:%M:%S").strftime(
+                    "%Y-%m-%d"
+                )
+            except Exception as e:
+                logger.error(f"xinyu search error: {e}")
+                publish_time = datetime.now().strftime("%Y-%m-%d")
+        else:
+            publish_time = datetime.now().strftime("%Y-%m-%d")
+
+        read_items = self.reader.get_memory([content], type="doc", info=info)
+
+        memory_items = []
+        for read_item_i in read_items[0]:
+            read_item_i.memory = (
+                f"Title: {title}\nNewsTime: {publish_time}\nSummary: {summary}\n"
+                f"Content: {read_item_i.memory}"
+            )
+            read_item_i.metadata.source = "web"
+            read_item_i.metadata.memory_type = "OuterMemory"
+            read_item_i.metadata.sources = [url] if url else []
+            read_item_i.metadata.visibility = "public"
+
+            memory_items.append(read_item_i)
+        return memory_items