feat: integrate azure deepseek r1 (#591)

* fix several task & integrate deepseek R1

* fix CI

---------

Co-authored-by: Xu Yang <xuyang1@microsoft.com>

pyproject.toml

@@ -84,6 +84,7 @@ ignore = [
   "ANN401",
   "D",
   "ERA001",
+  "EXE002",
   "FIX",
   "INP001",
   "PGH",
@@ -91,7 +92,7 @@ ignore = [
   "S101",
   "S301",
   "T20",
-  "TC003",
+  "TCH003",
   "TD",
 ]
 select = ["ALL"]

rdagent/app/data_science/loop.py

@@ -5,10 +5,15 @@ import fire
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
 from rdagent.components.coder.data_science.ensemble import EnsembleCoSTEER
+from rdagent.components.coder.data_science.ensemble.exp import EnsembleTask
 from rdagent.components.coder.data_science.feature import FeatureCoSTEER
+from rdagent.components.coder.data_science.feature.exp import FeatureTask
 from rdagent.components.coder.data_science.model import ModelCoSTEER
+from rdagent.components.coder.data_science.model.exp import ModelTask
 from rdagent.components.coder.data_science.raw_data_loader import DataLoaderCoSTEER
+from rdagent.components.coder.data_science.raw_data_loader.exp import DataLoaderTask
 from rdagent.components.coder.data_science.workflow import WorkflowCoSTEER
+from rdagent.components.coder.data_science.workflow.exp import WorkflowTask
 from rdagent.components.workflow.conf import BasePropSetting
 from rdagent.components.workflow.rd_loop import RDLoop
 from rdagent.core.exception import CoderError, RunnerError
@@ -70,15 +75,15 @@ class DataScienceRDLoop(RDLoop):
         exp = prev_out["direct_exp_gen"]
         for tasks in exp.pending_tasks_list:
             exp.sub_tasks = tasks
-            if exp.hypothesis.component == "DataLoadSpec":
+            if isinstance(exp.sub_tasks[0], DataLoaderTask):
                 exp = self.data_loader_coder.develop(exp)
-            elif exp.hypothesis.component == "FeatureEng":
+            elif isinstance(exp.sub_tasks[0], FeatureTask):
                 exp = self.feature_coder.develop(exp)
-            elif exp.hypothesis.component == "Model":
+            elif isinstance(exp.sub_tasks[0], ModelTask):
                 exp = self.model_coder.develop(exp)
-            elif exp.hypothesis.component == "Ensemble":
+            elif isinstance(exp.sub_tasks[0], EnsembleTask):
                 exp = self.ensemble_coder.develop(exp)
-            elif exp.hypothesis.component == "Workflow":
+            elif isinstance(exp.sub_tasks[0], WorkflowTask):
                 exp = self.workflow_coder.develop(exp)
             else:
                 raise NotImplementedError(f"Unsupported component in DataScienceRDLoop: {exp.hypothesis.component}")
@@ -118,7 +123,7 @@ class DataScienceRDLoop(RDLoop):
                     ExperimentFeedback.from_exception(e),
                 )
             )
-            if len(self.trace.hist) >= DS_RD_SETTING.consecutive_errors:
+            if self.trace.sota_experiment() is None and len(self.trace.hist) >= DS_RD_SETTING.consecutive_errors:
                 trace_exp_next_component_list = [
                     exp.next_component_required() for exp, _ in self.trace.hist[-DS_RD_SETTING.consecutive_errors :]
                 ]

rdagent/components/coder/data_science/ensemble/eval.py

@@ -50,7 +50,7 @@ class EnsembleCoSTEEREvaluator(CoSTEEREvaluator):
         }
         de = DockerEnv(conf=ds_docker_conf)
 
