Merge pull request #1295 from WojtAcht/main

fix: skill name is not defined

pandasai/pipelines/chat/code_cleaning.py

@@ -407,7 +407,7 @@ def _get_originals(self, dfs):
         return original_dfs
 
     def _extract_fix_dataframe_redeclarations(
-        self, node: ast.AST, code_lines: list[str]
+        self, node: ast.AST, code_lines: list[str], context: CodeExecutionContext
     ) -> ast.AST:
         if isinstance(node, ast.Assign):
             target_names, is_slice, target = self._get_target_names(node.targets)
@@ -417,6 +417,12 @@ def _extract_fix_dataframe_redeclarations(
                 code = "\n".join(code_lines)
                 env = get_environment(self._additional_dependencies)
                 env["dfs"] = copy.deepcopy(self._get_originals(self._dfs))
+                if context.skills_manager.used_skills:
+                    for skill_func_name in context.skills_manager.used_skills:
+                        skill = context.skills_manager.get_skill_by_func_name(
+                            skill_func_name
+                        )
+                        env[skill_func_name] = skill
                 exec(code, env)
 
                 df_generated = (
@@ -512,7 +518,9 @@ def _clean_code(self, code: str, context: CodeExecutionContext) -> str:
             clean_code_lines.append(astor.to_source(node))
 
             new_body.append(
-                self._extract_fix_dataframe_redeclarations(node, clean_code_lines)
+                self._extract_fix_dataframe_redeclarations(
+                    node, clean_code_lines, context
+                )
                 or node
             )
 

tests/unit_tests/pipelines/smart_datalake/test_code_cleaning.py

@@ -604,7 +604,7 @@ def test_fix_dataframe_redeclarations(self, mock_head, context: PipelineContext)
         clean_code = ["df1 = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})"]
 
         output = code_cleaning._extract_fix_dataframe_redeclarations(
-            tree.body[0], clean_code
+            tree.body[0], clean_code, context
         )
 
         assert isinstance(output, ast.Assign)
@@ -634,7 +634,9 @@ def test_fix_dataframe_multiline_redeclarations(
         ]
 
         outputs = [
-            code_cleaning._extract_fix_dataframe_redeclarations(node, clean_codes)
+            code_cleaning._extract_fix_dataframe_redeclarations(
+                node, clean_codes, context
+            )
             for node in tree.body
         ]
 
@@ -661,7 +663,7 @@ def test_fix_dataframe_no_redeclarations(self, mock_head, context: PipelineConte
         code_list = ["df1 = dfs[0]"]
 
         output = code_cleaning._extract_fix_dataframe_redeclarations(
-            tree.body[0], code_list
+            tree.body[0], code_list, context
         )
 
         assert output is None
@@ -686,7 +688,7 @@ def test_fix_dataframe_redeclarations_with_subscript(
         code_list = ["dfs[0] = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})"]
 
         output = code_cleaning._extract_fix_dataframe_redeclarations(
-            tree.body[0], code_list
+            tree.body[0], code_list, context
         )
 
         assert isinstance(output, ast.Assign)
@@ -720,7 +722,7 @@ def test_fix_dataframe_redeclarations_with_subscript_and_data_variable(
         ]
 
         output = code_cleaning._extract_fix_dataframe_redeclarations(
-            tree.body[1], code_list
+            tree.body[1], code_list, context
         )
 
         assert isinstance(output, ast.Assign)
@@ -754,7 +756,7 @@ def test_fix_dataframe_redeclarations_and_data_variable(
         ]
 
         output = code_cleaning._extract_fix_dataframe_redeclarations(
-            tree.body[1], code_list
+            tree.body[1], code_list, context
         )
 
         assert isinstance(output, ast.Assign)