MemoryOS 2.0.3__py3-none-any.whl

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

@@ -0,0 +1,364 @@
+import copy
+import time
+
+from typing import Any
+
+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 TextualMemoryItem, TextualMemoryMetadata
+from memos.memories.textual.tree_text_memory.retrieve.bm25_util import EnhancedBM25
+from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import (
+    FastTokenizer,
+    parse_structured_output,
+)
+from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
+from memos.reranker.base import BaseReranker
+from memos.templates.advanced_search_prompts import PROMPT_MAPPING
+from memos.types.general_types import SearchMode
+
+
+logger = get_logger(__name__)
+
+
+class AdvancedSearcher(Searcher):
+    def __init__(
+        self,
+        dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM,
+        graph_store: Neo4jGraphDB,
+        embedder: OllamaEmbedder,
+        reranker: BaseReranker,
+        bm25_retriever: EnhancedBM25 | None = None,
+        internet_retriever: None = None,
+        search_strategy: dict | None = None,
+        manual_close_internet: bool = True,
+        process_llm: Any | None = None,
+        tokenizer: FastTokenizer | None = None,
+        include_embedding: bool = False,
+    ):
+        super().__init__(
+            dispatcher_llm=dispatcher_llm,
+            graph_store=graph_store,
+            embedder=embedder,
+            reranker=reranker,
+            bm25_retriever=bm25_retriever,
+            internet_retriever=internet_retriever,
+            search_strategy=search_strategy,
+            manual_close_internet=manual_close_internet,
+            tokenizer=tokenizer,
+            include_embedding=include_embedding,
+        )
+
+        self.stage_retrieve_top = 3
+        self.process_llm = process_llm
+        self.thinking_stages = 3
+        self.max_retry_times = 2
+        self.deep_search_top_k_bar = 2
+
+    def load_template(self, template_name: str) -> str:
+        if template_name not in PROMPT_MAPPING:
+            logger.error("Prompt template is not found!")
+        prompt = PROMPT_MAPPING[template_name]
+        return prompt
+
+    def build_prompt(self, template_name: str, **kwargs) -> str:
+        template = self.load_template(template_name)
+        if not template:
+            raise FileNotFoundError(f"Prompt template `{template_name}` not found.")
+        return template.format(**kwargs)
+
+    def stage_retrieve(
+        self,
+        stage_id: int,
+        query: str,
+        previous_retrieval_phrases: list[str],
+        text_memories: str,
+    ) -> tuple[bool, str, list[str]]:
+        """Run a retrieval-expansion stage and parse structured LLM output.
+
+        Returns a tuple of:
+        - can_answer: whether current memories suffice to answer
+        - reason: brief reasoning or hypotheses
+        - context: synthesized context summary
+        - retrieval_phrases: list of phrases to retrieve next
+        """
+
+        # Format previous phrases as bullet list to align with prompt expectations
+        prev_phrases_text = (
+            "- " + "\n- ".join(previous_retrieval_phrases) if previous_retrieval_phrases else ""
+        )
+
+        args = {
+            "template_name": f"stage{stage_id}_expand_retrieve",
+            "query": query,
+            "previous_retrieval_phrases": prev_phrases_text,
+            "memories": text_memories,
+        }
+        prompt = self.build_prompt(**args)
+
+        max_attempts = max(0, self.max_retry_times) + 1
+        for attempt in range(1, max_attempts + 1):
+            try:
+                llm_response = self.process_llm.generate(
+                    [{"role": "user", "content": prompt}]
+                ).strip()
+                result = parse_structured_output(content=llm_response)
+
+                # Parse booleans and fallbacks robustly
+                can_answer_str = str(result.get("can_answer", "")).strip().lower()
+                can_answer = can_answer_str in {"true", "yes", "y", "1"}
+
+                reason = result.get("reason", "")
+
+                phrases_val = result.get("retrieval_phrases", result.get("retrival_phrases", []))
+                if isinstance(phrases_val, list):
+                    retrieval_phrases = [str(p).strip() for p in phrases_val if str(p).strip()]
+                elif isinstance(phrases_val, str) and phrases_val.strip():
+                    retrieval_phrases = [p.strip() for p in phrases_val.splitlines() if p.strip()]
+                else:
+                    retrieval_phrases = []
+
+                return can_answer, reason, retrieval_phrases
+
+            except Exception as e:
+                if attempt < max_attempts:
+                    logger.debug(f"[stage_retrieve]🔁 retry {attempt}/{max_attempts} failed: {e!s}")
+                    time.sleep(1)
+                else:
+                    logger.error(
+                        f"[stage_retrieve]❌ all {max_attempts} attempts failed: {e!s}; \nprompt: {prompt}",
+                        exc_info=True,
+                    )
+                    raise e
+
+    def judge_memories(self, query: str, text_memories: str):
+        args = {
+            "template_name": "memory_judgement",
+            "query": query,
+            "memories": text_memories,
+        }
+
+        prompt = self.build_prompt(**args)
+
+        max_attempts = max(0, self.max_retry_times) + 1
+        for attempt in range(1, max_attempts + 1):
+            try:
+                llm_response = self.process_llm.generate([{"role": "user", "content": prompt}])
+                result = parse_structured_output(content=llm_response)
+                reason, can_answer = (
+                    result["reason"],
+                    result["can_answer"],
+                )
+
+                return reason, can_answer
+            except Exception as e:
+                if attempt < max_attempts:
+                    logger.debug(
+                        f"[summarize_and_eval]🔁 retry {attempt}/{max_attempts} failed: {e!s}"
+                    )
+                    time.sleep(1)
+                else:
+                    logger.error(
+                        f"[summarize_and_eval]❌ all {max_attempts} attempts failed: {e!s}; \nprompt: {prompt}",
+                        exc_info=True,
+                    )
+                    raise e
+
+    def tree_memories_to_text_memories(self, memories: list[TextualMemoryItem]):
+        mem_list = []
+        source_documents = []
+        for mem in memories:
+            source_documents.extend(
+                [f"({one.chat_time}) {one.content}" for one in mem.metadata.sources]
+            )
+            mem_list.append(mem.memory)
+        mem_list = list(set(mem_list))
+        source_documents = list(set(source_documents))
+        return mem_list, source_documents
+
+    def get_final_memories(self, user_id: str, top_k: int, mem_list: list[str]):
+        enhanced_memories = []
+        for new_mem in mem_list:
+            enhanced_memories.append(
+                TextualMemoryItem(memory=new_mem, metadata=TextualMemoryMetadata(user_id=user_id))
+            )
+        if len(enhanced_memories) > top_k:
+            logger.info(
+                f"Result count {len(enhanced_memories)} exceeds requested top_k {top_k}, truncating to top {top_k} memories"
+            )
+        result_memories = enhanced_memories[:top_k]
+        return result_memories
+
+    def memory_recreate_enhancement(
+        self,
+        query: str,
+        top_k: int,
+        text_memories: list[str],
+        retries: int,
+    ) -> list:
+        attempt = 0
+        text_memories = "\n".join([f"- [{i}] {mem}" for i, mem in enumerate(text_memories)])
+        prompt_name = "memory_recreate_enhancement"
+        prompt = self.build_prompt(
+            template_name=prompt_name, query=query, top_k=top_k, memories=text_memories
+        )
+
+        llm_response = None
+        while attempt <= max(0, retries) + 1:
+            try:
+                llm_response = self.process_llm.generate([{"role": "user", "content": prompt}])
+                processed_text_memories = parse_structured_output(content=llm_response)
+                logger.debug(
+                    f"[memory_recreate_enhancement]\n "
+                    f"- original memories: \n"
+                    f"{text_memories}\n"
+                    f"- final memories: \n"
+                    f"{processed_text_memories['answer']}"
+                )
+                return processed_text_memories["answer"]
+            except Exception as e:
+                attempt += 1
+                time.sleep(1)
+                logger.debug(
+                    f"[memory_recreate_enhancement] 🔁 retry {attempt}/{max(1, retries) + 1} failed: {e}"
+                )
+        logger.error(
+            f"Fail to run memory enhancement; prompt: {prompt};\n llm_response: {llm_response}",
+            exc_info=True,
+        )
+        raise ValueError("Fail to run memory enhancement")
+
+    def deep_search(
+        self,
+        query: str,
+        top_k: int,
+        info=None,
+        memory_type="All",
+        search_filter: dict | None = None,
+        user_name: str | None = None,
+        **kwargs,
+    ):
+        previous_retrieval_phrases = [query]
+        retrieved_memories = self.retrieve(
+            query=query,
+            user_name=user_name,
+            top_k=top_k,
+            mode=SearchMode.FAST,
+            memory_type=memory_type,
+            search_filter=search_filter,
+            info=info,
+        )
+        memories = self.post_retrieve(
+            retrieved_results=retrieved_memories,
+            top_k=top_k,
+            user_name=user_name,
+            info=info,
+        )
+        if len(memories) == 0:
+            logger.warning("Requirements not met; returning memories as-is.")
+            return memories
+
+        user_id = memories[0].metadata.user_id
+
+        mem_list, _ = self.tree_memories_to_text_memories(memories=memories)
+        retrieved_memories = copy.deepcopy(retrieved_memories)
+        rewritten_flag = False
+        for current_stage_id in range(self.thinking_stages + 1):
+            try:
+                # at last
+                if current_stage_id == self.thinking_stages:
+                    # eval to finish
+                    reason, can_answer = self.judge_memories(
+                        query=query,
+                        text_memories="- " + "\n- ".join(mem_list) + "\n",
+                    )
+
+                    logger.info(
+                        f"Final Stage: Stage {current_stage_id}; "
+                        f"previous retrieval phrases have been tried: {previous_retrieval_phrases}; "
+                        f"final can_answer: {can_answer}; reason: {reason}"
+                    )
+                    if rewritten_flag:
+                        enhanced_memories = self.get_final_memories(
+                            user_id=user_id, top_k=top_k, mem_list=mem_list
+                        )
+                    else:
+                        enhanced_memories = memories
+                    return enhanced_memories[:top_k]
+
+                can_answer, reason, retrieval_phrases = self.stage_retrieve(
+                    stage_id=current_stage_id + 1,
+                    query=query,
+                    previous_retrieval_phrases=previous_retrieval_phrases,
+                    text_memories="- " + "\n- ".join(mem_list) + "\n",
+                )
+                if can_answer:
+                    logger.info(
+                        f"Stage {current_stage_id}: determined answer can be provided, creating enhanced memories; reason: {reason}",
+                    )
+                    if rewritten_flag:
+                        enhanced_memories = self.get_final_memories(
+                            user_id=user_id, top_k=top_k, mem_list=mem_list
+                        )
+                    else:
+                        enhanced_memories = memories
+                    return enhanced_memories[:top_k]
+                else:
+                    previous_retrieval_phrases.extend(retrieval_phrases)
+                    logger.info(
+                        f"Start complementary retrieval for Stage {current_stage_id}; "
+                        f"previous retrieval phrases have been tried: {previous_retrieval_phrases}; "
+                        f"can_answer: {can_answer}; reason: {reason}"
+                    )
+                    logger.info(
+                        "Stage %d - Found %d new retrieval phrases",
+                        current_stage_id,
+                        len(retrieval_phrases),
+                    )
+                    # Search for additional memories based on retrieval phrases
+                    additional_retrieved_memories = []
+                    for phrase in retrieval_phrases:
+                        _retrieved_memories = self.retrieve(
+                            query=phrase,
+                            user_name=user_name,
+                            top_k=self.stage_retrieve_top,
+                            mode=SearchMode.FAST,
+                            memory_type=memory_type,
+                            search_filter=search_filter,
+                            info=info,
+                        )
+                        logger.info(
+                            "Found %d additional memories for phrase: '%s'",
+                            len(_retrieved_memories),
+                            phrase[:30] + "..." if len(phrase) > 30 else phrase,
+                        )
+                        additional_retrieved_memories.extend(_retrieved_memories)
+                    merged_memories = self.post_retrieve(
+                        retrieved_results=retrieved_memories + additional_retrieved_memories,
+                        top_k=top_k * 2,
+                        user_name=user_name,
+                        info=info,
+                    )
+                    rewritten_flag = True
+                    _mem_list, _ = self.tree_memories_to_text_memories(memories=merged_memories)
+                    mem_list = _mem_list
+                    mem_list = list(set(mem_list))
+                    mem_list = self.memory_recreate_enhancement(
+                        query=query,
+                        top_k=top_k,
+                        text_memories=mem_list,
+                        retries=self.max_retry_times,
+                    )
+                    logger.info(
+                        "After stage %d, total memories in list: %d",
+                        current_stage_id,
+                        len(mem_list),
+                    )
+
+            except Exception as e:
+                logger.error("Error in stage %d: %s", current_stage_id, str(e), exc_info=True)
+                # Continue to next stage instead of failing completely
+                continue
+        logger.error("Deep search failed, returning original memories")
+        return memories

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

