diff --git a/src/memos/embedders/universal_api.py b/src/memos/embedders/universal_api.py index fc51cf073..583a02acb 100644 --- a/src/memos/embedders/universal_api.py +++ b/src/memos/embedders/universal_api.py @@ -26,7 +26,7 @@ def __init__(self, config: UniversalAPIEmbedderConfig): else: raise ValueError(f"Embeddings unsupported provider: {self.provider}") - @timed(log=True, log_prefix="EmbedderAPI") + @timed(log=True, log_prefix="model_timed_embedding") def embed(self, texts: list[str]) -> list[list[float]]: if self.provider == "openai" or self.provider == "azure": try: diff --git a/src/memos/llms/openai.py b/src/memos/llms/openai.py index 1a1703340..da55ae593 100644 --- a/src/memos/llms/openai.py +++ b/src/memos/llms/openai.py @@ -1,5 +1,6 @@ import hashlib import json +import time from collections.abc import Generator from typing import ClassVar @@ -57,12 +58,15 @@ def clear_cache(cls): cls._instances.clear() logger.info("OpenAI LLM instance cache cleared") - @timed(log=True, log_prefix="OpenAI LLM") + @timed(log=True, log_prefix="model_timed_openai") def generate(self, messages: MessageList, **kwargs) -> str: """Generate a response from OpenAI LLM, optionally overriding generation params.""" temperature = kwargs.get("temperature", self.config.temperature) max_tokens = kwargs.get("max_tokens", self.config.max_tokens) top_p = kwargs.get("top_p", self.config.top_p) + start_time = time.time() + logger.info(f"openai model request start, model_name: {self.config.model_name_or_path}") + response = self.client.chat.completions.create( model=self.config.model_name_or_path, messages=messages, @@ -71,7 +75,11 @@ def generate(self, messages: MessageList, **kwargs) -> str: max_tokens=max_tokens, top_p=top_p, ) - logger.info(f"Response from OpenAI: {response.model_dump_json()}") + + end_time = time.time() + logger.info( + f"openai model request end, time_cost: {end_time - start_time:.0f} ms, response from OpenAI: {response.model_dump_json()}" + ) response_content = response.choices[0].message.content if self.config.remove_think_prefix: return remove_thinking_tags(response_content) diff --git a/src/memos/reranker/http_bge.py b/src/memos/reranker/http_bge.py index 41011df14..db5a51fc2 100644 --- a/src/memos/reranker/http_bge.py +++ b/src/memos/reranker/http_bge.py @@ -119,7 +119,7 @@ def __init__( self.warn_unknown_filter_keys = bool(warn_unknown_filter_keys) self._warned_missing_keys: set[str] = set() - @timed(log=True, log_prefix="RerankerAPI") + @timed(log=True, log_prefix="model_timed_rerank") def rerank( self, query: str, diff --git a/src/memos/utils.py b/src/memos/utils.py index 9ae27bb81..4b1a59834 100644 --- a/src/memos/utils.py +++ b/src/memos/utils.py @@ -17,8 +17,9 @@ 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:.2f} seconds") + logger.info(f"[TIMER] {log_prefix or fn.__name__} took {elapsed_ms:.0f} ms") return result return wrapper