-        fname = "ensemble_test.txt"
+        fname = "test/ensemble_test.txt"
         test_code = (DIRNAME / "eval_tests" / "ensemble_test.txt").read_text()
         test_code = (
             Environment(undefined=StrictUndefined)

rdagent/components/coder/data_science/ensemble/exp.py

@@ -7,4 +7,7 @@ from typing import Dict, Optional
 from rdagent.components.coder.CoSTEER.task import CoSTEERTask
 from rdagent.core.utils import cache_with_pickle
 
-EnsembleTask = CoSTEERTask
+
+# Because we use isinstance to distinguish between different types of tasks, we need to use sub classes to represent different types of tasks
+class EnsembleTask(CoSTEERTask):
+    pass

rdagent/components/coder/data_science/feature/eval.py

@@ -52,7 +52,7 @@ class FeatureCoSTEEREvaluator(CoSTEEREvaluator):
         de = DockerEnv(conf=ds_docker_conf)
 
         # TODO: do we need to clean the generated temporary content?
-        fname = "feature_test.py"
+        fname = "test/feature_test.py"
         test_code = (DIRNAME / "eval_tests" / "feature_test.txt").read_text()
         implementation.inject_files(**{fname: test_code})
 

rdagent/components/coder/data_science/feature/exp.py

@@ -7,4 +7,7 @@ from typing import Dict, Optional
 from rdagent.components.coder.CoSTEER.task import CoSTEERTask
 from rdagent.core.utils import cache_with_pickle
 
-FeatureTask = CoSTEERTask
+
+# Because we use isinstance to distinguish between different types of tasks, we need to use sub classes to represent different types of tasks
+class FeatureTask(CoSTEERTask):
+    pass

rdagent/components/coder/data_science/model/eval.py

@@ -62,7 +62,7 @@ class ModelGeneralCaseSpecEvaluator(CoSTEEREvaluator):
         }
         de = DockerEnv(conf=ds_docker_conf)
 
-        fname = "model_test.py"
+        fname = "test/model_test.py"
         test_code = (
             (DIRNAME / "eval_tests" / "model_test.txt").read_text().replace("model01", target_task.name)
         )  # only check the model changed this time

rdagent/components/coder/data_science/model/exp.py

@@ -11,6 +11,7 @@ from rdagent.oai.llm_utils import md5_hash
 from rdagent.utils.env import DockerEnv, DSDockerConf
 
 
+# Because we use isinstance to distinguish between different types of tasks, we need to use sub classes to represent different types of tasks
 class ModelTask(CoSTEERTask):
     def __init__(
         self,

rdagent/components/coder/data_science/raw_data_loader/eval.py

@@ -54,7 +54,7 @@ class DataLoaderCoSTEEREvaluator(CoSTEEREvaluator):
         de = DockerEnv(conf=ds_docker_conf)
 
         # TODO: do we need to clean the generated temporary content?
-        fname = "data_loader_test.py"
+        fname = "test/data_loader_test.py"
         test_code = (DIRNAME / "eval_tests" / "data_loader_test.txt").read_text()
         implementation.inject_files(**{fname: test_code})
         stdout = implementation.execute(env=de, entry=f"python {fname}")

rdagent/components/coder/data_science/raw_data_loader/exp.py

@@ -12,3 +12,8 @@ from rdagent.utils.agent.tpl import T
 from rdagent.utils.env import DockerEnv, DSDockerConf
 
 DataLoaderTask = CoSTEERTask
+
+
+# Because we use isinstance to distinguish between different types of tasks, we need to use sub classes to represent different types of tasks
+class DataLoaderTask(CoSTEERTask):
+    pass

rdagent/components/coder/data_science/workflow/eval.py

@@ -98,9 +98,9 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator):
             stdout += "\nSubmission file (submission.csv) is not generated."
         else:
             base_check_code = (DIRNAME / "eval_tests" / "submission_format_test.txt").read_text()
-            implementation.inject_files(**{"submission_format_test.py": base_check_code})
+            implementation.inject_files(**{"test/submission_format_test.py": base_check_code})
             # stdout += "----Submission Check 1-----\n"
-            stdout += implementation.execute(env=de, entry="python submission_format_test.py")
+            stdout += implementation.execute(env=de, entry="python test/submission_format_test.py")
 
             # MLEBench Check
             # !!! Since we are running on a sampled dataset, mlebench check is not required.
@@ -109,9 +109,9 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator):
             #     .read_text()
             #     .replace("<competition_id>", self.scen.competition)
             # )
-            # implementation.inject_files(**{"mle_submission_format_test.py": mle_check_code})
+            # implementation.inject_files(**{"test/mle_submission_format_test.py": mle_check_code})
             # stdout += "----Submission Check 2-----\n"
