remove other commits and focus on add more package info for persistent problems

Author: Hoder-zyf

rdagent/scenarios/data_science/proposal/exp_gen/package_info.py

@@ -6,6 +6,72 @@ def get_installed_packages():
     return {dist.metadata["Name"].lower(): dist.version for dist in distributions()}
 
 
+# Kaggle competition packages - based on usage frequency
+PYTHON_BASE_PACKAGES = ["catboost", "lightgbm", "numpy", "optuna", "pandas", "scikit-learn", "scipy", "shap", "xgboost"]
+
+PYTHON_ADVANCED_PACKAGES = [
+    "accelerate",
+    "albumentations",
+    "category_encoders",
+    "cudf-cu12",
+    "cuml-cu12",
+    "datasets",
+    "featuretools",
+    "imbalanced-learn",
+    "opencv-python",
+    "pillow",
+    "polars",
+    "sentence-transformers",
+    "spacy",
+    "tensorflow",
+    "timm",
+    "tokenizers",
+    "torch",
+    "torchvision",
+    "transformers",
+]
+
+PYTHON_AUTO_ML_PACKAGES = ["autogluon"]
+
+
+def get_available_packages_prompt():
+    """Generate prompt template for dynamically detected available packages"""
+    installed_packages = get_installed_packages()
+
+    # Check which packages are actually installed
+    base_available = [pkg for pkg in PYTHON_BASE_PACKAGES if pkg.lower() in installed_packages]
+    advanced_available = [pkg for pkg in PYTHON_ADVANCED_PACKAGES if pkg.lower() in installed_packages]
+    automl_available = [pkg for pkg in PYTHON_AUTO_ML_PACKAGES if pkg.lower() in installed_packages]
+
+    # Build prompt
+    prompt_parts = ["Available packages in environment:\n"]
+
+    if base_available:
+        prompt_parts.append("【Basic Libraries】(core tools for most competitions):")
+        prompt_parts.append(f"- {', '.join(base_available)}")
+        prompt_parts.append("")
+
+    if advanced_available:
+        prompt_parts.append("【Advanced Tools】(specialized for specific domains):")
+        prompt_parts.append(f"- {', '.join(advanced_available)}")
+        prompt_parts.append("")
+
+    if automl_available:
+        prompt_parts.append("【AutoML Tools】(automated machine learning):")
+        prompt_parts.append(f"- {', '.join(automl_available)}")
+        prompt_parts.append("")
+
+    prompt_parts.append("Choose appropriate tool combinations based on the competition type.")
+
+    return "\n".join(prompt_parts).strip()
+
+
+def get_all_available_packages():
+    """Get flattened list of all packages"""
+    all_packages = PYTHON_BASE_PACKAGES + PYTHON_ADVANCED_PACKAGES + PYTHON_AUTO_ML_PACKAGES
+    return sorted(set(all_packages))
+
+
 def print_filtered_packages(installed_packages, filtered_packages):
     to_print = []
     for package_name in filtered_packages:
@@ -26,24 +92,8 @@ def get_python_packages():
     # Example: `python package_info.py pandas torch scikit-learn`
     # If no extra arguments are provided we fall back to the original default list
     # to keep full backward-compatibility.
-    packages_list = [  # default packages
-        "transformers",
-        "accelerate",
-        "torch",
-        "tensorflow",
-        "pandas",
-        "numpy",
-        "scikit-learn",
-        "scipy",
-        "xgboost",
-        "sklearn",
-        "lightgbm",
-        "vtk",
-        "opencv-python",
-        "keras",
-        "matplotlib",
-        "pydicom",
-    ]
+    # Use our Kaggle-optimized package list as default
+    packages_list = get_all_available_packages()
     if len(sys.argv) > 1:
         packages_list = list(set(packages_list) | set(sys.argv[1:]))
 
@@ -61,4 +111,8 @@ def get_python_packages():
 
 
 if __name__ == "__main__":
-    get_python_packages()
+    # Check for special argument to get prompt instead of package list
+    if len(sys.argv) > 1 and sys.argv[1] == "--prompt":
+        print(get_available_packages_prompt())
+    else:
+        get_python_packages()

rdagent/scenarios/data_science/proposal/exp_gen/proposal.py

@@ -23,6 +23,9 @@
     DSDraftExpGen,  # TODO: DSDraftExpGen should be moved to router in the further
 )
 from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSIdea
+from rdagent.scenarios.data_science.proposal.exp_gen.package_info import (
+    get_available_packages_prompt,
+)
 from rdagent.scenarios.data_science.proposal.exp_gen.planner import (
     DSExperimentPlan,
     RD_Agent_TIMER_wrapper,
@@ -601,6 +604,12 @@ def hypothesis_gen(
         for i, (problem_name, problem_dict) in enumerate(problems.items()):
             problem_formatted_str += f"## {i+1}. {problem_name}\n"
             problem_formatted_str += f"{problem_dict['problem']}\n"
+
+            # Add package information for persistent problems
+            if problem_dict.get("label") == "PERSISTENT_PROBLEM":
+                packages_prompt = get_available_packages_prompt()
+                problem_formatted_str += f"\n{packages_prompt}\n"
+
             if "idea" in problem_dict:
                 idea_formatted_str = DSIdea(problem_dict["idea"]).to_formatted_str()
                 problem_formatted_str += f"Sampled Idea by user: \n{idea_formatted_str}\n"

rdagent/scenarios/data_science/proposal/exp_gen/utils.py

@@ -103,3 +103,16 @@ def get_packages(pkgs: list[str] | None = None) -> str:
     pkg_args = " ".join(pkgs) if pkgs else ""
     stdout = implementation.execute(env=env, entry=f"python {fname} {pkg_args}")
     return stdout
+
+
+def get_packages_prompt() -> str:
+    """Return available packages prompt information."""
+    # Reuse package prompt cached during Draft stage when available.
+
+    env = get_ds_env()
+    implementation = FBWorkspace()
+    fname = "package_info.py"
+    implementation.inject_files(**{fname: (Path(__file__).absolute().resolve().parent / "package_info.py").read_text()})
+
+    stdout = implementation.execute(env=env, entry=f"python {fname} --prompt")
+    return stdout