DB Chat for pro users only

python_code_executor_service.py

@@ -32,7 +32,6 @@ class ChartSpecification(BaseModel):
     """Details about requested charts"""
     image_description: str
     code: Optional[str] = None
-    media_pointer: Optional[str] = None  # Added media pointer field
 
 
 class AnalysisOperation(BaseModel):
@@ -188,62 +187,62 @@ class PythonExecutor:
     def _format_result(self, result: Any) -> str:
         """Format the result for display"""
         if isinstance(result, (pd.DataFrame, pd.Series)):
-            return result.to_string()
+            # Convert DataFrame to a string, then to a list of dicts (handles NumPy types)
+            json_str = result.to_json(orient='records', date_format='iso')
+            json.dumps(json.loads(json_str), indent=2)  # Re-parse for pretty formatting
         elif isinstance(result, (dict, list)):
             return json.dumps(result, indent=2)
         return str(result)
     
     async def process_response(self, response: CsvChatResult, chat_id: str) -> str:
-        """Process the response with proper variable handling"""
-        output_parts = [response.casual_response]
-        
-        # Process analysis operations first
-        for operation in response.analysis_operations:
-            execution_result = self.execute_code(operation.code.code)
-            result = self.exec_locals.get(operation.result_var)
-            
-            if execution_result['error']:
-                output_parts.append(f"\nError in operation '{operation.result_var}':")
-                output_parts.append("```python\n" + execution_result['error']['message'] + "\n```")
-            elif result is not None:
-                if result is None or (hasattr(result, '__len__') and len(result) == 0):
-                    output_parts.append(f"\nNo data - Operation '{operation.result_var}' returned empty results")
-                else:
-                    output_parts.append(f"\nResult for '{operation.result_var}':")
-                    output_parts.append("```python\n" + self._format_result(result) + "\n```")
-            else:
-                output_parts.append(f"\nOperation '{operation.result_var}' output:")
-                if execution_result['output'].strip():
-                    output_parts.append("```\n" + execution_result['output'].strip() + "\n```")
-                else:
-                    output_parts.append("No output - Operation didn't produce any results")
+     """Process the response with proper variable handling"""
+     output_parts = [response.casual_response]
+    
+     # Process analysis operations first
+     for operation in response.analysis_operations:
+         execution_result = self.execute_code(operation.code.code)
         
-        # Process charts after all operations
-        if response.charts:
-            output_parts.append("\nVisualizations:")
-            for chart in response.charts:
-                # Skip processing if media pointer indicates blank chart
-                if chart.media_pointer and "blank" in chart.media_pointer.lower():
-                    output_parts.append(f"\nSkipped blank chart: {chart.image_description}")
-                    continue
-                    
-                if chart.code:
-                    chart_result = self.execute_code(chart.code)
-                    if chart_result['plots']:
-                        for plot_data in chart_result['plots']:
-                            try:
-                                public_url = await self.save_plot_to_supabase(
-                                    plot_data=plot_data,
-                                    description=chart.image_description,
-                                    chat_id=chat_id
-                                )
-                                output_parts.append(f"\n{chart.image_description}")
-                                output_parts.append(f"![{chart.image_description}]({public_url})")
-                            except Exception as e:
-                                output_parts.append(f"\nChart error - Couldn't upload: {str(e)}")
-                    elif chart_result['error']:
-                        output_parts.append("```python\n" + f"Chart error: {chart_result['error']['message']}" + "\n```")
-                    else:
-                        output_parts.append(f"\nNo chart - '{chart.image_description}' couldn't be generated")
+         # Get the result from locals
+         result = self.exec_locals.get(operation.result_var)
         
-        return "\n".join(output_parts)
+         if execution_result['error']:
+             output_parts.append(f"\n❌ Error in operation '{operation.result_var}':")
+             output_parts.append("```python\n" + execution_result['error']['message'] + "\n```")
+         elif result is not None:
+             # Handle empty/None results
+             if result is None or (hasattr(result, '__len__') and len(result) == 0):
+                 output_parts.append(f"\n⚠️ Values are missing - Operation '{operation.result_var}' returned no data")
+             else:
+                 output_parts.append(f"\n🔹 Result for '{operation.result_var}':")
+                 output_parts.append("```python\n" + self._format_result(result) + "\n```")
+         else:
+             output_str = execution_result['output'].strip()
+             if output_str:
+                 output_parts.append(f"\nOutput for '{operation.result_var}':")
+                 output_parts.append("```\n" + output_str + "\n```")
+                
+    
+     # Process charts after all operations
+     if response.charts:
+         output_parts.append("\n📊 Visualizations:")
+         for chart in response.charts:
+             if chart.code:
+                 chart_result = self.execute_code(chart.code)
+                 if chart_result['plots']:
+                     for plot_data in chart_result['plots']:
+                         try:
+                             public_url = await self.save_plot_to_supabase(
+                                 plot_data=plot_data,
+                                 description=chart.image_description,
+                                 chat_id=chat_id
+                             )
+                             output_parts.append(f"\n🖼️ {chart.image_description}")
+                             output_parts.append(f"![{chart.image_description}]({public_url})")
+                         except Exception as e:
+                            output_parts.append(f"\n⚠️ Values are missing - Error uploading chart: {str(e)}")
+                 elif chart_result['error']:
+                     output_parts.append("```python\n" + f"Error generating {chart.image_description}: {chart_result['error']['message']}" + "\n```")
+                 else:
+                     output_parts.append(f"\n⚠️ Values are missing - No chart generated for '{chart.image_description}'")
+    
+     return "\n".join(output_parts)