fix: last_code_generated retrieval logic in multi-turn conversations

[kingdomad]

Fix last_code_generated retrieval logic in multi-turn conversations

Issue Description

In multi-turn conversations, previously generated code was not being correctly passed to new code generation processes, causing a lack of continuity in code generation. This issue affected the user experience when conducting multi-step analyses that required building upon previous code.

Root Cause

The issue stemmed from a mismatch in how last_code_generated was stored and retrieved:

1. The generated code was stored in the AgentState.last_code_generated attribute:

   self._context.last_code_generated = code

2. However, when creating a new prompt, the code attempted to retrieve it from the intermediate_values dictionary:

   # In get_chat_prompt_for_sql function
   last_code_generated=context.get("last_code_generated")

3. The AgentState.get method only fetches values from the intermediate_values dictionary:

   def get(self, key: str, default: Any = "") -> Any:
       """Fetches a value from intermediate values or returns a default."""
       return self.intermediate_values.get(key, default)

This disconnect meant that in multi-turn conversations, each new code generation would start from scratch rather than building upon previous code.

Changes Made

This PR fixes the issue by:

1. Modifying get_chat_prompt_for_sql to directly use the last_code_generated attribute:

   # Changed from:
   last_code_generated=context.get("last_code_generated")
   
   # To:
   last_code_generated=context.last_code_generated

2. Ensuring the cleaned code is also stored in last_code_generated:

   # Changed from:
   return self.validate_and_clean_code(code)
   
   # To:
   cleaned_code = self.validate_and_clean_code(code)
   self._context.last_code_generated = cleaned_code
   return cleaned_code

3. Adding a test case to verify the fix:

   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 prompt using get_chat_prompt_for_sql
       from pandasai.core.prompts import get_chat_prompt_for_sql
       prompt = get_chat_prompt_for_sql(agent._state)
       
       # Verify prompt uses correct last_code_generated
       assert prompt.props["last_code_generated"] == test_code
       
       # Verify it's not retrieved from intermediate_values
       agent._state.add("last_code_generated", "Wrong code")
       prompt = get_chat_prompt_for_sql(agent._state)
       
       # Should still use last_code_generated attribute, not intermediate_values
       assert prompt.props["last_code_generated"] == test_code
       assert prompt.props["last_code_generated"] != "Wrong code"

Benefits

This fix ensures:

1. Code continuity in multi-turn conversations
2. Better user experience when building complex analyses step by step
3. More efficient code generation by leveraging previously generated code

Testing

A new test case test_last_code_generated_retrieval was added to verify the fix, which passes successfully.

---

Important

Fixes last_code_generated retrieval logic in multi-turn conversations to ensure code continuity.

• Behavior:
  • Fixes last_code_generated retrieval in get_chat_prompt_for_sql to use context.last_code_generated instead of context.get().
  • Ensures cleaned code is stored in last_code_generated in generate_code() in base.py.
• Testing:
  • Adds test_last_code_generated_retrieval in test_agent.py to verify correct retrieval of last_code_generated.
• Benefits:
  • Ensures code continuity in multi-turn conversations.
  • Improves user experience for multi-step analyses.

^{This description was created by for 3180cc6. You can customize this summary. It will automatically update as commits are pushed.}

[ellipsis-dev]

Caution

Changes requested ❌

Reviewed everything up to 3180cc6 in 1 minute and 45 seconds. Click for details.

• Reviewed 61 lines of code in 3 files
• Skipped 0 files when reviewing.
• Skipped posting 2 draft comments. View those below.
• Modify your settings and rules to customize what types of comments Ellipsis leaves. And don't forget to react with 👍 or 👎 to teach Ellipsis.

1. pandasai/core/code_generation/base.py:38

• Draft comment:
  Good update using the cleaned code to update last_code_generated. Consider whether the initial assignment of last_code_generated (before cleaning) is necessary, as it gets immediately overwritten by the cleaned version. If the raw code isn't needed later, removing the redundant assignment could reduce confusion.
• Reason this comment was not posted:
  Decided after close inspection that this draft comment was likely wrong and/or not actionable: usefulness confidence = 20% vs. threshold = 50% While the comment identifies a potential redundancy, we don't have enough context to know if the raw code needs to be temporarily stored for logging or debugging purposes. The initial assignment happens before logging the raw code, so it may serve a purpose. Without understanding the full usage of last_code_generated throughout the codebase, removing it could be risky. The comment makes assumptions about the usage of last_code_generated without full context. The initial assignment might be intentional for logging or debugging purposes. Given that the raw code is logged immediately after the first assignment, the initial assignment may serve a purpose in the logging flow. The comment should be deleted as it makes assumptions without sufficient context and the current implementation may be intentional for logging purposes.

2. pandasai/core/prompts/__init__.py:20

• Draft comment:
  Retrieving last_code_generated directly from the context attribute (instead of using context.get(...)) is a clear improvement for multi-turn conversation continuity.
• Reason this comment was not posted:
  Comment did not seem useful. Confidence is useful = 0% <= threshold 50% This comment is purely informative and does not provide any actionable feedback or suggestions for improvement. It simply states an opinion about the code change without offering any specific guidance or questions for the PR author.

Workflow ID: wflow_shanL1yxu5U6Xrmq

^{You can customize by changing your verbosity settings, reacting with 👍 or 👎, replying to comments, or adding code review rules.}

[ellipsis-dev]

The new test for last_code_generated retrieval is thorough. Consider using a consistent language in comments (currently mixing Chinese and English) to improve clarity and maintainability.

Suggested change

	# 设置 last_code_generated
	# Set last_code_generated

^{This comment was generated because it violated a code review rule: mrule_kO2c9rTcwkbqZtg4.}

[kingdomad]

Hi @ArslanSaleem , could you please take a look at this PR when you have time? Thanks!

[kingdomad]

Hi there, I've looked into the logs for the failed CI check. The failure seems to be caused by an OSError while trying to download the GPTCache/paraphrase-albert-small-v2 model from Hugging Face, likely due to rate limiting (HTTP 429).

This seems to be a transient network issue within the CI environment and is unrelated to the changes in this PR. Since I don't have permissions to re-run the job, could a maintainer please help restart the check? Thanks!