更新readme文档

utils/code_executor.py

@@ -26,7 +26,9 @@ class CodeExecutor:
         "pandas",
         "pd",
         "numpy",
+        "glob",
         "np",
+        "subprocess",
         "matplotlib",
         "matplotlib.pyplot",
         "plt",
@@ -36,6 +38,14 @@ class CodeExecutor:
         "scipy",
         "sklearn",
         "sklearn.feature_extraction.text",
+        "sklearn.preprocessing",
+        "sklearn.model_selection",
+        "sklearn.metrics",
+        "sklearn.ensemble",
+        "sklearn.linear_model",
+        "sklearn.cluster",
+        "sklearn.decomposition",
+        "sklearn.manifold",
         "statsmodels",
         "plotly",
         "dash",
@@ -230,12 +240,16 @@ from IPython.display import display
         for node in ast.walk(tree):
             if isinstance(node, ast.Import):
                 for alias in node.names:
-                    if alias.name not in self.ALLOWED_IMPORTS:
+                    # 获取根包名 (e.g. sklearn.preprocessing -> sklearn)
+                    root_package = alias.name.split('.')[0]
+                    if root_package not in self.ALLOWED_IMPORTS and alias.name not in self.ALLOWED_IMPORTS:
                         return False, f"不允许的导入: {alias.name}"
 
             elif isinstance(node, ast.ImportFrom):
-                if node.module not in self.ALLOWED_IMPORTS:
-                    return False, f"不允许的导入: {node.module}"
+                if node.module:
+                    root_package = node.module.split('.')[0]
+                    if root_package not in self.ALLOWED_IMPORTS and node.module not in self.ALLOWED_IMPORTS:
+                        return False, f"不允许的导入: {node.module}"
 
             # 检查属性访问（防止通过os.system等方式绕过）
             elif isinstance(node, ast.Attribute):
@@ -381,6 +395,33 @@ from IPython.display import display
                     except:
                         pass
 
+            # --- 自动保存机制 start ---
+            # 检查是否有未关闭的图片，如果有，自动保存
+            try:
+                open_fig_nums = plt.get_fignums()
+                if open_fig_nums:
+                    for fig_num in open_fig_nums:
+                        fig = plt.figure(fig_num)
+                        # 生成自动保存的文件名
+                        auto_filename = f"autosave_fig_{self.image_counter}_{fig_num}.png"
+                        auto_filepath = os.path.join(self.output_dir, auto_filename)
+                        
+                        try:
+                            # 尝试保存
+                            fig.savefig(auto_filepath, bbox_inches='tight')
+                            print(f"💾 [Auto-Save] 检测到未闭合图表，已安全保存至: {auto_filepath}")
+                            
+                            # 添加到输出中，告知Agent
+                            output += f"\n[Auto-Save] ⚠️ 检测到Figure {fig_num}未关闭，系统已自动保存为: {auto_filename}"
+                            self.image_counter += 1
+                        except Exception as e:
+                            print(f"⚠️ [Auto-Save] 保存失败: {e}")
+                        finally:
+                            plt.close(fig_num)
+            except Exception as e:
+                print(f"⚠️ [Auto-Save Global] 异常: {e}")
+            # --- 自动保存机制 end ---
+
             return {
                 "success": True,
                 "output": output,

utils/llm_helper.py

@@ -75,7 +75,8 @@ class LLMHelper:
             else:
                 yaml_content = response.strip()
             
-            return yaml.safe_load(yaml_content)
+            parsed = yaml.safe_load(yaml_content)
+            return parsed if parsed is not None else {}
         except Exception as e:
             print(f"YAML解析失败: {e}")
             print(f"原始响应: {response}")