Merge branch 'sinaptik-ai:main' into template_improvements

Author: matteocacciola

docs/v3/semantic-layer/views.mdx

@@ -4,7 +4,8 @@ description: "Learn how to work with views in PandasAI"
 ---
 
 <Note title="Beta Notice">
-PandasAI 3.0 is currently in beta. This documentation reflects the latest features and functionality, which may evolve before the final release.
+  PandasAI 3.0 is currently in beta. This documentation reflects the latest
+  features and functionality, which may evolve before the final release.
 </Note>
 
 ## What are views?
@@ -68,21 +69,21 @@ view_relations = [
 # Select relevant columns for the sales analytics view
 view_columns = [
     # Order details
-    "orders.id",
-    "orders.order_date",
-    "orders.total_amount",
-    "orders.status",
-    
+    {"name": "orders.id", "type": "integer"},
+    {"name": "orders.order_date", "type": "date"},
+    {"name": "orders.total_amount", "type": "float"},
+    {"name": "orders.status", "type": "string"},
+
     # Product information
-    "products.name",
-    "products.category",
-    "products.unit_price",
-    "products.stock_level",
-    
+    {"name": "products.name", "type": "string"},
+    {"name": "products.category", "type": "string"},
+    {"name": "products.unit_price", "type": "float"},
+    {"name": "products.stock_level", "type": "integer"},
+
     # Customer information
-    "customers.segment",
-    "customers.country",
-    "customers.join_date"
+    {"name": "customers.segment", "type": "string"},
+    {"name": "customers.country", "type": "string"},
+    {"name": "customers.join_date", "type": "date"},
 ]
 
 # Create a comprehensive sales analytics view
@@ -125,16 +126,19 @@ relations:
 #### Constraints
 
 1. **Mutual Exclusivity**:
+
    - A schema cannot define both `table` and `view` simultaneously.
    - If `view` is `true`, then the schema represents a view.
 
 2. **Column Format**:
+
    - For views:
      - All columns must follow the format `[table].[column]`.
      - `from` and `to` fields in `relations` must follow the `[table].[column]` format.
      - Example: `loans.payment_amount`, `heart.condition`.
 
 3. **Relationships for Views**:
+
    - Each table referenced in `columns` must have at least one relationship defined in `relations`.
    - Relationships must specify `from` and `to` attributes in the `[table].[column]` format.
    - Relations define how different tables in your view are connected.

extensions/llms/litellm/pandasai_litellm/litellm.py

@@ -48,7 +48,7 @@ def type(self) -> str:
             str: The type of the model."""
         return f"litellm"
 
-    def call(self, instruction: BasePrompt, _: AgentState = None) -> str:
+    def call(self, instruction: BasePrompt, context: AgentState = None) -> str:
         """Generates a completion response based on the provided instruction.
 
         This method converts the given instruction into a user prompt string and
@@ -57,17 +57,18 @@ def call(self, instruction: BasePrompt, _: AgentState = None) -> str:
 
         Args:
             instruction (BasePrompt): The instruction to convert into a prompt.
-            _ (AgentState, optional): An optional state of the agent. Defaults to None.
+            context (AgentState, optional): An optional state of the agent. Defaults to None.
 
         Returns:
             str: The content of the model's response to the user prompt."""
 
-        user_prompt = instruction.to_string()
+        memory = context.memory if context else None
+        self.last_prompt = self.prepend_system_prompt(instruction.to_string(), memory)
 
         return (
             completion(
                 model=self.model,
-                messages=[{"content": user_prompt, "role": "user"}],
+                messages=[{"content": self.last_prompt, "role": "user"}],
                 **self.params,
             )
             .choices[0]

extensions/llms/openai/pandasai_openai/base.py

@@ -115,6 +115,7 @@ def completion(self, prompt: str, memory: Memory) -> str:
 
         Args:
             prompt (str): A string representation of the prompt.
+            memory (Memory): Memory object containing conversation history.
 
         Returns:
             str: LLM response.
@@ -139,6 +140,7 @@ def chat_completion(self, value: str, memory: Memory) -> str:
 
         Args:
             value (str): Prompt
+            memory (Memory): Memory object containing conversation history.
 
         Returns:
             str: LLM response.

pandasai/agent/base.py

@@ -270,7 +270,7 @@ def _process_query(self, query: str, output_type: Optional[str] = None):
             self._state.assign_prompt_id()
 
             # Generate code
-            code = self.generate_code_with_retries(query)
+            code = self.generate_code_with_retries(str(query))
 
             # Execute code with retries
             result = self.execute_with_retries(code)
@@ -296,7 +296,7 @@ def _regenerate_code_after_error(self, code: str, error: Exception) -> str:
 
         return self._code_generator.generate_code(prompt)
 
-    def _handle_exception(self, code: str) -> str:
+    def _handle_exception(self, code: str) -> ErrorResponse:
         """Handle exceptions and return an error message."""
         error_message = traceback.format_exc()
         self._state.logger.log(f"Processing failed with error: {error_message}")

pandasai/core/code_generation/base.py

@@ -18,7 +18,6 @@ def generate_code(self, prompt: BasePrompt) -> str:
         Generates code using a given LLM and performs validation and cleaning steps.
 
         Args:
-            context (PipelineContext): The pipeline context containing dataframes and logger.
             prompt (BasePrompt): The prompt to guide code generation.
 
         Returns:
@@ -32,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,
     )
 

pandasai/llm/base.py

@@ -119,11 +119,11 @@ def _extract_code(self, response: str, separator: str = "```") -> str:
 
         return code
 
-    def prepend_system_prompt(self, prompt: BasePrompt, memory: Memory):
+    def prepend_system_prompt(self, prompt: str, memory: Memory) -> str | Any:
         """
         Append system prompt to the chat prompt, useful when model doesn't have messages for chat history
         Args:
-            prompt (BasePrompt): prompt for chat method
+            prompt (str): prompt for chat method
             memory (Memory): user conversation history
         """
         return self.get_system_prompt(memory) + prompt if memory else prompt
@@ -164,6 +164,7 @@ def generate_code(self, instruction: BasePrompt, context: AgentState) -> str:
 
         Args:
             instruction (BasePrompt): Prompt with instruction for LLM.
+            context (AgentState): Context to pass.
 
         Returns:
             str: A string of Python code.

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"