@@ -0,0 +1,186 @@
+import threading
+
+import numpy as np
+
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+from memos.dependency import require_python_package
+from memos.log import get_logger
+from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import FastTokenizer
+from memos.utils import timed
+
+
+logger = get_logger(__name__)
+# Global model cache
+_CACHE_LOCK = threading.Lock()
+
+
+class EnhancedBM25:
+    """Enhanced BM25 with Spacy tokenization and TF-IDF reranking"""
+
+    @require_python_package(import_name="cachetools", install_command="pip install cachetools")
+    def __init__(self, tokenizer=None, en_model="en_core_web_sm", zh_model="zh_core_web_sm"):
+        """
+        Initialize Enhanced BM25 with memory management
+        """
+        if tokenizer is None:
+            self.tokenizer = FastTokenizer()
+        else:
+            self.tokenizer = tokenizer
+        self._current_tfidf = None
+
+        global _BM25_CACHE
+        from cachetools import LRUCache
+
+        _BM25_CACHE = LRUCache(maxsize=100)
+
+    def _tokenize_doc(self, text):
+        """
+        Tokenize a single document using SpacyTokenizer
+        """
+        return self.tokenizer.tokenize_mixed(text, lang="auto")
+
+    @require_python_package(import_name="rank_bm25", install_command="pip install rank_bm25")
+    def _prepare_corpus_data(self, corpus, corpus_name="default"):
+        from rank_bm25 import BM25Okapi
+
+        with _CACHE_LOCK:
+            if corpus_name in _BM25_CACHE:
+                print("hit::", corpus_name)
+                return _BM25_CACHE[corpus_name]
+            print("not hit::", corpus_name)
+
+            tokenized_corpus = [self._tokenize_doc(doc) for doc in corpus]
+            bm25_model = BM25Okapi(tokenized_corpus)
+            _BM25_CACHE[corpus_name] = bm25_model
+            return bm25_model
+
+    def clear_cache(self, corpus_name=None):
+        """Clear cache for specific corpus or clear all cache"""
+        with _CACHE_LOCK:
+            if corpus_name:
+                if corpus_name in _BM25_CACHE:
+                    del _BM25_CACHE[corpus_name]
+            else:
+                _BM25_CACHE.clear()
+
+    def get_cache_info(self):
+        """Get current cache information"""
+        with _CACHE_LOCK:
+            return {
+                "cache_size": len(_BM25_CACHE),
+                "max_cache_size": 100,
+                "cached_corpora": list(_BM25_CACHE.keys()),
+            }
+
+    def _search_docs(
+        self,
+        query: str,
+        corpus: list[str],
+        corpus_name="test",
+        top_k=50,
+        use_tfidf=False,
+        rerank_candidates_multiplier=2,
+        cleanup=False,
+    ):
+        """
+        Args:
+            query: Search query string
+            corpus: List of document texts
+            top_k: Number of top results to return
+            rerank_candidates_multiplier: Multiplier for candidate selection
+            cleanup: Whether to cleanup memory after search (default: True)
+        """
+        if not corpus:
+            return []
+
+        logger.info(f"Searching {len(corpus)} documents for query: '{query}'")
+
+        try:
+            # Prepare BM25 model
+            bm25_model = self._prepare_corpus_data(corpus, corpus_name=corpus_name)
+            tokenized_query = self._tokenize_doc(query)
+            tokenized_query = list(dict.fromkeys(tokenized_query))
+
+            # Get BM25 scores
+            bm25_scores = bm25_model.get_scores(tokenized_query)
+
+            # Select candidates
+            candidate_count = min(top_k * rerank_candidates_multiplier, len(corpus))
+            candidate_indices = np.argsort(bm25_scores)[-candidate_count:][::-1]
+            combined_scores = bm25_scores[candidate_indices]
+
+            if use_tfidf:
+                # Create TF-IDF for this search
+                tfidf = TfidfVectorizer(
+                    tokenizer=self._tokenize_doc, lowercase=False, token_pattern=None
+                )
+                tfidf_matrix = tfidf.fit_transform(corpus)
+
+                # TF-IDF reranking
+                query_vec = tfidf.transform([query])
+                tfidf_similarities = (
+                    (tfidf_matrix[candidate_indices] * query_vec.T).toarray().flatten()
+                )
+
+                # Combine scores
+                combined_scores = 0.7 * bm25_scores[candidate_indices] + 0.3 * tfidf_similarities
+
+            sorted_candidate_indices = candidate_indices[np.argsort(combined_scores)[::-1][:top_k]]
+            sorted_combined_scores = np.sort(combined_scores)[::-1][:top_k]
+
+            # build result list
+            bm25_recalled_results = []
+            for rank, (doc_idx, combined_score) in enumerate(
+                zip(sorted_candidate_indices, sorted_combined_scores, strict=False), 1
+            ):
+                bm25_score = bm25_scores[doc_idx]
+
+                candidate_pos = np.where(candidate_indices == doc_idx)[0][0]
+                tfidf_score = tfidf_similarities[candidate_pos] if use_tfidf else 0
+
+                bm25_recalled_results.append(
+                    {
+                        "text": corpus[doc_idx],
+                        "bm25_score": float(bm25_score),
+                        "tfidf_score": float(tfidf_score),
+                        "combined_score": float(combined_score),
+                        "rank": rank,
+                        "doc_index": int(doc_idx),
+                    }
+                )
+
+            logger.debug(f"Search completed: found {len(bm25_recalled_results)} results")
+            return bm25_recalled_results
+
+        except Exception as e:
+            logger.error(f"BM25 search failed: {e}")
+            return []
+        finally:
+            # Always cleanup if requested
+            if cleanup:
+                self._cleanup_memory()
+
+    @timed
+    def search(self, query: str, node_dicts: list[dict], corpus_name="default", **kwargs):
+        """
+        Search with BM25 and optional TF-IDF reranking
+        """
+        try:
+            corpus_list = []
+            for node_dict in node_dicts:
+                corpus_list.append(
+                    " ".join([node_dict["metadata"]["key"]] + node_dict["metadata"]["tags"])
+                )
+
+            recalled_results = self._search_docs(
+                query, corpus_list, corpus_name=corpus_name, **kwargs
+            )
+            bm25_searched_nodes = []
+            for item in recalled_results:
+                doc_idx = item["doc_index"]
+                bm25_searched_nodes.append(node_dicts[doc_idx])
+            return bm25_searched_nodes
+        except Exception as e:
+            logger.error(f"Error in bm25 search: {e}")
+            return []