From 44d1ad80c57d0196002619f23372d79718bb63bd Mon Sep 17 00:00:00 2001 From: you-n-g Date: Wed, 5 Mar 2025 22:46:43 +0800 Subject: [PATCH] fix: target json type & round (#662) * refactor: Rename direct_exp_gen to json_target_type in DSExpGen class * fix type * fix: Adjust loop iterations and update json_target_type for nested dicts --- rdagent/oai/backend/base.py | 2 +- rdagent/scenarios/data_science/proposal/exp_gen.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/rdagent/oai/backend/base.py b/rdagent/oai/backend/base.py index 5a0b83ce..c9a5789a 100644 --- a/rdagent/oai/backend/base.py +++ b/rdagent/oai/backend/base.py @@ -378,7 +378,7 @@ class APIBackend(ABC): all_response = "" new_messages = deepcopy(messages) - for _ in range(3): + for _ in range(6): # for some long code, 3 times may not enough for reasoning models if "json_mode" in kwargs: del kwargs["json_mode"] response, finish_reason = self._create_chat_completion_add_json_in_prompt( diff --git a/rdagent/scenarios/data_science/proposal/exp_gen.py b/rdagent/scenarios/data_science/proposal/exp_gen.py index 8c3dcdb1..96de6361 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen.py @@ -420,7 +420,10 @@ class DSExpGen(ExpGen): user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True, - json_target_type=dict[str, dict[str, str] | str], # workflow_update may be a string + # NOTE: corner cases. + # workflow_update may be a string + # model could have 2 level nested dict. + json_target_type=dict[str, dict[str, str | dict] | str], ) ) assert "hypothesis_proposal" in resp_dict, "Hypothesis proposal not provided."