DB Chat for pro users only

python_code_executor_service.py

@@ -16,30 +16,33 @@ import seaborn as sns
 import scipy.stats as stats
 from pydantic import BaseModel
 from tabulate import tabulate
-from PIL import Image  # New import
 
 from supabase_service import upload_file_to_supabase
 
-# Load environment variables
+# Load environment variables from .env file
 load_dotenv()
 
-
 class CodeResponse(BaseModel):
+    """Container for code-related responses"""
     language: str = "python"
     code: str
 
 
 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):
+    """Container for a single analysis operation with its code and result"""
     code: CodeResponse
     result_var: str
 
 
 class CsvChatResult(BaseModel):
+    """Structured response for CSV-related AI interactions"""
     response_type: Literal["casual", "data_analysis", "visualization", "mixed"]
     casual_response: str
     analysis_operations: List[AnalysisOperation]
@@ -47,149 +50,200 @@ class CsvChatResult(BaseModel):
 
 
 class PythonExecutor:
+    """Handles execution of Python code with comprehensive data analysis libraries"""
+    
     def __init__(self, df: pd.DataFrame, charts_folder: str = "generated_charts"):
+        """
+        Initialize the PythonExecutor with a DataFrame
+        
+        Args:
+            df (pd.DataFrame): The DataFrame to operate on
+            charts_folder (str): Folder to save charts in
+        """
         self.df = df
         self.charts_folder = Path(charts_folder)
         self.charts_folder.mkdir(exist_ok=True)
         self.exec_locals = {}
-
+        
     def execute_code(self, code: str) -> Dict[str, Any]:
+        """
+        Execute Python code with full data analysis context and return results
+        
+        Args:
+            code (str): Python code to execute
+            
+        Returns:
+            dict: Dictionary containing execution results and any generated plots
+        """
         output = ""
         error = None
         plots = []
-
+        
+        # Capture stdout
         stdout = io.StringIO()
+        
+        # Monkey patch plt.show() to save figures
         original_show = plt.show
-
+        
         def custom_show():
+            """Custom show function that saves plots instead of displaying them"""
             for i, fig in enumerate(plt.get_fignums()):
                 figure = plt.figure(fig)
+                # Save plot to bytes buffer
                 buf = io.BytesIO()
                 figure.savefig(buf, format='png', bbox_inches='tight')
                 buf.seek(0)
                 plots.append(buf.read())
             plt.close('all')
-
+        
         try:
+            # Create comprehensive execution context with data analysis libraries
             exec_globals = {
+                # Core data analysis
                 'pd': pd,
                 'np': np,
                 'df': self.df,
+                
+                # Visualization
                 'plt': plt,
                 'sns': sns,
                 'tabulate': tabulate,
+                
+                # Statistics
                 'stats': stats,
+                
+                # Date/time
                 'datetime': datetime,
                 'timedelta': timedelta,
                 'time': time,
+                
+                # Utilities
                 'json': json,
                 '__builtins__': __builtins__,
             }
-
+            
+            # Replace plt.show with custom implementation
             plt.show = custom_show
-
+            
+            # Execute code and capture output
             with contextlib.redirect_stdout(stdout):
                 exec(code, exec_globals, self.exec_locals)
-
+            
             output = stdout.getvalue()
-
+            
         except Exception as e:
             error = {
                 "message": str(e),
                 "traceback": traceback.format_exc()
             }
         finally:
+            # Restore original plt.show
             plt.show = original_show
-
+            
         return {
             'output': output,
             'error': error,
             'plots': plots,
             'locals': self.exec_locals
         }
-
-    def is_blank_image(self, image_bytes: bytes) -> bool:
-        """Check if the image is blank (all pixels same intensity)"""
-        try:
-            img = Image.open(io.BytesIO(image_bytes)).convert("L")  # grayscale
-            extrema = img.getextrema()
-            return extrema[0] == extrema[1]  # all pixels same
-        except Exception:
-            return False
-
+    
     async def save_plot_to_supabase(self, plot_data: bytes, description: str, chat_id: str) -> str:
+        """
+        Save plot to Supabase storage and return the public URL
+        
+        Args:
+            plot_data (bytes): Image data in bytes
+            description (str): Description of the plot
+            chat_id (str): ID of the chat session
+            
+        Returns:
+            str: Public URL of the uploaded chart
+        """
+        # Generate unique filename
         filename = f"chart_{uuid.uuid4().hex}.png"
         filepath = self.charts_folder / filename
-
+        
+        # Save the plot locally first
         with open(filepath, 'wb') as f:
             f.write(plot_data)
-
+            
         try:
+            # Upload to Supabase
             public_url = await upload_file_to_supabase(
                 file_path=str(filepath),
                 file_name=filename,
                 chat_id=chat_id
             )
+            
+            # Remove the local file after upload
             os.remove(filepath)
+            
             return public_url
         except Exception as e:
+            # Clean up local file if upload fails
             if os.path.exists(filepath):
                 os.remove(filepath)
             raise Exception(f"Failed to upload plot to Supabase: {e}")
-
+    
     def _format_result(self, result: Any) -> str:
+        """Format the result for display"""
         if isinstance(result, (pd.DataFrame, pd.Series)):
-            json_str = result.to_json(orient='records', date_format='iso')
-            return json.dumps(json.loads(json_str), indent=2)
+            return result.to_string()
         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"\n❌ Error in operation '{operation.result_var}':")
+                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"\n⚠️ Values are missing - Operation '{operation.result_var}' returned no data")
+                    output_parts.append(f"\nNo data - Operation '{operation.result_var}' returned empty results")
                 else:
-                    output_parts.append(f"\n🔹 Result for '{operation.result_var}':")
+                    output_parts.append(f"\nResult 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```")
-
+                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 charts after all operations
         if response.charts:
-            output_parts.append("\n📊 Visualizations:")
+            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']:
-                            if not self.is_blank_image(plot_data):
-                                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)}")
-                            else:
-                                output_parts.append(f"\n⚠️ Skipped blank chart: '{chart.image_description}'")
+                            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"Error generating {chart.image_description}: {chart_result['error']['message']}" + "\n```")
+                        output_parts.append("```python\n" + f"Chart error: {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)
+                        output_parts.append(f"\nNo chart - '{chart.image_description}' couldn't be generated")
+        
+        return "\n".join(output_parts)

requirements.txt

@@ -16,4 +16,3 @@ langchain-google-genai==2.0.7
 supabase==2.13.0
 langchain_openai==0.3.11
 tabulate==0.9.0
-pillow==11.2.1