-            # stdout += implementation.execute(env=mde, entry=f"python mle_submission_format_test.py")
+            # stdout += implementation.execute(env=mde, entry=f"python test/mle_submission_format_test.py")
 
         system_prompt = T(".prompts:workflow_eval.system").r(
             scenario=self.scen.get_scenario_all_desc(),

rdagent/components/coder/data_science/workflow/exp.py

@@ -7,4 +7,7 @@ from typing import Dict, Optional
 from rdagent.components.coder.CoSTEER.task import CoSTEERTask
 from rdagent.core.utils import cache_with_pickle
 
-WorkflowTask = CoSTEERTask
+
+# Because we use isinstance to distinguish between different types of tasks, we need to use sub classes to represent different types of tasks
+class WorkflowTask(CoSTEERTask):
+    pass

rdagent/oai/backend/deprec.py

@@ -44,6 +44,19 @@ except ImportError:
     if LLM_SETTINGS.use_llama2:
         logger.warning("llama is not installed.")
 
+try:
+    from azure.ai.inference import ChatCompletionsClient
+    from azure.ai.inference.models import (
+        AssistantMessage,
+        ChatRequestMessage,
+        SystemMessage,
+        UserMessage,
+    )
+    from azure.core.credentials import AzureKeyCredential
+except ImportError:
+    if LLM_SETTINGS.chat_use_azure_deepseek:
+        logger.warning("azure.ai.inference or azure.core.credentials is not installed.")
+
 
 class ConvManager:
     """
@@ -299,6 +312,15 @@ class DeprecBackend(APIBackend):
             self.chat_model_map = json.loads(LLM_SETTINGS.chat_model_map)
             self.chat_model = LLM_SETTINGS.chat_model if chat_model is None else chat_model
             self.encoder = None
+        elif LLM_SETTINGS.chat_use_azure_deepseek:
+            self.client = ChatCompletionsClient(
+                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.encoder = None
+            self.chat_model = "deepseek-R1"
+            self.chat_stream = LLM_SETTINGS.chat_stream
         else:
             self.chat_use_azure = LLM_SETTINGS.chat_use_azure or LLM_SETTINGS.use_azure
             self.embedding_use_azure = LLM_SETTINGS.embedding_use_azure or LLM_SETTINGS.use_azure
@@ -389,6 +411,7 @@ class DeprecBackend(APIBackend):
         # transfer the config to the class if the config is not supposed to change during the runtime
         self.use_llama2 = LLM_SETTINGS.use_llama2
         self.use_gcr_endpoint = LLM_SETTINGS.use_gcr_endpoint
+        self.chat_use_azure_deepseek = LLM_SETTINGS.chat_use_azure_deepseek
         self.retry_wait_seconds = LLM_SETTINGS.retry_wait_seconds
 
     def _get_encoder(self) -> tiktoken.Encoding:
@@ -510,25 +533,60 @@ class DeprecBackend(APIBackend):
             return resp[0]
         return resp
 
-    def _create_chat_completion_auto_continue(self, messages: list[dict[str, Any]], *args, **kwargs) -> str:  # type: ignore[no-untyped-def]
+    def _create_chat_completion_auto_continue(
+        self,
+        messages: list[dict[str, Any]],
+        *args: Any,
+        json_mode: bool = False,
+        chat_cache_prefix: str = "",
+        seed: Optional[int] = None,
+        **kwargs: Any,
+    ) -> str:
         """
         Call the chat completion function and automatically continue the conversation if the finish_reason is length.
-        TODO: This function only continues once, maybe need to continue more than once in the future.
         """
-        response, finish_reason = self._create_chat_completion_inner_function(messages, *args, **kwargs)
+        if seed is None and LLM_SETTINGS.use_auto_chat_cache_seed_gen:
+            seed = LLM_CACHE_SEED_GEN.get_next_seed()
+        input_content_json = json.dumps(messages)
+        input_content_json = (
+            chat_cache_prefix + input_content_json + f"<seed={seed}/>"
+        )  # FIXME this is a hack to make sure the cache represents the round index
+        if self.use_chat_cache:
+            cache_result = self.cache.chat_get(input_content_json)
+            if cache_result is not None:
+                if LLM_SETTINGS.log_llm_chat_content:
+                    logger.info(f"{LogColors.CYAN}Response:{cache_result}{LogColors.END}", tag="llm_messages")
+                return cache_result
 
-        if finish_reason == "length":
-            new_message = deepcopy(messages)
-            new_message.append({"role": "assistant", "content": response})
-            new_message.append(
-                {
-                    "role": "user",
-                    "content": "continue the former output with no overlap",
-                },
-            )
-            new_response, finish_reason = self._create_chat_completion_inner_function(new_message, *args, **kwargs)
-            return response + new_response
-        return response
+        all_response = ""
+        new_messages = deepcopy(messages)
+        for _ in range(10):
+            if "json_mode" in kwargs:
+                del kwargs["json_mode"]
+            response, finish_reason = self._create_chat_completion_inner_function(
+                new_messages, json_mode=json_mode, *args, **kwargs
+            )  # type: ignore[misc]
+            all_response += response
+            if finish_reason is None or finish_reason != "length":
+                if self.chat_use_azure_deepseek:
+                    match = re.search(r"<think>(.*?)</think>(.*)", all_response, re.DOTALL)
+                    think_part, all_response = match.groups() if match else ("", all_response)
+                    if LLM_SETTINGS.log_llm_chat_content:
+                        logger.info(f"{LogColors.CYAN}Think:{think_part}{LogColors.END}", tag="llm_messages")
+                        logger.info(f"{LogColors.CYAN}Response:{all_response}{LogColors.END}", tag="llm_messages")
+
+                if json_mode:
+                    try:
+                        json.loads(all_response)
+                    except:
+                        match = re.search(r"```json(.*?)```", all_response, re.DOTALL)
+                        all_response = match.groups()[0] if match else all_response
+                        json.loads(all_response)
+                if self.dump_chat_cache:
+                    self.cache.chat_set(input_content_json, all_response)
+                return all_response
+            new_messages.append({"role": "assistant", "content": response})
+        return all_response
 
     def _try_create_chat_completion_or_embedding(  # type: ignore[no-untyped-def]
         self,
@@ -618,12 +676,10 @@ class DeprecBackend(APIBackend):
         messages: list[dict[str, Any]],
         temperature: float | None = None,
         max_tokens: int | None = None,
-        chat_cache_prefix: str = "",
         frequency_penalty: float | None = None,
         presence_penalty: float | None = None,
         json_mode: bool = False,
         add_json_in_prompt: bool = False,
-        seed: Optional[int] = None,
         *args,
         **kwargs,
     ) -> tuple[str, str | None]:
@@ -633,23 +689,11 @@ class DeprecBackend(APIBackend):
             To make retries useful, we need to enable a seed.
             This seed is different from `self.chat_seed` for GPT. It is for the local cache mechanism enabled by RD-Agent locally.
         """
-        if seed is None and LLM_SETTINGS.use_auto_chat_cache_seed_gen:
-            seed = LLM_CACHE_SEED_GEN.get_next_seed()
 
         # TODO: we can add this function back to avoid so much `self.cfg.log_llm_chat_content`
         if LLM_SETTINGS.log_llm_chat_content:
             logger.info(self._build_log_messages(messages), tag="llm_messages")
         # TODO: fail to use loguru adaptor due to stream response
-        input_content_json = json.dumps(messages)
-        input_content_json = (
-            chat_cache_prefix + input_content_json + f"<seed={seed}/>"
-        )  # FIXME this is a hack to make sure the cache represents the round index
-        if self.use_chat_cache:
-            cache_result = self.cache.chat_get(input_content_json)
-            if cache_result is not None:
-                if LLM_SETTINGS.log_llm_chat_content:
-                    logger.info(f"{LogColors.CYAN}Response:{cache_result}{LogColors.END}", tag="llm_messages")
-                return cache_result, None
 
         if temperature is None:
             temperature = LLM_SETTINGS.chat_temperature
@@ -700,6 +744,46 @@ class DeprecBackend(APIBackend):
             resp = json.loads(response.read().decode())["output"]
             if LLM_SETTINGS.log_llm_chat_content:
                 logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages")
+        elif self.chat_use_azure_deepseek:
+            azure_style_message: list[ChatRequestMessage] = []
+            for message in messages:
+                if message["role"] == "system":
+                    azure_style_message.append(SystemMessage(content=message["content"]))
+                elif message["role"] == "user":
+                    azure_style_message.append(UserMessage(content=message["content"]))
+                elif message["role"] == "assistant":
+                    azure_style_message.append(AssistantMessage(content=message["content"]))
+
+            response = self.client.complete(
+                messages=azure_style_message,
+                stream=self.chat_stream,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                frequency_penalty=frequency_penalty,
+                presence_penalty=presence_penalty,
+            )
+            if self.chat_stream:
+                resp = ""
+                # TODO: with logger.config(stream=self.chat_stream): and add a `stream_start` flag to add timestamp for first message.
+                if LLM_SETTINGS.log_llm_chat_content:
+                    logger.info(f"{LogColors.CYAN}Response:{LogColors.END}", tag="llm_messages")
+
+                for chunk in response:
+                    content = (
+                        chunk.choices[0].delta.content
+                        if len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None
+                        else ""
+                    )
+                    if LLM_SETTINGS.log_llm_chat_content:
+                        logger.info(LogColors.CYAN + content + LogColors.END, raw=True, tag="llm_messages")
+                    resp += content
+                    if len(chunk.choices) > 0 and chunk.choices[0].finish_reason is not None:
+                        finish_reason = chunk.choices[0].finish_reason
+            else:
+                resp = response.choices[0].message.content
+                finish_reason = response.choices[0].finish_reason
+                if LLM_SETTINGS.log_llm_chat_content:
+                    logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages")
         else:
             call_kwargs = dict(
                 model=model,
@@ -758,13 +842,11 @@ class DeprecBackend(APIBackend):
                         ),
                         tag="llm_messages",
                     )
