fixing LiteLLM call for the follow up (#1799)

Author: matteocacciola

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:

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.