From f909b1b6bcb5665ab0213f29db6908d8670e75e4 Mon Sep 17 00:00:00 2001 From: XianBW <36835909+XianBW@users.noreply.github.com> Date: Thu, 24 Apr 2025 18:22:51 +0800 Subject: [PATCH] feat: using different chat model in different part (#822) * using model in the chat_model_map in one tag * add replace timer to DS loop * fix CI * fix CI * add more custom config in chat_model_map * fix CI * fix CI * fix CI --------- Co-authored-by: Xu Yang --- rdagent/app/data_science/loop.py | 11 ++++++++--- rdagent/oai/backend/deprec.py | 26 +++++++++++++------------- rdagent/oai/backend/litellm.py | 32 +++++++++++++++++++++++--------- rdagent/oai/llm_conf.py | 2 +- rdagent/utils/workflow.py | 8 ++++++-- 5 files changed, 51 insertions(+), 28 deletions(-) diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py index f7cdeffe..79e282b1 100644 --- a/rdagent/app/data_science/loop.py +++ b/rdagent/app/data_science/loop.py @@ -221,9 +221,13 @@ class DataScienceRDLoop(RDLoop): @classmethod def load( - cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None, do_truncate: bool = False + cls, + path: Union[str, Path], + output_path: Optional[Union[str, Path]] = None, + do_truncate: bool = False, + replace_timer: bool = True, ) -> "LoopBase": - session = super().load(path, output_path, do_truncate) + session = super().load(path, output_path, do_truncate, replace_timer) if ( DS_RD_SETTING.enable_knowledge_base and DS_RD_SETTING.knowledge_base_version == "v1" @@ -255,6 +259,7 @@ def main( competition="bms-molecular-translation", do_truncate=True, timeout=None, + replace_timer=True, ): """ @@ -297,7 +302,7 @@ def main( if path is None: kaggle_loop = DataScienceRDLoop(DS_RD_SETTING) else: - kaggle_loop = DataScienceRDLoop.load(path, output_path, do_truncate) + kaggle_loop = DataScienceRDLoop.load(path, output_path, do_truncate, replace_timer) kaggle_loop.run(step_n=step_n, loop_n=loop_n, all_duration=timeout) diff --git a/rdagent/oai/backend/deprec.py b/rdagent/oai/backend/deprec.py index a96d8814..96bcf496 100644 --- a/rdagent/oai/backend/deprec.py +++ b/rdagent/oai/backend/deprec.py @@ -152,7 +152,7 @@ class DeprecBackend(APIBackend): self.gcr_endpoint_max_token = LLM_SETTINGS.gcr_endpoint_max_token if not os.environ.get("PYTHONHTTPSVERIFY", "") and hasattr(ssl, "_create_unverified_context"): ssl._create_default_https_context = ssl._create_unverified_context # noqa: SLF001 - self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map) + self.chat_model_map = LLM_SETTINGS.chat_model_map self.chat_model = LLM_SETTINGS.chat_model self.encoder = None elif LLM_SETTINGS.chat_use_azure_deepseek: @@ -160,7 +160,7 @@ class DeprecBackend(APIBackend): endpoint=LLM_SETTINGS.chat_azure_deepseek_endpoint, credential=AzureKeyCredential(LLM_SETTINGS.chat_azure_deepseek_key), ) - self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map) + self.chat_model_map = LLM_SETTINGS.chat_model_map self.encoder = None self.chat_model = "deepseek-R1" self.chat_stream = LLM_SETTINGS.chat_stream @@ -181,7 +181,7 @@ class DeprecBackend(APIBackend): ) self.chat_model = LLM_SETTINGS.chat_model - self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map) + self.chat_model_map = LLM_SETTINGS.chat_model_map self.encoder = self._get_encoder() self.chat_openai_base_url = LLM_SETTINGS.chat_openai_base_url self.embedding_openai_base_url = LLM_SETTINGS.embedding_openai_base_url @@ -303,19 +303,20 @@ class DeprecBackend(APIBackend): logger.info(self._build_log_messages(messages), tag="llm_messages") # TODO: fail to use loguru adaptor due to stream response + model = LLM_SETTINGS.chat_model temperature = LLM_SETTINGS.chat_temperature max_tokens = LLM_SETTINGS.chat_max_tokens frequency_penalty = LLM_SETTINGS.chat_frequency_penalty presence_penalty = LLM_SETTINGS.chat_presence_penalty - # Use index 4 to skip the current function and intermediate calls, - # and get the locals of the caller's frame. - caller_locals = inspect.stack()[4].frame.f_locals - if "self" in caller_locals: - tag = caller_locals["self"].__class__.__name__ - else: - tag = inspect.stack()[4].function - model = self.chat_model_map.get(tag, self.chat_model) + if self.chat_model_map: + for t, mc in self.chat_model_map.items(): + if t in logger._tag: + model = mc.get("model", model) + temperature = float(mc.get("temperature", temperature)) + if "max_tokens" in mc: + max_tokens = int(mc["max_tokens"]) + break finish_reason = None if self.use_llama2: @@ -394,7 +395,7 @@ class DeprecBackend(APIBackend): logger.info(f"{LogColors.CYAN}Think:{think_part}{LogColors.END}", tag="llm_messages") logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") else: - call_kwargs = dict( + call_kwargs: dict[str, Any] = dict( model=model, messages=messages, max_tokens=max_tokens, @@ -443,7 +444,6 @@ class DeprecBackend(APIBackend): logger.info( json.dumps( { - "tag": tag, "total_tokens": response.usage.total_tokens, "prompt_tokens": response.usage.prompt_tokens, "completion_tokens": response.usage.completion_tokens, diff --git a/rdagent/oai/backend/litellm.py b/rdagent/oai/backend/litellm.py index 65f272c4..9bd529e3 100644 --- a/rdagent/oai/backend/litellm.py +++ b/rdagent/oai/backend/litellm.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Literal, cast from litellm import ( completion, @@ -78,18 +78,32 @@ class LiteLLMAPIBackend(APIBackend): logger.info(self._build_log_messages(messages), tag="llm_messages") # Call LiteLLM completion + model = LITELLM_SETTINGS.chat_model + temperature = LITELLM_SETTINGS.chat_temperature + max_tokens = LITELLM_SETTINGS.chat_max_tokens + reasoning_effort = LITELLM_SETTINGS.reasoning_effort + + if LITELLM_SETTINGS.chat_model_map: + for t, mc in LITELLM_SETTINGS.chat_model_map.items(): + if t in logger._tag: + model = mc["model"] + if "temperature" in mc: + temperature = float(mc["temperature"]) + if "max_tokens" in mc: + max_tokens = int(mc["max_tokens"]) + if "reasoning_effort" in mc and mc["reasoning_effort"] in ["low", "medium", "high"]: + reasoning_effort = cast(Literal["low", "medium", "high"], mc["reasoning_effort"]) + break response = completion( - model=LITELLM_SETTINGS.chat_model, + model=model, messages=messages, stream=LITELLM_SETTINGS.chat_stream, - temperature=LITELLM_SETTINGS.chat_temperature, - max_tokens=LITELLM_SETTINGS.chat_max_tokens, - reasoning_effort=LITELLM_SETTINGS.reasoning_effort, + temperature=temperature, + max_tokens=max_tokens, + reasoning_effort=reasoning_effort, **kwargs, ) - logger.info( - f"{LogColors.GREEN}Using chat model{LogColors.END} {LITELLM_SETTINGS.chat_model}", tag="llm_messages" - ) + logger.info(f"{LogColors.GREEN}Using chat model{LogColors.END} {model}", tag="llm_messages") if LITELLM_SETTINGS.chat_stream: logger.info(f"{LogColors.BLUE}assistant:{LogColors.END}", tag="llm_messages") @@ -117,7 +131,7 @@ class LiteLLMAPIBackend(APIBackend): logger.info(f"{LogColors.BLUE}assistant:{LogColors.END} {finish_reason_str}\n{content}", tag="llm_messages") global ACC_COST - cost = completion_cost(model=LITELLM_SETTINGS.chat_model, messages=messages, completion=content) + cost = completion_cost(model=model, messages=messages, completion=content) ACC_COST += cost logger.info( f"Current Cost: ${float(cost):.10f}; Accumulated Cost: ${float(ACC_COST):.10f}; {finish_reason=}", diff --git a/rdagent/oai/llm_conf.py b/rdagent/oai/llm_conf.py index 474e8025..a3ecf83e 100644 --- a/rdagent/oai/llm_conf.py +++ b/rdagent/oai/llm_conf.py @@ -114,7 +114,7 @@ class LLMSettings(ExtendedBaseSettings): chat_azure_deepseek_endpoint: str = "" chat_azure_deepseek_key: str = "" - chat_model_map: str = "{}" + chat_model_map: dict[str, dict[str, str]] = {} LLM_SETTINGS = LLMSettings() diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py index 2abeacd8..da86b301 100644 --- a/rdagent/utils/workflow.py +++ b/rdagent/utils/workflow.py @@ -203,7 +203,11 @@ class LoopBase: @classmethod def load( - cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None, do_truncate: bool = False + cls, + path: Union[str, Path], + output_path: Optional[Union[str, Path]] = None, + do_truncate: bool = False, + replace_timer: bool = True, ) -> "LoopBase": path = Path(path) with path.open("rb") as f: @@ -223,7 +227,7 @@ class LoopBase: max_loop = max(session.loop_trace.keys()) logger.storage.truncate(time=session.loop_trace[max_loop][-1].end) - if session.timer.started: + if session.timer.started and replace_timer: RD_Agent_TIMER_wrapper.replace_timer(session.timer) RD_Agent_TIMER_wrapper.timer.restart_by_remain_time() return session