-            if json_mode:
-                json.loads(resp)
-        if self.dump_chat_cache:
-            self.cache.chat_set(input_content_json, resp)
         return resp, finish_reason
 
     def _calculate_token_from_messages(self, messages: list[dict[str, Any]]) -> int:
+        if self.chat_use_azure_deepseek:
+            return 0
         if self.encoder is None:
             raise ValueError("Encoder is not initialized.")
         if self.use_llama2 or self.use_gcr_endpoint:

rdagent/oai/llm_conf.py

@@ -100,6 +100,10 @@ class LLMSettings(ExtendedBaseSettings):
     gcr_endpoint_do_sample: bool = False
     gcr_endpoint_max_token: int = 100
 
+    chat_use_azure_deepseek: bool = False
+    chat_azure_deepseek_endpoint: str = ""
+    chat_azure_deepseek_key: str = ""
+
     chat_model_map: str = "{}"
 
 

rdagent/scenarios/data_science/dev/runner/__init__.py

@@ -117,9 +117,9 @@ class DSCoSTEERRunner(CoSTEER):
             .read_text()
             .replace("<competition_id>", self.scen.competition)
         )
-        exp.experiment_workspace.inject_files(**{"mle_submission_format_test.py": mle_check_code})
+        exp.experiment_workspace.inject_files(**{"test/mle_submission_format_test.py": mle_check_code})
         exp.format_check_result = exp.experiment_workspace.execute(
-            env=mde, entry=f"python mle_submission_format_test.py"
+            env=mde, entry=f"python test/mle_submission_format_test.py"
         )
 
         return exp

