From 623da17ac14502fbdc5553cf272bea58d8df895b Mon Sep 17 00:00:00 2001 From: harvey_xiang Date: Tue, 2 Dec 2025 19:12:05 +0800 Subject: [PATCH 1/3] feat: timer add log args --- src/memos/llms/openai.py | 2 +- src/memos/utils.py | 39 +++++++++++++++++++++++++++++++-------- 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/src/memos/llms/openai.py b/src/memos/llms/openai.py index 19d7a60fe..c45038e9d 100644 --- a/src/memos/llms/openai.py +++ b/src/memos/llms/openai.py @@ -28,7 +28,7 @@ def __init__(self, config: OpenAILLMConfig): ) logger.info("OpenAI LLM instance initialized") - @timed(log=True, log_prefix="OpenAI LLM") + @timed(log=True, log_prefix="OpenAI LLM", log_args=["model_name_or_path"]) def generate(self, messages: MessageList, **kwargs) -> str: """Generate a response from OpenAI LLM, optionally overriding generation params.""" response = self.client.chat.completions.create( diff --git a/src/memos/utils.py b/src/memos/utils.py index 4b1a59834..a330f0e79 100644 --- a/src/memos/utils.py +++ b/src/memos/utils.py @@ -1,3 +1,4 @@ +import functools import time from memos.log import get_logger @@ -6,20 +7,42 @@ logger = get_logger(__name__) -def timed(func=None, *, log=True, log_prefix=""): - """Decorator to measure and optionally log time of retrieval steps. - - Can be used as @timed or @timed(log=True) +def timed(func=None, *, log=True, log_prefix="", log_args=None): + """ + Parameters: + - log: enable timing logs (default True) + - log_prefix: prefix; falls back to function name + - log_args: names to include in logs (str or list/tuple of str). + Value priority: kwargs → args[0].config. (if available). + Non-string items are ignored. + + Examples: + - @timed(log=True, log_prefix="OpenAI LLM", log_args=["model_name_or_path", "temperature"]) + - @timed(log=True, log_prefix="OpenAI LLM", log_args=["temperature"]) + - @timed() # defaults """ def decorator(fn): + @functools.wraps(fn) def wrapper(*args, **kwargs): start = time.perf_counter() result = fn(*args, **kwargs) - elapsed = time.perf_counter() - start - elapsed_ms = elapsed * 1000.0 - if log: - logger.info(f"[TIMER] {log_prefix or fn.__name__} took {elapsed_ms:.0f} ms") + elapsed_ms = (time.perf_counter() - start) * 1000.0 + ctx_str = "" + + if log is not True: + return result + + if log_args: + ctx_parts = [] + for key in log_args: + val = kwargs.get(key) + ctx_parts.append(f"{key}={val}") + ctx_str = f" [{', '.join(ctx_parts)}]" + logger.info( + f"[TIMER] {log_prefix or fn.__name__} took {elapsed_ms:.0f} ms, args: {ctx_str}" + ) + return result return wrapper From d3d7602e028edf0fbae4ac7ee0a7b5e9ea0ce97a Mon Sep 17 00:00:00 2001 From: harvey_xiang Date: Tue, 2 Dec 2025 19:15:12 +0800 Subject: [PATCH 2/3] feat: timer add log args --- src/memos/utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/memos/utils.py b/src/memos/utils.py index a330f0e79..bdde53904 100644 --- a/src/memos/utils.py +++ b/src/memos/utils.py @@ -1,4 +1,3 @@ -import functools import time from memos.log import get_logger @@ -23,7 +22,6 @@ def timed(func=None, *, log=True, log_prefix="", log_args=None): """ def decorator(fn): - @functools.wraps(fn) def wrapper(*args, **kwargs): start = time.perf_counter() result = fn(*args, **kwargs) From 51ec395fda1fe327243fd88786d69028168b7b5a Mon Sep 17 00:00:00 2001 From: harvey_xiang Date: Tue, 2 Dec 2025 19:56:19 +0800 Subject: [PATCH 3/3] feat: timer add log args --- src/memos/embedders/universal_api.py | 6 +++++- src/memos/reranker/http_bge.py | 4 +++- src/memos/utils.py | 17 ++++++++++++----- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/src/memos/embedders/universal_api.py b/src/memos/embedders/universal_api.py index f39ffaa58..79a5d9ea6 100644 --- a/src/memos/embedders/universal_api.py +++ b/src/memos/embedders/universal_api.py @@ -30,7 +30,11 @@ def __init__(self, config: UniversalAPIEmbedderConfig): else: raise ValueError(f"Embeddings unsupported provider: {self.provider}") - @timed(log=True, log_prefix="model_timed_embedding") + @timed( + log=True, + log_prefix="model_timed_embedding", + log_extra_args={"model_name_or_path": "text-embedding-3-large"}, + ) def embed(self, texts: list[str]) -> list[list[float]]: if self.provider == "openai" or self.provider == "azure": try: diff --git a/src/memos/reranker/http_bge.py b/src/memos/reranker/http_bge.py index 764b53032..29f41e38f 100644 --- a/src/memos/reranker/http_bge.py +++ b/src/memos/reranker/http_bge.py @@ -119,7 +119,9 @@ def __init__( self.warn_unknown_filter_keys = bool(warn_unknown_filter_keys) self._warned_missing_keys: set[str] = set() - @timed(log=True, log_prefix="model_timed_rerank") + @timed( + log=True, log_prefix="model_timed_rerank", log_extra_args={"model_name_or_path": "reranker"} + ) def rerank( self, query: str, diff --git a/src/memos/utils.py b/src/memos/utils.py index bdde53904..6671d88b7 100644 --- a/src/memos/utils.py +++ b/src/memos/utils.py @@ -6,7 +6,7 @@ logger = get_logger(__name__) -def timed(func=None, *, log=True, log_prefix="", log_args=None): +def timed(func=None, *, log=True, log_prefix="", log_args=None, log_extra_args=None): """ Parameters: - log: enable timing logs (default True) @@ -27,19 +27,26 @@ def wrapper(*args, **kwargs): result = fn(*args, **kwargs) elapsed_ms = (time.perf_counter() - start) * 1000.0 ctx_str = "" + ctx_parts = [] if log is not True: return result if log_args: - ctx_parts = [] for key in log_args: val = kwargs.get(key) ctx_parts.append(f"{key}={val}") ctx_str = f" [{', '.join(ctx_parts)}]" - logger.info( - f"[TIMER] {log_prefix or fn.__name__} took {elapsed_ms:.0f} ms, args: {ctx_str}" - ) + + if log_extra_args: + ctx_parts.extend([f"{key}={val}" for key, val in log_extra_args.items()]) + + if ctx_parts: + ctx_str = f" [{', '.join(ctx_parts)}]" + + logger.info( + f"[TIMER] {log_prefix or fn.__name__} took {elapsed_ms:.0f} ms, args: {ctx_str}" + ) return result