fix: code suffix

Author: fridayL

src/memos/reranker/base.py

@@ -2,7 +2,8 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
+
 
 if TYPE_CHECKING:
     from memos.memories.textual.item import TextualMemoryItem
@@ -21,4 +22,4 @@ def rerank(
         **kwargs,
     ) -> list[tuple[TextualMemoryItem, float]]:
         """Return top_k (item, score) sorted by score desc."""
-        raise NotImplementedError
+        raise NotImplementedError

src/memos/reranker/factory.py

@@ -8,8 +8,8 @@
 
 from .cosine_local import CosineLocalReranker
 from .http_bge import HTTPBGEReranker
-from .noop import NoopReranker
 from .http_bge_strategy import HTTPBGERerankerStrategy
+from .noop import NoopReranker
 
 
 if TYPE_CHECKING:

src/memos/reranker/http_bge_strategy.py

@@ -9,10 +9,9 @@
 import requests
 
 from memos.log import get_logger
-
-from .base import BaseReranker
 from memos.reranker.strategies import RerankerStrategyFactory
 
+from .base import BaseReranker
 
 
 logger = get_logger(__name__)
@@ -151,7 +150,9 @@ def rerank(
         if not graph_results:
             return []
 
-        tracker, original_items, documents = self.reranker_strategy.prepare_documents(query, graph_results, top_k)
+        tracker, original_items, documents = self.reranker_strategy.prepare_documents(
+            query, graph_results, top_k
+        )
 
         logger.info(
             f"[HTTPBGEWithSourceReranker] strategy: {self.reranker_strategy}, "
@@ -167,9 +168,7 @@ def rerank(
 
         try:
             # Make the HTTP request to the reranker service
-            resp = requests.post(
-                self.reranker_url, headers=headers, json=payload, timeout=30
-            )
+            resp = requests.post(self.reranker_url, headers=headers, json=payload, timeout=self.timeout)
             resp.raise_for_status()
             data = resp.json()
 
@@ -192,13 +191,13 @@ def rerank(
                         ranked_indices.append(idx)
                         scores.append(raw_score)
                 reconstructed_items = self.reranker_strategy.reconstruct_items(
-                    ranked_indices=ranked_indices, 
-                    scores=scores, 
-                    tracker=tracker, 
-                    original_items=original_items, 
+                    ranked_indices=ranked_indices,
+                    scores=scores,
+                    tracker=tracker,
+                    original_items=original_items,
                     top_k=top_k,
                     graph_results=graph_results,
-                    documents=documents
+                    documents=documents,
                 )
                 return reconstructed_items
 

src/memos/reranker/strategies/__init__.py

@@ -1,4 +1,4 @@
 from .factory import RerankerStrategyFactory
 
 
-__all__ = ["RerankerStrategyFactory"]
+__all__ = ["RerankerStrategyFactory"]

src/memos/reranker/strategies/base.py

@@ -1,8 +1,11 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
+from typing import Any
+
 from memos.memories.textual.item import TextualMemoryItem
+
 from .dialogue_common import DialogueRankingTracker
 
+
 class BaseRerankerStrategy(ABC):
     """Abstract interface for memory rerankers with concatenation strategy."""
 
@@ -16,21 +19,21 @@ def prepare_documents(
     ) -> tuple[DialogueRankingTracker, dict[str, Any], list[str]]:
         """
         Prepare documents for ranking based on the strategy.
-        
+
         Args:
             query: The search query
             graph_results: List of TextualMemoryItem objects to process
             top_k: Maximum number of items to return
             **kwargs: Additional strategy-specific parameters
-            
+
         Returns:
-            tuple[DialogueRankingTracker, dict[str, Any], list[str]]: 
+            tuple[DialogueRankingTracker, dict[str, Any], list[str]]:
             - Tracker: DialogueRankingTracker instance
             - original_items: Dict mapping memory_id to original TextualMemoryItem
             - documents: List of text documents ready for ranking
         """
         raise NotImplementedError
-            
+
     @abstractmethod
     def reconstruct_items(
         self,
@@ -43,16 +46,16 @@ def reconstruct_items(
     ) -> list[tuple[TextualMemoryItem, float]]:
         """
         Reconstruct TextualMemoryItem objects from ranked results.
-        
+
         Args:
             ranked_indices: List of indices sorted by relevance
             scores: Corresponding relevance scores
             tracker: DialogueRankingTracker instance
             original_items: Dict mapping memory_id to original TextualMemoryItem
             top_k: Maximum number of items to return
             **kwargs: Additional strategy-specific parameters
-            
+
         Returns:
             List of (reconstructed_memory_item, aggregated_score) tuples
         """
-        raise NotImplementedError
+        raise NotImplementedError

src/memos/reranker/strategies/concat_background.py

@@ -1,22 +1,25 @@
 # memos/reranker/strategies/single_turn.py
 from __future__ import annotations
+
 import re
+
 from typing import Any
-from collections import defaultdict
-from copy import deepcopy
+
 from .base import BaseRerankerStrategy
-from .dialogue_common import DialogueRankingTracker, strip_memory_tags, extract_content
+from .dialogue_common import DialogueRankingTracker
+
 
 _TAG1 = re.compile(r"^\s*\[[^\]]*\]\s*")
 
+
 class ConcatBackgroundStrategy(BaseRerankerStrategy):
     """
     Concat background strategy.
-    
+
     This strategy processes dialogue pairs by concatenating background and
     user and assistant messages into single strings for ranking. Each dialogue pair becomes a
     separate document for ranking.
-    """ 
+    """
 
     def prepare_documents(
         self,
@@ -27,33 +30,33 @@ def prepare_documents(
     ) -> tuple[DialogueRankingTracker, dict[str, Any], list[str]]:
         """
         Prepare documents based on single turn concatenation strategy.
-        
+
         Args:
             query: The search query
             graph_results: List of graph results
             top_k: Maximum number of items to return
-            
+
         Returns:
-            tuple[DialogueRankingTracker, dict[str, Any], list[str]]: 
+            tuple[DialogueRankingTracker, dict[str, Any], list[str]]:
             - Tracker: DialogueRankingTracker instance
             - original_items: Dict mapping memory_id to original TextualMemoryItem
             - documents: List of text documents ready for ranking
         """
-        
+
         original_items = {}
         tracker = DialogueRankingTracker()
         documents = []
         for item in graph_results:
             memory = getattr(item, "memory", None)
             if isinstance(memory, str):
                 memory = _TAG1.sub("", memory)
-            
+
             background = ""
             if hasattr(item, "metadata") and hasattr(item.metadata, "background"):
                 background = getattr(item.metadata, "background", "")
                 if not isinstance(background, str):
                     background = ""
-            
+
             documents.append(f"{memory}\n{background}")
         return tracker, original_items, documents
 
@@ -68,26 +71,24 @@ def reconstruct_items(
     ) -> list[tuple[Any, float]]:
         """
         Reconstruct TextualMemoryItem objects from ranked dialogue pairs.
-        
+
         Args:
             ranked_indices: List of dialogue pair indices sorted by relevance
             scores: Corresponding relevance scores
             tracker: DialogueRankingTracker instance
             original_items: Dict mapping memory_id to original TextualMemoryItem
             top_k: Maximum number of items to return
-            
+
         Returns:
             List of (reconstructed_memory_item, aggregated_score) tuples
         """
-        graph_results = kwargs.get("graph_results", None)
-        documents = kwargs.get("documents", None)
+        graph_results = kwargs.get("graph_results")
+        documents = kwargs.get("documents")
         reconstructed_items = []
         for idx in ranked_indices:
             item = graph_results[idx]
             item.memory = f"{item.memory}\n{documents[idx]}"
             reconstructed_items.append((item, scores[idx]))
-            
+
         reconstructed_items.sort(key=lambda x: x[1], reverse=True)
         return reconstructed_items[:top_k]
-
-

src/memos/reranker/strategies/dialogue_common.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 import re
+
 from typing import Any, Literal
+
 from pydantic import BaseModel
-from memos.memories.textual.item import SourceMessage
-from memos.memories.textual.item import TextualMemoryItem
+
+from memos.memories.textual.item import SourceMessage, TextualMemoryItem
+
 
 # Strip a leading "[...]" tag (e.g., "[2025-09-01] ..." or "[meta] ...")
 # before sending text to the reranker. This keeps inputs clean and
@@ -17,10 +20,11 @@ def strip_memory_tags(item: TextualMemoryItem) -> str:
     memory = _TAG1.sub("", m) if isinstance((m := getattr(item, "memory", None)), str) else m
     return memory
 
+
 def extract_content(msg: dict[str, Any] | str) -> str:
     """Extract content from message, handling both string and dict formats."""
     if isinstance(msg, dict):
-        return msg.get('content', str(msg))
+        return msg.get("content", str(msg))
     if isinstance(msg, SourceMessage):
         return msg.content
     return str(msg)
@@ -33,7 +37,7 @@ class DialoguePair(BaseModel):
     memory_id: str  # ID of the source TextualMemoryItem
     memory: str
     pair_index: int  # Index of this pair within the source memory's dialogue
-    user_msg: str | dict[str, Any] | SourceMessage # User message content
+    user_msg: str | dict[str, Any] | SourceMessage  # User message content
     assistant_msg: str | dict[str, Any] | SourceMessage  # Assistant message content
     combined_text: str  # The concatenated text used for ranking
     chat_time: str | None = None
@@ -56,14 +60,14 @@ def __init__(self):
         self.dialogue_pairs: list[DialoguePair] = []
 
     def add_dialogue_pair(
-        self, 
-        memory_id: str, 
+        self,
+        memory_id: str,
         pair_index: int,
-        user_msg: str | dict[str, Any], 
+        user_msg: str | dict[str, Any],
         assistant_msg: str | dict[str, Any],
         memory: str,
         chat_time: str | None = None,
-        concat_format: Literal["user_assistant", "user_only"] = "user_assistant"
+        concat_format: Literal["user_assistant", "user_only"] = "user_assistant",
     ) -> str:
         """Add a dialogue pair and return its unique ID."""
         user_content = extract_content(user_msg)
@@ -85,7 +89,7 @@ def add_dialogue_pair(
             assistant_msg=assistant_msg,
             combined_text=combined_text,
             memory=memory,
-            chat_time=chat_time
+            chat_time=chat_time,
         )
 
         self.dialogue_pairs.append(dialogue_pair)

src/memos/reranker/strategies/factory.py

@@ -1,15 +1,18 @@
 # memos/reranker/factory.py
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any,ClassVar
-from .single_turn import SingleTurnStrategy
+from typing import TYPE_CHECKING, Any, ClassVar
+
 from .concat_background import ConcatBackgroundStrategy
+from .single_turn import SingleTurnStrategy
 from .singleturn_outmem import SingleTurnOutMemStrategy
 
+
 if TYPE_CHECKING:
     from .base import BaseRerankerStrategy
 
-class RerankerStrategyFactory():
+
+class RerankerStrategyFactory:
     """Factory class for creating reranker strategy instances."""
 
     backend_to_class: ClassVar[dict[str, Any]] = {
@@ -19,9 +22,7 @@ class RerankerStrategyFactory():
     }
 
     @classmethod
-    def from_config(
-        cls, config_factory: str = "single_turn"
-    ) -> BaseRerankerStrategy:
+    def from_config(cls, config_factory: str = "single_turn") -> BaseRerankerStrategy:
         if config_factory not in cls.backend_to_class:
             raise ValueError(f"Invalid backend: {config_factory}")
         strategy_class = cls.backend_to_class[config_factory]