Fix TypeError in fin_quant coding step (#1224)

rdagent/components/coder/CoSTEER/config.py

@@ -1,13 +1,14 @@
 from typing import Union
-
+from pydantic_settings import BaseSettings
+from pydantic import ConfigDict  
 from rdagent.core.conf import ExtendedBaseSettings
 
-
 class CoSTEERSettings(ExtendedBaseSettings):
     """CoSTEER settings, this setting is supposed not to be used directly!!!"""
 
-    class Config:
-        env_prefix = "CoSTEER_"
+    model_config = ConfigDict(
+           env_prefix="CoSTEER_"
+    )
 
     coder_use_cache: bool = False
     """Indicates whether to use cache for the coder"""
@@ -38,5 +39,5 @@ class Config:
 
     max_seconds_multiplier: int = 10**6
 
-
 CoSTEER_SETTINGS = CoSTEERSettings()
+

rdagent/components/coder/factor_coder/evolving_strategy.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import logging
 import re
 from typing import Dict
 
@@ -19,6 +20,7 @@
 from rdagent.oai.llm_utils import APIBackend
 from rdagent.utils.agent.tpl import T
 
+logger = logging.getLogger(__name__)
 
 class FactorMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
     def __init__(self, *args, **kwargs) -> None:
@@ -98,7 +100,6 @@ def implement_one_task(
         )
         queried_similar_successful_knowledge_to_render = queried_similar_successful_knowledge
         queried_similar_error_knowledge_to_render = queried_similar_error_knowledge
-        # 动态地防止prompt超长
         for _ in range(10):  # max attempt to reduce the length of user_prompt
             # 总结error（可选）
             if (
@@ -164,11 +165,30 @@ def implement_one_task(
             return ""  # return empty code if failed to get code after 10 attempts
 
     def assign_code_list_to_evo(self, code_list, evo):
+        for index in range(len(code_list)):
+            item = code_list[index]
+            if isinstance(item, str):
+                continue  # Already a string, skip
+            elif isinstance(item, dict):
+                code_str = item.get("factor.py") or item.get("code") or item.get("output") or item.get("content", "")
+                if isinstance(code_str, str):
+                    code_list[index] = code_str.strip()
+                    logger.info(f"Extracted code from dict at index {index} using key: {list(item.keys())}")
+                else:
+                    code_list[index] = "" 
+                    logger.warning(f"No valid str code in dict at index {index}: {item}")
+            elif item is None:
+                continue  
+            else:
+                logger.error(f"Invalid type at index {index}: {type(item)}")
+                code_list[index] = ""  
+
         for index in range(len(evo.sub_tasks)):
-            if code_list[index] is None:
-                continue
+            if code_list[index] is None or code_list[index] == "":
+                continue 
             if evo.sub_workspace_list[index] is None:
                 evo.sub_workspace_list[index] = FactorFBWorkspace(target_task=evo.sub_tasks[index])
+
             # Since the `implement_one_task` method is not standardized and the `code_list` has both `str` and `dict` data types,
             # we ended up getting an `TypeError` here, so we chose to fix the problem temporarily with this dirty method.
             if isinstance(code_list[index], dict):

rdagent/oai/backend/litellm.py

@@ -12,7 +12,8 @@
     supports_response_schema,
     token_counter,
 )
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
+from pydantic_settings import BaseSettings
 
 from rdagent.log import LogColors
 from rdagent.log import rdagent_logger as logger
@@ -31,14 +32,12 @@ def _reduce_no_init(exc: Exception) -> tuple:
 copyreg.pickle(BadRequestError, _reduce_no_init)
 
 
-class LiteLLMSettings(LLMSettings):
-
-    class Config:
-        env_prefix = "LITELLM_"
-        """Use `LITELLM_` as prefix for environment variables"""
-
-    # Placeholder for LiteLLM specific settings, so far it's empty
-
+class LiteLLMSettings(LLMSettings, BaseSettings):
+    """LiteLLM settings configuration."""
+    model_config = ConfigDict(
+        env_prefix = "LITELLM_",
+        extra = "ignore",
+    )
 
 LITELLM_SETTINGS = LiteLLMSettings()
 ACC_COST = 0.0

test/oai/test_evolving_strategy_fix.py

@@ -0,0 +1,92 @@
+import unittest
+from unittest.mock import Mock, patch
+import sys
+import os
+import pytest
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+from rdagent.components.coder.factor_coder.evolving_strategy import FactorMultiProcessEvolvingStrategy
+from rdagent.components.coder.factor_coder.factor import FactorFBWorkspace, FactorTask
+from rdagent.core.experiment import FBWorkspace
+
+@pytest.fixture(autouse=True)
+def mock_strategy_init():
+    with patch('rdagent.components.coder.factor_coder.evolving_strategy.MultiProcessEvolvingStrategy.__init__', return_value=None):
+        yield
+
+@pytest.fixture(autouse=True)
+def mock_workspace():
+    with patch('rdagent.components.coder.factor_coder.factor.FactorFBWorkspace') as mock_workspace_class:
+        mock_workspace_instance = Mock(spec=FBWorkspace)
+        mock_workspace_class.return_value = mock_workspace_instance
+        yield mock_workspace_class
+
+class TestAssignCodeListToEvoFix(unittest.TestCase):
+    def setUp(self):
+        self.mock_task = Mock(spec=FactorTask)
+        self.mock_task.get_task_information.return_value = "Mock task info"
+        self.strategy = FactorMultiProcessEvolvingStrategy(Mock(), Mock())
+        self.strategy.scen = Mock()
+        self.strategy.scen.get_scenario_all_desc.return_value = "Mock scenario"
+
+    @patch('rdagent.components.coder.factor_coder.evolving_strategy.FactorFBWorkspace.inject_files')
+    def test_happy_path_strings_only(self, mock_inject_files):
+        code_list = ['def factor1(): return 1 + 2', None, 'print("Another factor")']
+        evo = Mock(sub_tasks=[self.mock_task] * 3, sub_workspace_list=[None] * 3)
+        result = self.strategy.assign_code_list_to_evo(code_list, evo)
+        self.assertIs(evo, result)
+        self.assertEqual(mock_inject_files.call_count, 2)  # Only non-None items
+
+    @patch('rdagent.components.coder.factor_coder.evolving_strategy.FactorFBWorkspace.inject_files')
+    def test_dict_with_factor_py_key(self, mock_inject_files):
+        code_list = [{'factor.py': 'def alpha(): return close / open - 1'}, {'factor.py': ''}, 'fallback str code', {'no_factor.py': 'invalid'}]
+        evo = Mock(sub_tasks=[self.mock_task] * 4, sub_workspace_list=[None] * 4)
+        result = self.strategy.assign_code_list_to_evo(code_list, evo)
+        self.assertEqual(code_list[0], 'def alpha(): return close / open - 1')
+        self.assertEqual(mock_inject_files.call_count, 2)  # Only indices 0 and 2 injected
+
+    @patch('rdagent.components.coder.factor_coder.evolving_strategy.FactorFBWorkspace.inject_files')
+    def test_dict_without_factor_py_key(self, mock_inject_files):
+        code_list = [{'code': 'def momentum(): return vwap / ref(1)'}, {'output': 'model code here'}, None]
+        evo = Mock(sub_tasks=[self.mock_task] * 3, sub_workspace_list=[None] * 3)
+        result = self.strategy.assign_code_list_to_evo(code_list, evo)
+        self.assertEqual(code_list[0], 'def momentum(): return vwap / ref(1)')
+        self.assertEqual(mock_inject_files.call_count, 2)  # Indices 0 and 1 injected
+
+    @patch('rdagent.components.coder.factor_coder.evolving_strategy.FactorFBWorkspace.inject_files')
+    def test_empty_or_invalid_list(self, mock_inject_files):
+        code_list = []
+        evo = Mock(sub_tasks=[], sub_workspace_list=[])
+        result = self.strategy.assign_code_list_to_evo(code_list, evo)
+        self.assertIs(evo, result)
+        mock_inject_files.assert_not_called()
+
+        code_list = [None, None]
+        evo = Mock(sub_tasks=[self.mock_task] * 2, sub_workspace_list=[None] * 2)
+        result = self.strategy.assign_code_list_to_evo(code_list, evo)
+        mock_inject_files.assert_not_called()
+
+    def test_unfixed_behavior_simulation(self):
+        def broken_assign(self, code_list, evo):
+            for index in range(len(evo.sub_tasks)):
+                if code_list[index] is None:
+                    continue
+                if evo.sub_workspace_list[index] is None:
+                    evo.sub_workspace_list[index] = FactorFBWorkspace(target_task=evo.sub_tasks[index])
+                evo.sub_workspace_list[index].inject_files(**{"factor.py": code_list[index]})
+            return evo
+
+        original_method = self.strategy.assign_code_list_to_evo
+        self.strategy.assign_code_list_to_evo = broken_assign.__get__(self.strategy, FactorMultiProcessEvolvingStrategy)
+
+        code_list = [{'code': 'bad dict'}]
+        evo = Mock(sub_tasks=[self.mock_task], sub_workspace_list=[None])
+
+        with self.assertRaises(TypeError):
+            self.strategy.assign_code_list_to_evo(code_list, evo)
+
+        self.strategy.assign_code_list_to_evo = original_method
+
+if __name__ == '__main__':
+    unittest.main(verbosity=2)