diff --git a/src/memos/api/handlers/formatters_handler.py b/src/memos/api/handlers/formatters_handler.py index 81b957005..29e376d33 100644 --- a/src/memos/api/handlers/formatters_handler.py +++ b/src/memos/api/handlers/formatters_handler.py @@ -34,7 +34,7 @@ def to_iter(running: Any) -> list[Any]: def format_memory_item( - memory_data: Any, include_embedding: bool = False, save_sources: bool = False + memory_data: Any, include_embedding: bool = False, save_sources: bool = True ) -> dict[str, Any]: """ Format a single memory item for API response. @@ -186,6 +186,13 @@ def rerank_knowledge_mem( # rerank set unuse reranked_knowledge_mem = knowledge_mem + # Sort by relativity in descending order + reranked_knowledge_mem = sorted( + reranked_knowledge_mem, + key=lambda item: item.get("metadata", {}).get("relativity", 0.0), + reverse=True, + ) + # TODO revoke sources replace memory value for item in reranked_knowledge_mem: item["memory"] = item["metadata"]["sources"][0]["content"] diff --git a/src/memos/api/handlers/search_handler.py b/src/memos/api/handlers/search_handler.py index 3774410dc..e5af52f87 100644 --- a/src/memos/api/handlers/search_handler.py +++ b/src/memos/api/handlers/search_handler.py @@ -5,9 +5,12 @@ using dependency injection for better modularity and testability. """ +import time + from typing import Any from memos.api.handlers.base_handler import BaseHandler, HandlerDependencies +from memos.api.handlers.formatters_handler import rerank_knowledge_mem from memos.api.product_models import APISearchRequest, SearchResponse from memos.log import get_logger from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import ( @@ -69,6 +72,18 @@ def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse # Restore original top_k for downstream logic or response metadata search_req.top_k = original_top_k + start_time = time.time() + text_mem = results["text_mem"] + results["text_mem"] = rerank_knowledge_mem( + self.reranker, + query=search_req.query, + text_mem=text_mem, + top_k=original_top_k, + file_mem_proportion=0.5, + ) + rerank_time = time.time() - start_time + + self.logger.info(f"[Knowledge_replace_memory_time] Rerank time: {rerank_time} seconds") self.logger.info( f"[SearchHandler] Final search results: count={len(results)} results={results}" )