fix: last_code_generated retrieval logic in multi-turn conversations (#1784)

* Fix last_code_generated retrieval logic in multi-turn conversations

* refine comments

Author: kingdomad

pandasai/core/code_generation/base.py

@@ -31,10 +31,16 @@ def generate_code(self, prompt: BasePrompt) -> str:
 
             # Generate the code
             code = self._context.config.llm.generate_code(prompt, self._context)
+            # Store the original generated code (for logging purposes)
             self._context.last_code_generated = code
             self._context.logger.log(f"Code Generated:\n{code}")
 
-            return self.validate_and_clean_code(code)
+            # Validate and clean the code
+            cleaned_code = self.validate_and_clean_code(code)
+            # Update with the final cleaned code (for subsequent processing and multi-turn conversations)
+            self._context.last_code_generated = cleaned_code
+
+            return cleaned_code
 
         except Exception as e:
             error_message = f"An error occurred during code generation: {e}"

pandasai/core/prompts/__init__.py

@@ -19,7 +19,7 @@
 def get_chat_prompt_for_sql(context: AgentState) -> BasePrompt:
     return GeneratePythonCodeWithSQLPrompt(
         context=context,
-        last_code_generated=context.get("last_code_generated"),
+        last_code_generated=context.last_code_generated,
         output_type=context.output_type,
     )
 

tests/unit_tests/agent/test_agent.py

@@ -542,3 +542,25 @@ def test_handle_exception(self, agent):
         # Verify the error was logged
         mock_logger.log.assert_called_once()
         assert "Processing failed with error" in mock_logger.log.call_args[0][0]
+
+    def test_last_code_generated_retrieval(self, agent: Agent):
+        """Test that last_code_generated is correctly retrieved in get_chat_prompt_for_sql."""
+        # Set last_code_generated
+        test_code = "print('Test code')"
+        agent._state.last_code_generated = test_code
+
+        # 使用 get_chat_prompt_for_sql 获取提示
+        from pandasai.core.prompts import get_chat_prompt_for_sql
+
+        prompt = get_chat_prompt_for_sql(agent._state)
+
+        # 验证提示中使用了正确的 last_code_generated
+        assert prompt.props["last_code_generated"] == test_code
+
+        # 验证不是从 intermediate_values 中获取的
+        agent._state.add("last_code_generated", "Wrong code")
+        prompt = get_chat_prompt_for_sql(agent._state)
+
+        # 应该仍然使用 last_code_generated 属性，而不是 intermediate_values 中的值
+        assert prompt.props["last_code_generated"] == test_code
+        assert prompt.props["last_code_generated"] != "Wrong code"