feat: pass reranker from tree config

Author: CaralHsi

examples/basic_modules/tree_textual_memory_reranker.py

@@ -7,6 +7,7 @@
 from memos.configs.graph_db import GraphDBConfigFactory
 from memos.configs.internet_retriever import InternetRetrieverConfigFactory
 from memos.configs.llm import LLMConfigFactory
+from memos.configs.reranker import RerankerConfigFactory
 from memos.configs.vec_db import VectorDBConfigFactory
 from memos.exceptions import ConfigurationError
 
@@ -151,6 +152,10 @@ class TreeTextMemoryConfig(BaseTextMemoryConfig):
         default_factory=EmbedderConfigFactory,
         description="Embedder configuration for the memory embedding",
     )
+    reranker: RerankerConfigFactory | None = Field(
+        None,
+        description="Reranker configuration (optional, defaults to cosine_local).",
+    )
     graph_db: GraphDBConfigFactory = Field(
         ...,
         default_factory=GraphDBConfigFactory,

src/memos/configs/memory.py

@@ -8,6 +8,7 @@
 from typing import Any
 
 from memos.configs.memory import TreeTextMemoryConfig
+from memos.configs.reranker import RerankerConfigFactory
 from memos.embedders.factory import EmbedderFactory, OllamaEmbedder
 from memos.graph_dbs.factory import GraphStoreFactory, Neo4jGraphDB
 from memos.llms.factory import AzureLLM, LLMFactory, OllamaLLM, OpenAILLM
@@ -19,6 +20,7 @@
     InternetRetrieverFactory,
 )
 from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
+from memos.reranker.factory import RerankerFactory
 from memos.types import MessageList
 
 
@@ -39,6 +41,20 @@ def __init__(self, config: TreeTextMemoryConfig):
         )
         self.embedder: OllamaEmbedder = EmbedderFactory.from_config(config.embedder)
         self.graph_store: Neo4jGraphDB = GraphStoreFactory.from_config(config.graph_db)
+        if config.reranker is None:
+            default_cfg = RerankerConfigFactory.model_validate(
+                {
+                    "backend": "cosine_local",
+                    "config": {
+                        "level_weights": {"topic": 1.0, "concept": 1.0, "fact": 1.0},
+                        "level_field": "background",
+                    },
+                }
+            )
+            self.reranker = RerankerFactory.from_config(default_cfg)
+        else:
+            self.reranker = RerankerFactory.from_config(config.reranker)
+
         self.is_reorganize = config.reorganize
 
         self.memory_manager: MemoryManager = MemoryManager(
@@ -131,6 +147,7 @@ def search(
                 self.dispatcher_llm,
                 self.graph_store,
                 self.embedder,
+                self.reranker,
                 internet_retriever=None,
                 moscube=moscube,
             )
@@ -139,6 +156,7 @@ def search(
                 self.dispatcher_llm,
                 self.graph_store,
                 self.embedder,
+                self.reranker,
                 internet_retriever=self.internet_retriever,
                 moscube=moscube,
             )

src/memos/memories/textual/tree.py

@@ -8,12 +8,12 @@
 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
 
 
@@ -26,6 +26,7 @@ def __init__(
         dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM,
         graph_store: Neo4jGraphDB,
         embedder: OllamaEmbedder,
+        reranker: BaseReranker,
         internet_retriever: InternetRetrieverFactory | None = None,
         moscube: bool = False,
     ):
@@ -34,7 +35,7 @@ def __init__(
 
         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