rdagent/scenarios/data_science/dev/runner/eval.py

@@ -67,18 +67,15 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator):
                 .read_text()
                 .replace("<competition_id>", self.scen.competition)
             )
-            implementation.inject_files(**{"mle_submission_format_test.py": mle_check_code})
+            implementation.inject_files(**{"test/mle_submission_format_test.py": mle_check_code})
             stdout += f"\n MLEBench submission check:"
-            stdout += implementation.execute(env=mde, entry="python mle_submission_format_test.py")
+            stdout += implementation.execute(env=mde, entry="python test/mle_submission_format_test.py")
 
         # remove unused files
         implementation.execute(env=de, entry="coverage json -o coverage.json")
         if Path(implementation.workspace_path / "coverage.json").exists():
             with open(implementation.workspace_path / "coverage.json") as f:
-                used_files = set(json.load(f)["files"].keys()) | {
-                    "submission_format_test.py",
-                    "mle_submission_format_test.py",
-                }
+                used_files = set(json.load(f)["files"].keys())
                 logger.info("All used scripts: {}".format(used_files))
                 all_python_files = set(Path(implementation.workspace_path).rglob("*.py"))
                 unused_files = [

requirements.txt

@@ -49,3 +49,4 @@ nbformat
 # tool
 setuptools-scm
 seaborn
+azure.ai.inference